use std::sync::Arc;

use anyhow::{Result, anyhow};
use etcd_client::WatchOptions;
use log::warn;
use tokio::sync::{Mutex, RwLock};

use crate::{etcd_client::EtcdClient, watch_supervisor::WatchSupervisor};

/// If no watch response (including progress notifications) has been
/// received for this long, consider the watch connection potentially dead
/// and fall back to querying etcd directly.
const WATCH_LIVENESS_THRESHOLD: std::time::Duration = std::time::Duration::from_secs(30);

#[derive(Default)]
struct LeaderInfo {
    address: String,
    lease_id: i64,
}

impl LeaderInfo {
    fn update(&mut self, address: String, lease_id: i64) {
        self.address = address;
        self.lease_id = lease_id;
    }

    fn reset(&mut self) {
        self.address = String::new();
        self.lease_id = 0;
    }
}

/// High-Availability leader term that can handle etcd cluster changes.
///
/// Uses a watch supervisor that automatically restarts watch sessions
/// with exponential backoff on failures. When the watch is temporarily
/// unavailable, `get_leader_address()` falls back to querying etcd directly.
pub struct LeaderTerm {
    leader_prefix: &'static str,
    client: Arc<EtcdClient>,
    supervisor: WatchSupervisor,
    watcher: Arc<Mutex<Option<etcd_client::Watcher>>>,
    leader_info: Arc<RwLock<LeaderInfo>>,
    /// Updated on every watch response (events + progress notifications).
    /// Used to detect silently-dead connections.
    last_watch_response: Arc<std::sync::Mutex<std::time::Instant>>,
}

fn parse_kv(key: &[u8], value: &[u8], prefix: &'static str) -> Result<(i64, String)> {
    if key.starts_with(prefix.as_bytes()) {
        let lease_id_str = std::str::from_utf8(&key[prefix.len() + 1..])?;
        let lease_id = i64::from_str_radix(lease_id_str, 16)?;
        let addr = std::str::from_utf8(value)?.to_string();
        return Ok((lease_id, addr));
    }
    Err(anyhow!("Invalid KV: {:?} -> {:?}", key, value))
}

impl LeaderTerm {
    pub fn new(leader_prefix: &'static str, client: Arc<EtcdClient>) -> Self {
        Self {
            leader_prefix,
            client,
            supervisor: WatchSupervisor::new("leader"),
            watcher: Arc::new(Mutex::new(None)),
            leader_info: Arc::new(RwLock::new(LeaderInfo::default())),
            last_watch_response: Arc::new(std::sync::Mutex::new(std::time::Instant::now())),
        }
    }

    /// Query etcd for the current leader and update leader_info.
    async fn get_leader(&self) -> Result<()> {
        let mut client = self.client.get_client().await?;
        match client.leader(self.leader_prefix).await {
            Ok(resp) => match resp.kv() {
                Some(kv) => {
                    let (lease_id, addr) = parse_kv(kv.key(), kv.value(), self.leader_prefix)?;
                    self.leader_info.write().await.update(addr, lease_id);
                    Ok(())
                }
                None => Err(anyhow!("Unexpected empty response")),
            },
            Err(e) => Err(anyhow!("Unexpected err={}", e)),
        }
    }

    pub async fn start(&mut self) -> Result<()> {
        // Ensure the client is started
        self.client.start().await?;

        // Init leader_info from etcd
        self.get_leader().await?;

        // Cancel existing watcher if any (avoid leaking server-side watch)
        if let Some(mut watcher) = self.watcher.lock().await.take() {
            let _ = watcher.cancel().await;
        }

        // Start the watch supervisor
        let prefix = self.leader_prefix;
        let client = self.client.clone();
        let info = self.leader_info.clone();
        let watcher_slot = self.watcher.clone();
        let session_flag = self.supervisor.session_active_flag();
        let liveness = self.last_watch_response.clone();

        self.supervisor
            .start(Box::new(move |shutdown| {
                let client = client.clone();
                let info = info.clone();
                let watcher_slot = watcher_slot.clone();
                let session_flag = session_flag.clone();
                let liveness = liveness.clone();
                Box::pin(async move {
                    Self::run_watch_session(
                        prefix,
                        client,
                        info,
                        watcher_slot,
                        session_flag,
                        liveness,
                        shutdown,
                    )
                    .await
                })
            }))
            .await;

        // Wait for the first watch session to establish before returning,
        // so callers can trust that the watch is active.
        if let Err(e) = self
            .supervisor
            .wait_session_active(std::time::Duration::from_secs(10))
            .await
        {
            // Clean up the supervisor to avoid leaking a reconnect loop.
            self.supervisor.stop().await;
            return Err(e);
        }
        Ok(())
    }

