use std::{
    collections::HashSet,
    sync::{Arc, atomic::Ordering},
    time::Duration,
};

use async_trait::async_trait;
use etcd_client::{GetOptions, PutOptions, WatchOptions};
use log::{error, info};
use tokio::{
    sync::{Mutex, RwLock},
    task::JoinHandle,
};

use crate::{
    META_SERVICE_ELECTION_KEY,
    discovery::client::{AllPeers, DiscoveryClient},
    election::etcd_lease::EtcdLease,
    etcd_client::EtcdClient,
    watch_supervisor::WatchSupervisor,
};

/// If no watch response (including progress notifications) has been
/// received for this long, consider the watch unhealthy and refresh
/// the cached node list directly from etcd.
const WATCH_LIVENESS_THRESHOLD: Duration = Duration::from_secs(30);

struct RegistrationLease {
    key: String,
    lease: Arc<EtcdLease>,
    keep_alive_handle: JoinHandle<()>,
}

// EtcdDiscoveryClient implements the DiscoveryClient trait for service
// registration and discovery using etcd. It allows services to register
// themselves, deregister, and discover other services in the cluster.
// It can only register one service at a time.
pub struct EtcdDiscoveryClient {
    prefix: &'static str,
    client: Arc<EtcdClient>,
    supervisor: WatchSupervisor,
    nodes: Arc<RwLock<HashSet<String>>>,
    watcher: Arc<Mutex<Option<etcd_client::Watcher>>>,
    registration: Option<RegistrationLease>,
    registration_lease_ttl: Duration,
    registration_keep_alive_interval: Duration,
    /// Updated on every watch response (events + progress notifications).
    /// Used to detect silently-dead connections.
    last_watch_response: Arc<std::sync::Mutex<std::time::Instant>>,
}

impl EtcdDiscoveryClient {
    pub fn new(
        prefix: &'static str,
        etcd_client: Arc<EtcdClient>,
        registration_lease_ttl: Duration,
        registration_keep_alive_interval: Duration,
    ) -> Self {
        Self {
            prefix,
            client: etcd_client,
            supervisor: WatchSupervisor::new("discovery"),
            nodes: Arc::new(RwLock::new(HashSet::new())),
            watcher: Arc::new(Mutex::new(None)),
            registration: None,
            registration_lease_ttl,
            registration_keep_alive_interval,
            last_watch_response: Arc::new(std::sync::Mutex::new(std::time::Instant::now())),
        }
    }

