use flaredb_proto::kvrpc::kv_cas_server::KvCas;
use flaredb_proto::kvrpc::{CasRequest, GetRequest};
use flaredb_server::config::ServerConfig;
use flaredb_server::service::KvServiceImpl;
use flaredb_server::store::Store;
use flaredb_storage::rocks_engine::RocksEngine;
use flaredb_types::RegionMeta;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::time::Duration;
use tonic::Request;

fn test_peer_addrs() -> Arc<HashMap<u64, String>> {
    let mut addrs = HashMap::new();
    addrs.insert(1, "127.0.0.1:50051".to_string());
    Arc::new(addrs)
}

#[tokio::test]
async fn strong_write_succeeds_after_leader_election() {
    let dir = tempfile::TempDir::new().unwrap();
    let engine = Arc::new(RocksEngine::new(dir.path().to_str().unwrap()).unwrap());
    let config = Arc::new(ServerConfig::default_with_namespace());
    let store = Arc::new(Store::new(1, engine.clone(), config.clone(), test_peer_addrs()));
    store
        .bootstrap_regions(vec![(
            RegionMeta {
                id: 1,
                start_key: Vec::new(),
                end_key: Vec::new(),
            },
            vec![1],
        )])
        .await
        .unwrap();
    let service = KvServiceImpl::new(engine.clone(), config, store.clone());

    // Ensure leader using OpenRaft (single-node cluster elects itself quickly)
    if let Some(node) = store.get_raft_node(1).await {
        node.trigger_election().await.expect("trigger election");
        // Wait for leader election
        tokio::time::sleep(Duration::from_millis(500)).await;
    }

    // CAS succeeds once leader
    let resp = service
        .compare_and_swap(Request::new(CasRequest {
            namespace: "default".into(),
            key: b"k".to_vec(),
            value: b"v".to_vec(),
            expected_version: 0,
        }))
        .await
        .unwrap()
        .into_inner();
    assert!(resp.success);

    // Read requires leader too
    let get = service
        .get(Request::new(GetRequest {
            namespace: "default".into(),
            key: b"k".to_vec(),
        }))
        .await
        .unwrap()
        .into_inner();
    assert!(get.found);
    assert_eq!(get.value, b"v");
}