    /// Refresh leader_info by querying etcd directly.
    /// Used to catch leader changes that occurred during the gap between
    /// watch sessions. Returns Err on transient failures so the caller
    /// can abort the session (and the supervisor will retry).
    async fn refresh_leader_info(
        prefix: &'static str,
        client: &mut etcd_client::Client,
        info: &RwLock<LeaderInfo>,
    ) -> Result<()> {
        match client.leader(prefix).await {
            Ok(resp) => {
                if let Some(kv) = resp.kv() {
                    let (lease_id, addr) = parse_kv(kv.key(), kv.value(), prefix)?;
                    info.write().await.update(addr, lease_id);
                } else {
                    // No leader currently elected — clear cache so callers
                    // get an error rather than talking to a stale leader.
                    info.write().await.reset();
                }
                Ok(())
            }
            Err(e) => {
                // "election: no leader" means nobody has campaigned yet —
                // clear cache, but this is not a transient failure.
                if let etcd_client::Error::GRpcStatus(ref status) = e {
                    if status.message() == "election: no leader" {
                        info.write().await.reset();
                        return Ok(());
                    }
                }
                Err(anyhow!("Failed to refresh leader info: {}", e))
            }
        }
    }

    /// Runs a single watch session. Returns when the stream ends or an error
    /// occurs. Parse errors on individual events are logged but do not kill
    /// the session.
    async fn run_watch_session(
        prefix: &'static str,
        client: Arc<EtcdClient>,
        info: Arc<RwLock<LeaderInfo>>,
        watcher_slot: Arc<Mutex<Option<etcd_client::Watcher>>>,
        session_flag: std::sync::Arc<std::sync::atomic::AtomicBool>,
        liveness: Arc<std::sync::Mutex<std::time::Instant>>,
        shutdown: tokio_util::sync::CancellationToken,
    ) -> Result<()> {
        let mut etcd = client.get_client().await?;
        let option = Some(WatchOptions::new().with_prefix().with_progress_notify());
        let (mut watcher, mut stream) = etcd.watch(prefix, option).await?;
        watcher.request_progress().await?;

        // Store watcher so it can be cancelled externally on stop()
        *watcher_slot.lock().await = Some(watcher);

        // Refresh leader_info to catch any changes that occurred during
        // the gap between the previous watch session ending and this one
        // starting. The new watch only sees future events, so without this
        // refresh a leader change during the gap would be invisible.
        // If the refresh fails, abort this session so the supervisor retries
        // — we must not mark the session active with potentially stale data.
        Self::refresh_leader_info(prefix, &mut etcd, &info).await?;

        // Mark session as active only after watch is established and
        // gap-refresh is done, so callers can trust the cached leader_info.
        session_flag.store(true, std::sync::atomic::Ordering::Release);
        *liveness.lock().unwrap() = std::time::Instant::now();

        loop {
            let resp = tokio::select! {
                msg = stream.message() => {
                    match msg {
                        Ok(Some(resp)) => resp,
                        Ok(None) => return Ok(()),  // stream closed
                        Err(e) => return Err(e.into()),
                    }
                }
                _ = shutdown.cancelled() => return Ok(()),
            };

            // Update liveness on every watch response (events, progress
            // notifications, etc.) to track connection health.
            *liveness.lock().unwrap() = std::time::Instant::now();

            if resp.canceled() {
                return Ok(());
            }

            for event in resp.events() {
                let skip_event = (|| {
                    fail::fail_point!("leader_term_skip_watch_event", |_| true);
                    false
                })();
                if skip_event {
                    continue;
                }

                if let Some(kv) = event.kv() {
                    match event.event_type() {
                        etcd_client::EventType::Put => {
                            match parse_kv(kv.key(), kv.value(), prefix) {
                                Ok((lease_id, addr)) => {
                                    info.write().await.update(addr, lease_id);
                                }
                                Err(e) => {
                                    warn!("Failed to parse watch event: {}", e);
                                }
                            }
                        }
                        etcd_client::EventType::Delete => {
                            // Only clear cache if the deleted key belongs
                            // to the current leader. Non-leader waiters
                            // timing out also produce deletes under this
                            // prefix, so an unconditional reset would cause
                            // spurious leader loss.
                            if let Ok((lease_id, _)) = parse_kv(kv.key(), kv.value(), prefix)
                                .or_else(|_| {
                                    // Delete events may have empty values;
                                    // extract lease_id from the key alone.
                                    let suffix =
                                        std::str::from_utf8(&kv.key()[prefix.len() + 1..])?;
                                    let id = i64::from_str_radix(suffix, 16)?;
                                    Ok::<_, anyhow::Error>((id, String::new()))
                                })
                            {
                                if info.read().await.lease_id == lease_id {
                                    info.write().await.reset();
                                }
                            }
                        }
                    }
                }
            }
        }
    }