    /// Start the watch supervisor. It will continuously maintain a watch
    /// session, automatically restarting with exponential backoff on failure.
    async fn start_watch_supervisor(&mut self) {
        let prefix = self.prefix;
        let client = self.client.clone();
        let nodes = self.nodes.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 nodes = nodes.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,
                        nodes,
                        watcher_slot,
                        session_flag,
                        liveness,
                        shutdown,
                    )
                    .await
                })
            }))
            .await;
    }

    /// Runs a single watch session. Creates a new etcd watch, refreshes
    /// existing nodes (to catch changes during the gap between sessions),
    /// and processes watch events until the stream ends or shutdown is
    /// signalled.
    async fn run_watch_session(
        prefix: &'static str,
        client: Arc<EtcdClient>,
        nodes: Arc<RwLock<HashSet<String>>>,
        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,
    ) -> anyhow::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 existing nodes to catch changes during the gap between
        // the previous session ending and this watch starting.
        // If the refresh fails, abort this session so the supervisor retries
        // — we must not mark the session active with potentially stale data.
        Self::refresh_nodes(prefix, &mut etcd, &nodes).await?;

        // Mark session as active only after watch + gap-refresh are done.
        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(()),
                        Err(e) => return Err(e.into()),
                    }
                }
                _ = shutdown.cancelled() => return Ok(()),
            };

            *liveness.lock().unwrap() = std::time::Instant::now();

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

            for event in resp.events() {
                if let Some(kv) = event.kv() {
                    if !kv.key().starts_with(prefix.as_bytes()) {
                        continue;
                    }
                    let node = match std::str::from_utf8(&kv.key()[prefix.len()..]) {
                        Ok(k) => k.to_owned(),
                        Err(e) => {
                            error!(
                                "Invalid UTF-8 key in watch event, key={} err={:?}",
                                hex::encode(kv.key()),
                                e
                            );
                            continue;
                        }
                    };
                    match event.event_type() {
                        etcd_client::EventType::Put => {
                            info!("[discovery] Added node={}", node);
                            nodes.write().await.insert(node);
                        }
                        etcd_client::EventType::Delete => {
                            info!("[discovery] Deleted node={}", node);
                            nodes.write().await.remove(&node);
                        }
                    }
                }
            }
        }
    }

    /// Refresh the nodes set by querying etcd directly.
    /// Returns Err on failure so the caller can abort the session.
    async fn refresh_nodes(
        prefix: &str,
        client: &mut etcd_client::Client,
        nodes: &RwLock<HashSet<String>>,
    ) -> anyhow::Result<()> {
        let option = Some(GetOptions::new().with_prefix());
        match client.get(prefix, option).await {
            Ok(resp) => {
                let mut guard = nodes.write().await;
                guard.clear();
                for kv in resp.kvs() {
                    if let Ok(key) = std::str::from_utf8(kv.key()) {
                        if let Some(node) = key.strip_prefix(prefix) {
                            guard.insert(node.to_string());
                        }
                    }
                }
                Ok(())
            }
            Err(e) => Err(anyhow::anyhow!("Failed to refresh nodes: {}", e)),
        }
    }

    async fn refresh_nodes_if_watch_unhealthy(&self) -> Result<(), anyhow::Error> {
        let watch_stale =
            self.last_watch_response.lock().unwrap().elapsed() > WATCH_LIVENESS_THRESHOLD;
        if self.supervisor.is_session_active() && !watch_stale {
            return Ok(());
        }

        let mut client = self.client.get_client().await?;
        Self::refresh_nodes(self.prefix, &mut client, &self.nodes).await
    }

    async fn create_registration(
        &self,
        service_addr: &str,
        service_info: &str,
    ) -> Result<RegistrationLease, anyhow::Error> {
        let mut client = self.client.get_client().await?;
        let lease = Arc::new(EtcdLease::new(
            Arc::new(RwLock::new(client.clone())),
            self.registration_keep_alive_interval,
        ));
        lease.grant(self.registration_lease_ttl).await?;

        let key = format!("{}{}", self.prefix, service_addr);
        let lease_id = lease.lease_id.load(Ordering::SeqCst);
        let put_res = client
            .put(
                key.clone(),
                service_info,
                Some(PutOptions::new().with_lease(lease_id)),
            )
            .await;
        if let Err(e) = put_res {
            let _ = lease.close().await;
            return Err(e.into());
        }

        let keep_alive_lease = lease.clone();
        let keep_alive_handle = tokio::spawn(async move {
            keep_alive_lease.keep_alive().await;
        });

        Ok(RegistrationLease {
            key,
            lease,
            keep_alive_handle,
        })
    }

    async fn stop_registration(&mut self) -> Result<(), anyhow::Error> {
        let Some(registration) = self.registration.take() else {
            return Ok(());
        };

        let RegistrationLease {
            key,
            lease,
            keep_alive_handle,
        } = registration;

        keep_alive_handle.abort();
        let _ = keep_alive_handle.await;

        let close_result = lease.close().await;

        if let Ok(mut client) = self.client.get_client().await {
            let _ = client.delete(key, None).await;
        }

        close_result
    }
}

#[async_trait]
impl DiscoveryClient for EtcdDiscoveryClient {
    async fn register(
        &mut self,
        service_addr: &str,
        service_info: &str,
    ) -> Result<(), anyhow::Error> {
        if self.registration.is_some() {
            return Err(anyhow::anyhow!("Service already registered"));
        }

        // Ensure client is started
        self.client.start().await?;

        let registration = self.create_registration(service_addr, service_info).await?;
        self.registration = Some(registration);

        // Start the watch supervisor (handles restarts automatically)
        self.start_watch_supervisor().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: stop supervisor and delete the registered key.
            self.supervisor.stop().await;
            let _ = self.stop_registration().await;
            return Err(e);
        }

