use std::{env, net::TcpListener, sync::Arc, time::Duration};

use anyhow::{Result, anyhow};
use tici::{
    config::{MetaServiceConfig, make_default_test_meta_config, make_default_test_worker_config},
    meta::{server::MetaServiceServer, service::MetaServiceCore},
};
use tici_common::{config::GrpcConfig, index::IndexKey};
use tici_grpc::{GrpcMetaServiceClient, GrpcMetaServiceInternalClient};
use tici_service_discovery::{
    EtcdClient, EtcdDiscoveryClient, META_SERVICE_ELECTION_KEY, META_SERVICE_REGISTER_KEY_PREFIX,
    discovery::client::DiscoveryClient,
};
use tici_shard::config::WorkerNodeConfig;
use tokio::sync::oneshot::Sender;
use tonic::transport::Channel;

pub struct RunningMeta {
    pub addr: String,
    pub grpc_config: GrpcConfig,
    pub meta_service: Arc<MetaServiceCore>,
    server: Option<MetaServiceServer>,
    shutdown_tx: Option<Sender<()>>,
}

impl RunningMeta {
    pub async fn start(config: MetaServiceConfig) -> Result<Self> {
        let addr = config.server.get_advertise_addr().to_string();
        let grpc_config = config.server.grpc_client.clone();
        let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel();
        let mut server = MetaServiceServer::new(config).await?;
        let meta_service = server.meta_service.clone();
        server.start(shutdown_rx).await?;

        Ok(Self {
            addr,
            grpc_config,
            meta_service,
            server: Some(server),
            shutdown_tx: Some(shutdown_tx),
        })
    }

    pub async fn wait_ready(&self) {
        let _ = MetaServiceServer::get_internal_client(&self.addr, &self.grpc_config).await;
    }

    pub async fn meta_client(&self) -> GrpcMetaServiceClient<Channel> {
        MetaServiceServer::get_client(&self.addr, &self.grpc_config).await
    }

    pub async fn meta_internal_client(&self) -> GrpcMetaServiceInternalClient<Channel> {
        MetaServiceServer::get_internal_client(&self.addr, &self.grpc_config).await
    }

    pub async fn stop(&mut self) -> Result<()> {
        let Some(server) = self.server.take() else {
            return Ok(());
        };
        let shutdown_tx = self
            .shutdown_tx
            .take()
            .ok_or_else(|| anyhow!("missing meta shutdown channel"))?;
        server.stop(shutdown_tx).await
    }
}

pub fn pick_free_port() -> Result<u16> {
    let listener = TcpListener::bind("127.0.0.1:0")?;
    Ok(listener.local_addr()?.port())
}

pub fn configure_meta(addr: String) -> MetaServiceConfig {
    let mut config = make_default_test_meta_config(false);
    config.server.addr = addr.clone();
    config.server.set_advertise_addr(addr);
    if let Ok(pd_addr) = env::var("TICI_TEST_PD_ADDR") {
        config.server.pd_addr = pd_addr;
    }
    if let Ok(tidb_dsn) = env::var("TICI_TEST_TIDB_DSN") {
        config.tidb_server.dsns = vec![tidb_dsn];
    }
    if let Ok(endpoint) = env::var("TICI_TEST_S3_ENDPOINT") {
        config.s3.endpoint = endpoint;
    }
    if let Ok(bucket) = env::var("TICI_TEST_S3_BUCKET") {
        config.s3.bucket = bucket;
    }
    if let Ok(prefix) = env::var("TICI_TEST_S3_PREFIX") {
        config.s3.prefix = prefix;
    }
    config
}