    pub async fn stop(&mut self) -> Result<()> {
        // Stop the supervisor (cancels shutdown token + aborts task)
        self.supervisor.stop().await;

        // Cancel the watcher (server-side watch cleanup)
        if let Some(mut watcher) = self.watcher.lock().await.take() {
            let _ = watcher.cancel().await;
        }

        // Reset leader_info
        self.leader_info.write().await.reset();

        // Stop the client
        self.client.stop().await?;

        Ok(())
    }

    async fn cached_leader_address(&self) -> Result<String> {
        let addr = self.leader_info.read().await.address.clone();
        if addr.is_empty() {
            return Err(anyhow!("Leader address is empty"));
        }
        Ok(addr)
    }

    pub async fn refresh_leader_address(&self) -> Result<String> {
        if self.supervisor.is_stopped() {
            return Err(anyhow!("LeaderTerm is stopped"));
        }

        let mut client = self.client.get_client().await?;
        Self::refresh_leader_info(self.leader_prefix, &mut client, &self.leader_info).await?;
        self.cached_leader_address().await
    }

    pub async fn get_leader_address(&self) -> Result<String> {
        // After stop(), do not reactivate the etcd client via fallback.
        if self.supervisor.is_stopped() {
            return Err(anyhow!("LeaderTerm is stopped"));
        }

        // Fall back to querying etcd directly when the watch may be unhealthy:
        // - No active watch session (supervisor in backoff or terminated)
        // - No watch response received for WATCH_LIVENESS_THRESHOLD (covers the case
        //   where stream.message() hangs on a silently-dead connection)
        let watch_stale =
            self.last_watch_response.lock().unwrap().elapsed() > WATCH_LIVENESS_THRESHOLD;
        if !self.supervisor.is_session_active() || watch_stale {
            if let Err(e) = self.get_leader().await {
                return Err(anyhow!(
                    "Watch is not active and fallback query failed: {}",
                    e
                ));
            }
        }

        self.cached_leader_address().await
    }
}

#[cfg(test)]
mod tests {
    use std::time::Duration;

    use super::*;
    use crate::{
        META_SERVICE_ELECTION_KEY,
        election::{client::ElectionClient, etcd_client::EtcdElectionClient},
        test_utils::setup_etcd,
    };

    const TEST_KEEP_ALIVE_INTERVAL: Duration = Duration::from_millis(200);

    /// Poll a condition until it becomes true, or return false after timeout.
    async fn wait_until<F, Fut>(timeout: Duration, mut check: F) -> bool
    where
        F: FnMut() -> Fut,
        Fut: std::future::Future<Output = bool>,
    {
        let start = tokio::time::Instant::now();
        loop {
            if check().await {
                return true;
            }
            if start.elapsed() >= timeout {
                return false;
            }
            tokio::time::sleep(Duration::from_millis(100)).await;
        }
    }

    /// Helper: skip test with a visible warning when ETCD_ADDR is not set.
    async fn require_etcd() -> Option<(String, Arc<EtcdClient>)> {
        match setup_etcd().await {
            Ok(endpoint) => {
                let client = Arc::new(EtcdClient::new_for_test(vec![endpoint.clone()]));
                Some((endpoint, client))
            }
            Err(_) => {
                eprintln!(
                    "SKIPPED: ETCD_ADDR not set. \
                     Run with ETCD_ADDR=http://127.0.0.1:2379 to enable this test."
                );
                None
            }
        }
    }

    #[test]
    fn test_parse_kv() {
        let key: &[u8] = b"/tici/metaservice/election/2231978c288f0a16";
        let value: &[u8] = b"127.0.0.1:8500";
        let lease_id = 2463917099380050454i64;

        let res = parse_kv(key, value, META_SERVICE_ELECTION_KEY);
        assert!(res.is_ok());
        assert_eq!(lease_id, res.as_ref().unwrap().0);
        assert_eq!("127.0.0.1:8500".to_string(), res.as_ref().unwrap().1);
    }