        info!(
            "Registered service, addr={}, info={}",
            service_addr, service_info
        );
        Ok(())
    }

    async fn deregister(&mut self, service_addr: &str) -> Result<(), anyhow::Error> {
        if self.registration.is_none() {
            return Ok(());
        }

        let stop_registration_result = self.stop_registration().await;

        // remove from local nodes set
        self.nodes.write().await.remove(service_addr);

        info!("Deregistered service, addr={}", service_addr);
        stop_registration_result
    }

    async fn stop_watch(&mut self) -> Result<(), anyhow::Error> {
        // Stop lease keepalive and remove the registration key first.
        let stop_registration_result = self.stop_registration().await;

        // 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;
        }

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

    async fn get_all_peers(&self) -> Result<AllPeers, anyhow::Error> {
        self.refresh_nodes_if_watch_unhealthy().await?;

        let mut res = AllPeers {
            nodes: self.nodes.read().await.clone().into_iter().collect(),
            leader: None,
        };

        let mut client = self.client.get_client().await?;
        if let Ok(resp) = client.leader(META_SERVICE_ELECTION_KEY).await {
            if let Some(kv) = resp.kv() {
                let val = std::str::from_utf8(kv.value())
                    .unwrap_or_default()
                    .to_owned();
                if res.nodes.contains(&val) {
                    res.leader = Some(val);
                }
            }
        };
        Ok(res)
    }
}

#[cfg(test)]
mod tests {
    use tokio::time::sleep;

    use super::*;
    use crate::test_utils::setup_etcd;

    #[tokio::test(flavor = "multi_thread")]
    async fn test_register_and_deregister() -> Result<(), anyhow::Error> {
        let endpoint = match setup_etcd().await {
            Ok(endpoint) => endpoint,
            Err(e) => {
                eprint!("Failed to set up etcd, err={}", e);
                return Ok(());
            }
        };

        let etcd_client = Arc::new(EtcdClient::new_for_test(vec![endpoint]));
        let prefix = "/tici/test_register_and_deregister/registry/";

        // Clean up any existing test data
        etcd_client.start().await?;
        let mut client = etcd_client.get_client().await?;
        client
            .delete(
                prefix,
                Some(etcd_client::DeleteOptions::new().with_prefix()),
            )
            .await?;

        let mut discovery = EtcdDiscoveryClient::new(
            prefix,
            etcd_client.clone(),
            Duration::from_secs(10),
            Duration::from_secs(3),
        );

        // Register a service
        discovery.register("127.0.0.1:8080", "service_info").await?;

        // Wait a bit for the watch to pick up the change
        sleep(Duration::from_millis(100)).await;

        // Check if the service is discovered
        let peers = discovery.get_all_peers().await?;
        assert_eq!(peers.nodes.len(), 1);
        assert!(peers.nodes.contains(&"127.0.0.1:8080".to_string()));

        // Clean up
        discovery.deregister("127.0.0.1:8080").await?;
        sleep(Duration::from_millis(100)).await;

        let peers = discovery.get_all_peers().await?;
        assert!(!peers.nodes.contains(&"127.0.0.1:8080".to_string()));
        Ok(())
    }