pub fn configure_worker(addr: String) -> WorkerNodeConfig {
    let mut config = make_default_test_worker_config(false);
    config.server.addr = addr.clone();
    config.server.set_advertise_addr(addr);
    if let Ok(pd_addr) = env::var("TICI_TEST_PD_ADDR") {
        config.server.pd_addr = pd_addr;
    }
    if let Ok(endpoint) = env::var("TICI_TEST_S3_ENDPOINT") {
        config.s3.endpoint = endpoint;
    }
    if let Ok(bucket) = env::var("TICI_TEST_S3_BUCKET") {
        config.s3.bucket = bucket;
    }
    if let Ok(prefix) = env::var("TICI_TEST_S3_PREFIX") {
        config.s3.prefix = prefix;
    }
    config
}

pub fn pd_endpoints(pd_addr: &str) -> Vec<String> {
    vec![
        if pd_addr.starts_with("http://") || pd_addr.starts_with("https://") {
            pd_addr.to_string()
        } else {
            format!("http://{}", pd_addr)
        },
    ]
}

pub fn make_observer(pd_addr: &str) -> EtcdDiscoveryClient {
    EtcdDiscoveryClient::new(
        META_SERVICE_REGISTER_KEY_PREFIX,
        Arc::new(EtcdClient::new_for_test(pd_endpoints(pd_addr))),
        Duration::from_secs(10),
        Duration::from_secs(3),
    )
}

pub async fn current_leader(observer: &EtcdDiscoveryClient) -> Result<String> {
    let peers = observer.get_all_peers().await?;
    peers
        .leader
        .ok_or_else(|| anyhow!("meta leader is not available"))
}

pub async fn leader_has_worker(
    observer: &EtcdDiscoveryClient,
    meta1: &RunningMeta,
    meta2: &RunningMeta,
    worker_addr: &str,
) -> bool {
    let Ok(leader_addr) = current_leader(observer).await else {
        return false;
    };

    match leader_addr.as_str() {
        addr if addr == meta1.addr => {
            meta1
                .meta_service
                .worker_node_mgr
                .has_worker_node(worker_addr)
                .await
        }
        addr if addr == meta2.addr => {
            meta2
                .meta_service
                .worker_node_mgr
                .has_worker_node(worker_addr)
                .await
        }
        _ => false,
    }
}

pub async fn any_meta_has_worker(
    meta1: &RunningMeta,
    meta2: &RunningMeta,
    worker_addr: &str,
) -> bool {
    meta1
        .meta_service
        .worker_node_mgr
        .has_worker_node(worker_addr)
        .await
        || meta2
            .meta_service
            .worker_node_mgr
            .has_worker_node(worker_addr)
            .await
}

pub async fn cleanup_service_discovery_election(pd_addr: &str) -> Result<()> {
    let client = EtcdClient::new_for_test(pd_endpoints(pd_addr));
    let mut etcd = client.get_client().await?;
    etcd.delete(META_SERVICE_ELECTION_KEY, None).await?;
    Ok(())
}

pub async fn revoke_current_leader_lease(pd_addr: &str) -> Result<()> {
    let client = EtcdClient::new_for_test(pd_endpoints(pd_addr));
    let mut etcd = client.get_client().await?;
    let resp = etcd.leader(META_SERVICE_ELECTION_KEY).await?;
    let kv = resp
        .kv()
        .ok_or_else(|| anyhow!("missing leader kv when revoking current leader lease"))?;
    etcd.lease_revoke(kv.lease()).await?;
    Ok(())
}

pub async fn must_create_index_on_meta(meta: &RunningMeta, index_key: &IndexKey) -> Result<()> {
    let request = crate::test_utils::workload_ci::make_create_index_req(index_key);
    let mut client = meta.meta_client().await;
    for attempt in 0..10 {
        let response = client.create_index(request.clone()).await?;
        if response.get_ref().status == tici_grpc::ErrorCode::Success as i32 {
            return Ok(());
        }

        tokio::time::sleep(Duration::from_secs(1)).await;
        log::info!(
            "Retry create_index on {} after status={} attempt={}/10",
            meta.addr,
            response.get_ref().status,
            attempt + 1
        );
    }

    Err(anyhow!(
        "create_index did not succeed on {} after retries",
        meta.addr
    ))
}