    // Reproduces issue #739: when the watch session breaks, the supervisor
    // should automatically restart it and get_leader_address() should
    // continue to work (either via watch or via etcd fallback).
    #[tokio::test]
    async fn test_get_leader_address_survives_watch_break() {
        let Some((_, etcd_client)) = require_etcd().await else {
            return;
        };
        let leader_key: &'static str = "/tici/test_watch_fallback/election";
        let leader_addr = "127.0.0.1:19001";

        // Campaign to become leader
        let election_client = EtcdElectionClient::new(
            leader_key,
            leader_addr.to_string(),
            etcd_client.clone(),
            TEST_KEEP_ALIVE_INTERVAL,
        );
        let lease_id = election_client
            .campaign(Duration::from_secs(30))
            .await
            .unwrap();
        assert!(lease_id.is_some());

        // Start LeaderTerm watching the leader
        let mut leader_term = LeaderTerm::new(leader_key, etcd_client.clone());
        leader_term.start().await.unwrap();

        // Verify initial state works
        let addr = leader_term.get_leader_address().await.unwrap();
        assert_eq!(addr, leader_addr);

        // Break the watch by cancelling the watcher
        if let Some(mut w) = leader_term.watcher.lock().await.take() {
            let _ = w.cancel().await;
        }

        // Poll until supervisor recovers and leader address is available
        let recovered = wait_until(Duration::from_secs(10), || async {
            leader_term.get_leader_address().await.is_ok()
        })
        .await;
        assert!(
            recovered,
            "should survive watch break via supervisor restart"
        );
        assert_eq!(leader_term.get_leader_address().await.unwrap(), leader_addr);

        // Cleanup
        election_client.resign().await.unwrap();
        leader_term.stop().await.unwrap_or_default();
    }

    // After a watch break during leader transition A→B, the supervisor
    // should detect the new leader via gap refresh on session restart.
    #[tokio::test]
    async fn test_detects_new_leader_after_watch_break() {
        let Some((_, etcd_client)) = require_etcd().await else {
            return;
        };
        let leader_key: &'static str = "/tici/test_new_leader_fallback/election";
        let addr_a = "127.0.0.1:19002";
        let addr_b = "127.0.0.1:19003";

        // Campaign leader A
        let election_a = EtcdElectionClient::new(
            leader_key,
            addr_a.to_string(),
            etcd_client.clone(),
            TEST_KEEP_ALIVE_INTERVAL,
        );
        let lease_id = election_a.campaign(Duration::from_secs(30)).await.unwrap();
        assert!(lease_id.is_some());

        // Start LeaderTerm, verify it sees leader A
        let mut leader_term = LeaderTerm::new(leader_key, etcd_client.clone());
        leader_term.start().await.unwrap();
        assert_eq!(leader_term.get_leader_address().await.unwrap(), addr_a);

        // Break the watch
        if let Some(mut w) = leader_term.watcher.lock().await.take() {
            let _ = w.cancel().await;
        }

        // Leader A resigns, Leader B campaigns (during the watch gap)
        election_a.resign().await.unwrap();
        let election_b = EtcdElectionClient::new(
            leader_key,
            addr_b.to_string(),
            etcd_client.clone(),
            TEST_KEEP_ALIVE_INTERVAL,
        );
        let lease_id = election_b.campaign(Duration::from_secs(30)).await.unwrap();
        assert!(lease_id.is_some());

        // Poll until supervisor recovers and detects new leader B
        let detected = wait_until(Duration::from_secs(10), || async {
            leader_term.get_leader_address().await.ok().as_deref() == Some(addr_b)
        })
        .await;
        assert!(detected, "should detect new leader via supervisor restart");

        // Cleanup
        election_b.resign().await.unwrap();
        leader_term.stop().await.unwrap_or_default();
    }