    #[tokio::test(flavor = "multi_thread")]
    async fn test_multiple_services() -> Result<(), anyhow::Error> {
        let endpoint = match setup_etcd().await {
            Ok(endpoint) => endpoint,
            Err(e) => {
                eprint!("Failed to set up etcd, err={}", e);
                return Ok(());
            }
        };

        let prefix = "/tici/test_multiple_services/registry/";
        let etcd_client = Arc::new(EtcdClient::new_for_test(vec![endpoint]));

        // Clean up any existing test data
        etcd_client.start().await?;
        let mut client = etcd_client.get_client().await?;
        client
            .delete(
                prefix,
                Some(etcd_client::DeleteOptions::new().with_prefix()),
            )
            .await?;

        let mut discovery1 = EtcdDiscoveryClient::new(
            prefix,
            etcd_client.clone(),
            Duration::from_secs(10),
            Duration::from_secs(3),
        );
        // discovery1 registers
        discovery1.register("127.0.0.1:8082", "service1").await?;
        assert!(
            discovery1
                .register("127.0.0.1:8083", "service2")
                .await
                .is_err()
        );
        sleep(Duration::from_millis(100)).await;

        let mut discovery2 = EtcdDiscoveryClient::new(
            prefix,
            etcd_client.clone(),
            Duration::from_secs(10),
            Duration::from_secs(3),
        );
        // discovery2 registers
        discovery2.register("127.0.0.1:8084", "service1").await?;
        assert!(
            discovery2
                .register("127.0.0.1:8085", "service2")
                .await
                .is_err()
        );
        sleep(Duration::from_millis(100)).await;

        // Check all services are discovered
        {
            let peers = discovery1.get_all_peers().await?;
            assert_eq!(peers.nodes.len(), 2);
            assert!(peers.nodes.contains(&"127.0.0.1:8082".to_string()));
            assert!(peers.nodes.contains(&"127.0.0.1:8084".to_string()));
        }
        {
            let peers = discovery2.get_all_peers().await?;
            assert_eq!(peers.nodes.len(), 2);
            assert!(peers.nodes.contains(&"127.0.0.1:8082".to_string()));
            assert!(peers.nodes.contains(&"127.0.0.1:8084".to_string()));
        }

        // Clean up
        discovery1.deregister("127.0.0.1:8082").await?;
        {
            let peers = discovery1.get_all_peers().await?;
            assert_eq!(peers.nodes.len(), 1);
            assert!(peers.nodes.contains(&"127.0.0.1:8084".to_string()));
        }
        {
            let peers = discovery2.get_all_peers().await?;
            assert_eq!(peers.nodes.len(), 1);
            assert!(peers.nodes.contains(&"127.0.0.1:8084".to_string()));
        }
        discovery2.deregister("127.0.0.1:8084").await?;
        {
            let peers = discovery1.get_all_peers().await?;
            assert!(peers.nodes.is_empty());
        }
        {
            let peers = discovery2.get_all_peers().await?;
            assert!(peers.nodes.is_empty());
        }
        Ok(())
    }

    #[tokio::test(flavor = "multi_thread")]
    async fn test_watch_updates() -> Result<(), anyhow::Error> {
        let endpoint = match setup_etcd().await {
            Ok(endpoint) => endpoint,
            Err(e) => {
                eprint!("Failed to set up etcd, err={}", e);
                return Ok(());
            }
        };

        let prefix = "/tici/test_watch_updates/registry/";
        let etcd_client = Arc::new(EtcdClient::new_for_test(vec![endpoint]));

        // Clean up any existing test data
        etcd_client.start().await?;
        let mut client = etcd_client.get_client().await?;
        client
            .delete(
                prefix,
                Some(etcd_client::DeleteOptions::new().with_prefix()),
            )
            .await?;

        let mut discovery = EtcdDiscoveryClient::new(
            prefix,
            etcd_client.clone(),
            Duration::from_secs(10),
            Duration::from_secs(3),
        );

        // Register starts the watch supervisor automatically
        discovery.register("127.0.0.1:8090", "initial").await?;
        sleep(Duration::from_millis(200)).await;

        // Simulate another service joining the cluster via etcd directly
        let key = format!("{}{}", prefix, "127.0.0.1:8085");
        client.put(key, "service_info", None).await?;
        sleep(Duration::from_millis(200)).await;

        // Check the watch picked up the external registration
        let peers = discovery.get_all_peers().await?;
        assert!(peers.nodes.contains(&"127.0.0.1:8085".to_string()));

        // Clean up
        discovery.stop_watch().await?;
        Ok(())
    }

