use std::{panic::AssertUnwindSafe, time::Duration};

use anyhow::{Result, anyhow};
use futures::FutureExt;
use log::info;
use scopeguard::defer;
use tici::meta::Bootstrap;
use tici_service_discovery::discovery::client::DiscoveryClient;
use tici_test_utils::test_utils::{
    meta_failover::{
        RunningMeta, any_meta_has_worker, cleanup_service_discovery_election, configure_meta,
        configure_worker, current_leader, leader_has_worker, make_observer, pick_free_port,
        revoke_current_leader_lease,
    },
    worker_handle::WorkerServiceHandle,
};

async fn wait_until<F, Fut>(max_attempts: usize, label: &str, f: F) -> Result<()>
where
    F: Fn() -> Fut,
    Fut: std::future::Future<Output = bool>,
{
    for attempt in 0..max_attempts {
        if f().await {
            return Ok(());
        }
        tokio::time::sleep(Duration::from_secs(1)).await;
        info!(
            "Waiting for {} (attempt {}/{})",
            label,
            attempt + 1,
            max_attempts
        );
    }
    Err(anyhow!(
        "Timed out waiting for {} after {} attempts",
        label,
        max_attempts
    ))
}

#[test]
fn test_meta_failover_keeps_worker_heartbeats() -> Result<()> {
    fail::cfg("worker_node_client_fail_fast", "return(2)").map_err(anyhow::Error::msg)?;
    fail::cfg("leader_term_skip_watch_event", "return").map_err(anyhow::Error::msg)?;
    defer!({
        fail::remove("worker_node_client_fail_fast");
        fail::remove("leader_term_skip_watch_event");
    });

    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()?;

    rt.block_on(async {
        let mut meta1 = None;
        let mut meta2 = None;
        let mut worker = None;

        let test_result = AssertUnwindSafe(async {
            let meta_addr1 = format!("127.0.0.1:{}", pick_free_port()?);
            let meta_addr2 = format!("127.0.0.1:{}", pick_free_port()?);
            let worker_addr = format!("127.0.0.1:{}", pick_free_port()?);

            let mut meta_cfg1 = configure_meta(meta_addr1.clone());
            let mut meta_cfg2 = configure_meta(meta_addr2.clone());
            let mut worker_cfg = configure_worker(worker_addr.clone());
            let unique_s3_prefix = if meta_cfg1.s3.prefix.is_empty() {
                format!("service_discovery_{}", uuid::Uuid::new_v4())
            } else {
                format!(
                    "{}/service_discovery_{}",
                    meta_cfg1.s3.prefix,
                    uuid::Uuid::new_v4()
                )
            };
            meta_cfg1.s3.prefix = unique_s3_prefix.clone();
            meta_cfg2.s3.prefix = unique_s3_prefix.clone();
            worker_cfg.s3.prefix = unique_s3_prefix;

            Bootstrap::new_with_config(&meta_cfg1)
                .clear_all_meta(true)
                .await?;
            cleanup_service_discovery_election(&meta_cfg1.server.pd_addr).await?;

            let observer = make_observer(&meta_cfg1.server.pd_addr);

            let running_meta1 = RunningMeta::start(meta_cfg1.clone()).await?;
            let running_meta2 = RunningMeta::start(meta_cfg2.clone()).await?;
            running_meta1.wait_ready().await;
            running_meta2.wait_ready().await;

            wait_until(30, "two metas registered with a leader", || {
                let observer = &observer;
                let meta_addr1 = meta_addr1.clone();
                let meta_addr2 = meta_addr2.clone();
                async move {
                    let Ok(peers) = observer.get_all_peers().await else {
                        return false;
                    };
                    peers.nodes.contains(&meta_addr1)
                        && peers.nodes.contains(&meta_addr2)
                        && matches!(
                            peers.leader.as_deref(),
                            Some(addr) if addr == meta_addr1 || addr == meta_addr2
                        )
                }
            })
            .await?;

            worker = Some(WorkerServiceHandle::new(worker_cfg, None).await?);
            meta1 = Some(running_meta1);
            meta2 = Some(running_meta2);

            wait_until(30, "initial worker heartbeat registration", || {
                let meta1 = &meta1;
                let meta2 = &meta2;
                let worker_addr = worker_addr.clone();
                async move {
                    let (Some(meta1), Some(meta2)) = (meta1.as_ref(), meta2.as_ref()) else {
                        return false;
                    };
                    any_meta_has_worker(meta1, meta2, &worker_addr).await
                }
            })
            .await?;

            let old_leader = current_leader(&observer).await?;
            info!("Revoking current meta leader lease at {}", old_leader);
            revoke_current_leader_lease(&meta_cfg1.server.pd_addr).await?;

            wait_until(30, "new leader election after lease revoke", || {
                let observer = &observer;
                let old_leader = old_leader.clone();
                async move {
                    matches!(
                        current_leader(observer).await,
                        Ok(leader) if leader != old_leader
                    )
                }
            })
            .await?;

            wait_until(20, "worker heartbeat recovery after NOT_LEADER", || {
                let observer = &observer;
                let meta1 = &meta1;
                let meta2 = &meta2;
                let worker_addr = worker_addr.clone();
                async move {
                    let (Some(meta1), Some(meta2)) = (meta1.as_ref(), meta2.as_ref()) else {
                        return false;
                    };
                    leader_has_worker(observer, meta1, meta2, &worker_addr).await
                }
            })
            .await?;

            Ok::<(), anyhow::Error>(())
        })
        .catch_unwind()
        .await;

        if let Some(worker) = worker.take() {
            let _ = worker.stop().await;
        }
        if let Some(meta) = meta1.as_mut() {
            let _ = meta.stop().await;
        }
        if let Some(meta) = meta2.as_mut() {
            let _ = meta.stop().await;
        }

        match test_result {
            Ok(Ok(())) => Ok(()),
            Ok(Err(err)) => Err(err),
            Err(panic) => std::panic::resume_unwind(panic),
        }
    })
}