    // After stop(), get_leader_address() must not reactivate the etcd client.
    #[tokio::test]
    async fn test_get_leader_address_errors_after_stop() {
        let Some((_, etcd_client)) = require_etcd().await else {
            return;
        };
        let leader_key: &'static str = "/tici/test_stop_semantics/election";
        let leader_addr = "127.0.0.1:19004";

        let election_client = EtcdElectionClient::new(
            leader_key,
            leader_addr.to_string(),
            etcd_client.clone(),
            TEST_KEEP_ALIVE_INTERVAL,
        );
        let lease_id = election_client
            .campaign(Duration::from_secs(30))
            .await
            .unwrap();
        assert!(lease_id.is_some());

        let mut leader_term = LeaderTerm::new(leader_key, etcd_client.clone());
        leader_term.start().await.unwrap();
        assert_eq!(leader_term.get_leader_address().await.unwrap(), leader_addr);

        leader_term.stop().await.unwrap();

        // After stop, get_leader_address must return error, not silently reconnect.
        let result = leader_term.get_leader_address().await;
        assert!(result.is_err(), "should error after stop()");
        assert!(
            result.unwrap_err().to_string().contains("stopped"),
            "error should indicate LeaderTerm is stopped"
        );

        election_client.resign().await.unwrap();
    }

    // During the backoff window after a watch break, get_leader_address()
    // should fall back to querying etcd directly instead of returning
    // the stale cached leader_info.
    #[tokio::test]
    async fn test_stale_leader_during_backoff_window() {
        let Some((_, etcd_client)) = require_etcd().await else {
            return;
        };
        let leader_key: &'static str = "/tici/test_stale_backoff/election";
        let addr_a = "127.0.0.1:19005";
        let addr_b = "127.0.0.1:19006";

        // Campaign leader A
        let election_a = EtcdElectionClient::new(
            leader_key,
            addr_a.to_string(),
            etcd_client.clone(),
            TEST_KEEP_ALIVE_INTERVAL,
        );
        assert!(
            election_a
                .campaign(Duration::from_secs(30))
                .await
                .unwrap()
                .is_some()
        );

        // Start LeaderTerm
        let mut leader_term = LeaderTerm::new(leader_key, etcd_client.clone());
        leader_term.start().await.unwrap();
        assert_eq!(leader_term.get_leader_address().await.unwrap(), addr_a);

        // Break the watch
        if let Some(mut w) = leader_term.watcher.lock().await.take() {
            let _ = w.cancel().await;
        }

        // Switch leader A→B while supervisor is in backoff
        election_a.resign().await.unwrap();
        let election_b = EtcdElectionClient::new(
            leader_key,
            addr_b.to_string(),
            etcd_client.clone(),
            TEST_KEEP_ALIVE_INTERVAL,
        );
        assert!(
            election_b
                .campaign(Duration::from_secs(30))
                .await
                .unwrap()
                .is_some()
        );

        // Small sleep to let watcher cancel propagate, but NOT long enough
        // for the supervisor to fully restart.
        tokio::time::sleep(Duration::from_millis(200)).await;

        // get_leader_address() should detect no active session and query etcd
        // directly, returning the new leader B instead of stale A.
        let result = leader_term.get_leader_address().await;
        assert!(
            result.is_ok(),
            "should fall back to etcd query during backoff, got: {:?}",
            result.err()
        );
        assert_eq!(result.unwrap(), addr_b);

        // Cleanup
        election_b.resign().await.unwrap();
        leader_term.stop().await.unwrap_or_default();
    }

    // When the old leader resigns during the watch gap and no new leader
    // exists, get_leader_address() should return an error instead of the
    // stale old leader address.
    #[tokio::test]
    async fn test_no_leader_during_gap() {
        let Some((_, etcd_client)) = require_etcd().await else {
            return;
        };
        let leader_key: &'static str = "/tici/test_no_leader_gap/election";
        let addr_a = "127.0.0.1:19007";

        // Campaign leader A
        let election_a = EtcdElectionClient::new(
            leader_key,
            addr_a.to_string(),
            etcd_client.clone(),
            TEST_KEEP_ALIVE_INTERVAL,
        );
        assert!(
            election_a
                .campaign(Duration::from_secs(30))
                .await
                .unwrap()
                .is_some()
        );

        // Start LeaderTerm, verify it sees leader A
        let mut leader_term = LeaderTerm::new(leader_key, etcd_client.clone());
        leader_term.start().await.unwrap();
        assert_eq!(leader_term.get_leader_address().await.unwrap(), addr_a);

        // Break the watch
        if let Some(mut w) = leader_term.watcher.lock().await.take() {
            let _ = w.cancel().await;
        }

        // Leader A resigns, no new leader campaigns
        election_a.resign().await.unwrap();

        // Poll until get_leader_address returns an error (no leader)
        let errored = wait_until(Duration::from_secs(10), || async {
            leader_term.get_leader_address().await.is_err()
        })
        .await;
        assert!(errored, "should error when no leader exists");

        // Cleanup
        leader_term.stop().await.unwrap_or_default();
    }
}