    /// Poll a condition until it becomes true, or fail 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;
            }
            sleep(Duration::from_millis(100)).await;
        }
    }

    // Regression test: when the watch stream breaks (etcd restart, network
    // partition, etc.), new node registrations should still be detected after
    // the watch automatically recovers.
    #[tokio::test(flavor = "multi_thread")]
    async fn test_discovery_watch_recovers_after_break() -> Result<(), anyhow::Error> {
        let endpoint = match setup_etcd().await {
            Ok(e) => e,
            Err(_) => {
                eprintln!(
                    "SKIPPED: ETCD_ADDR not set. \
                     Run with ETCD_ADDR=http://127.0.0.1:2379 to enable this test."
                );
                return Ok(());
            }
        };

        let prefix = "/tici/test_watch_recovery/registry/";
        let etcd_client = Arc::new(EtcdClient::new_for_test(vec![endpoint]));

        // Clean up any existing test data
        etcd_client.start().await?;
        let mut client = etcd_client.get_client().await?;
        client
            .delete(
                prefix,
                Some(etcd_client::DeleteOptions::new().with_prefix()),
            )
            .await?;

        let mut discovery = EtcdDiscoveryClient::new(
            prefix,
            etcd_client.clone(),
            Duration::from_secs(10),
            Duration::from_secs(3),
        );
        discovery.register("127.0.0.1:9001", "service1").await?;

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

        // Register a new node directly via etcd (bypassing discovery API)
        let key = format!("{}127.0.0.1:9002", prefix);
        client.put(key, "service2", None).await?;

        // Poll until the supervisor recovers and detects the new node.
        let nodes = discovery.nodes.clone();
        let found = wait_until(Duration::from_secs(10), || {
            let nodes = nodes.clone();
            async move { nodes.read().await.contains("127.0.0.1:9002") }
        })
        .await;
        assert!(
            found,
            "new node should be detected after watch recovery, got nodes: {:?}",
            discovery.nodes.read().await
        );

        // Cleanup
        discovery.stop_watch().await.ok();
        client
            .delete(
                prefix,
                Some(etcd_client::DeleteOptions::new().with_prefix()),
            )
            .await?;
        Ok(())
    }

    // When a node is deleted during the watch gap, the supervisor should
    // detect it via refresh_nodes() on session restart.
    #[tokio::test(flavor = "multi_thread")]
    async fn test_discovery_watch_recovers_missed_delete() -> Result<(), anyhow::Error> {
        let endpoint = match setup_etcd().await {
            Ok(e) => e,
            Err(_) => {
                eprintln!(
                    "SKIPPED: ETCD_ADDR not set. \
                     Run with ETCD_ADDR=http://127.0.0.1:2379 to enable this test."
                );
                return Ok(());
            }
        };

        let prefix = "/tici/test_watch_missed_delete/registry/";
        let etcd_client = Arc::new(EtcdClient::new_for_test(vec![endpoint]));

        // Clean up any existing test data
        etcd_client.start().await?;
        let mut client = etcd_client.get_client().await?;
        client
            .delete(
                prefix,
                Some(etcd_client::DeleteOptions::new().with_prefix()),
            )
            .await?;

        // Register two nodes
        let mut discovery = EtcdDiscoveryClient::new(
            prefix,
            etcd_client.clone(),
            Duration::from_secs(10),
            Duration::from_secs(3),
        );
        discovery.register("127.0.0.1:9010", "service1").await?;

        let key2 = format!("{}127.0.0.1:9011", prefix);
        client.put(key2.clone(), "service2", None).await?;

        // Wait for watch to pick up the second node
        let nodes = discovery.nodes.clone();
        let found = wait_until(Duration::from_secs(5), || {
            let nodes = nodes.clone();
            async move { nodes.read().await.contains("127.0.0.1:9011") }
        })
        .await;
        assert!(found, "second node should be detected via watch");

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

        // Delete node 9011 during the watch gap
        client.delete(key2, None).await?;

        // Poll until the supervisor recovers and the deleted node is gone
        let nodes = discovery.nodes.clone();
        let removed = wait_until(Duration::from_secs(10), || {
            let nodes = nodes.clone();
            async move { !nodes.read().await.contains("127.0.0.1:9011") }
        })
        .await;
        assert!(
            removed,
            "deleted node should be removed after watch recovery, got nodes: {:?}",
            discovery.nodes.read().await
        );
        assert!(discovery.nodes.read().await.contains("127.0.0.1:9010"));

        // Cleanup
        discovery.stop_watch().await.ok();
        client
            .delete(
                prefix,
                Some(etcd_client::DeleteOptions::new().with_prefix()),
            )
            .await?;
        Ok(())
    }

    #[tokio::test(flavor = "multi_thread")]
    async fn test_register_attaches_lease_and_expires_without_keepalive()
    -> Result<(), anyhow::Error> {
        let endpoint = match setup_etcd().await {
            Ok(endpoint) => endpoint,
            Err(e) => {
                eprint!("Failed to set up etcd, err={}", e);
                return Ok(());
            }
        };

        let prefix = "/tici/test_registration_lease/registry/";
        let etcd_client = Arc::new(EtcdClient::new_for_test(vec![endpoint]));

        etcd_client.start().await?;
        let mut client = etcd_client.get_client().await?;
        client
            .delete(
                prefix,
                Some(etcd_client::DeleteOptions::new().with_prefix()),
            )
            .await?;

        let lease_ttl = Duration::from_secs(2);
        let mut discovery = EtcdDiscoveryClient::new(
            prefix,
            etcd_client.clone(),
            lease_ttl,
            Duration::from_millis(200),
        );
        discovery.register("127.0.0.1:9100", "service_info").await?;

        let key = format!("{}{}", prefix, "127.0.0.1:9100");
        let resp = client.get(key.clone(), None).await?;
        let kv = resp.kvs().first().expect("registered key should exist");
        assert_ne!(
            kv.lease(),
            0,
            "registered key should be attached to a lease"
        );

        if let Some(registration) = discovery.registration.as_mut() {
            registration.keep_alive_handle.abort();
        }

        // PD/etcd may remove an expired lease key after the nominal TTL. Keep
        // this as a test wait bound instead of relying on exact expiry timing.
        let expired = wait_until(lease_ttl + Duration::from_secs(13), || {
            let etcd_client = etcd_client.clone();
            let key = key.clone();
            async move {
                let Ok(mut client) = etcd_client.get_client().await else {
                    return false;
                };
                match client.get(key, None).await {
                    Ok(resp) => resp.kvs().is_empty(),
                    Err(_) => false,
                }
            }
        })
        .await;
        assert!(
            expired,
            "registration key should expire after keepalive stops"
        );

        discovery.stop_watch().await.ok();
        Ok(())
    }

    #[tokio::test(flavor = "multi_thread")]
    async fn test_get_all_peers_refreshes_when_watch_is_stale() -> Result<(), anyhow::Error> {
        let endpoint = match setup_etcd().await {
            Ok(endpoint) => endpoint,
            Err(e) => {
                eprint!("Failed to set up etcd, err={}", e);
                return Ok(());
            }
        };

        let prefix = "/tici/test_stale_watch_refresh/registry/";
        let etcd_client = Arc::new(EtcdClient::new_for_test(vec![endpoint]));

        etcd_client.start().await?;
        let mut client = etcd_client.get_client().await?;
        client
            .delete(
                prefix,
                Some(etcd_client::DeleteOptions::new().with_prefix()),
            )
            .await?;

        let mut discovery = EtcdDiscoveryClient::new(
            prefix,
            etcd_client.clone(),
            Duration::from_secs(10),
            Duration::from_secs(3),
        );
        discovery.register("127.0.0.1:9110", "service1").await?;

        let key = format!("{}{}", prefix, "127.0.0.1:9111");
        client.put(key, "service2", None).await?;
        discovery.nodes.write().await.remove("127.0.0.1:9111");
        assert!(!discovery.nodes.read().await.contains("127.0.0.1:9111"));

        *discovery.last_watch_response.lock().unwrap() =
            std::time::Instant::now() - WATCH_LIVENESS_THRESHOLD - Duration::from_secs(1);

        let peers = discovery.get_all_peers().await?;
        assert!(peers.nodes.contains(&"127.0.0.1:9111".to_string()));

        discovery.stop_watch().await?;
        Ok(())
    }
}