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

use async_trait::async_trait;
use etcd_client::{LeaderKey, ResignOptions};
use tokio::sync::RwLock;

use super::{client::ElectionClient, etcd_lease::EtcdLease};
use crate::etcd_client::EtcdClient;

/// High-Availability etcd election client using etcd client
#[derive(Clone)]
pub struct EtcdElectionClient {
    // The key used for leader election in etcd.
    // All clients will use the same key to campaign for leadership.
    // Store `leader_key` here to allow run unittest concurrently.
    leader_key: &'static str,
    leader_value: String,
    lease_keep_alive_interval: Duration,
    // TODO: simplify the lease management
    lease: Arc<RwLock<Option<Arc<EtcdLease>>>>,
    client: Arc<EtcdClient>,
}

impl EtcdElectionClient {
    // Create a new etcd election client with the given client.
    pub fn new(
        leader_key: &'static str,
        leader_value: String,
        client: Arc<EtcdClient>,
        lease_keep_alive_interval: Duration,
    ) -> EtcdElectionClient {
        EtcdElectionClient {
            leader_key,
            leader_value,
            lease_keep_alive_interval,
            lease: Arc::new(RwLock::new(None)),
            client,
        }
    }
}

#[async_trait]
impl ElectionClient for EtcdElectionClient {
    // Try to campaign as leader.
    // 1. If already be leader, lease keep alive, block until not be leader.
    // 2. If not leader, try to be the leader, block until be the leader.
    // Return leader lease id if it is leader.
    async fn campaign(&self, lease_timeout: Duration) -> Result<Option<i64>, anyhow::Error> {
        // Ensure client is started
        self.client.start().await?;

        if self.is_leader().await? {
            // Already leader, lease keep alive
            let lease = self.lease.read().await.clone();
            if let Some(lease) = lease {
                lease.keep_alive().await;
                let mut guard = self.lease.write().await;
                if guard
                    .as_ref()
                    .is_some_and(|current| Arc::ptr_eq(current, &lease))
                {
                    *guard = None;
                }
            }
            return Ok(None);
        }

        // Not leader, try to campaign
        let mut client = self.client.get_client().await?;
        let new_lease = Arc::new(EtcdLease::new(
            Arc::new(RwLock::new(client.clone())),
            self.lease_keep_alive_interval,
        ));
        new_lease.grant(lease_timeout).await?;
        // lease granted from etcd, set it to self.lease
        *self.lease.write().await = Some(new_lease.clone());
        let lease_id = new_lease.lease_id.load(Ordering::SeqCst);
        // campaign the leader of `leader_key` by granted lease
        match client
            .campaign(self.leader_key, self.leader_value.as_str(), lease_id)
            .await
        {
            Ok(resp) => {
                if resp.leader().is_none() {
                    new_lease.close().await?;
                    Err(etcd_client::Error::ElectError("No leader was elected.".into()).into())
                } else {
                    // become leader successfully
                    Ok(Some(lease_id))
                }
            }
            Err(e) => {
                new_lease.close().await?;
                Err(etcd_client::Error::ElectError(format!("Failed to campaign, err={}", e)).into())
            }
        }
    }

    // Is leader or not.
    async fn is_leader(&self) -> Result<bool, anyhow::Error> {
        let lease = self.lease.read().await.clone();
        let lease_id: i64 = match lease {
            Some(l) => {
                if l.is_expired().await {
                    0
                } else {
                    l.lease_id.load(Ordering::SeqCst)
                }
            }
            None => 0,
        };
        if lease_id == 0 {
            return Ok(false);
        }

        let mut client = self.client.get_client().await?;
        match client.leader(self.leader_key).await {
            Ok(resp) => {
                if let Some(kv) = resp.kv() {
                    Ok(kv.value() == self.leader_value.as_bytes() && kv.lease() == lease_id)
                } else {
                    Err(anyhow::anyhow!("Unexpected empty response"))
                }
            }
            Err(e) => {
                if let etcd_client::Error::GRpcStatus(ref status) = e {
                    if status.message() == "election: no leader" {
                        return Ok(false);
                    }
                }
                Err(anyhow::anyhow!("Unexpected err={}", e))
            }
        }
    }

    // Releases election leadership.
    async fn resign(&self) -> Result<(), anyhow::Error> {
        let lease = self.lease.read().await.clone();
        let lease_id: i64 = match lease {
            Some(l) => match l.is_expired().await {
                true => 0,
                false => l.lease_id.load(Ordering::SeqCst),
            },
            None => 0,
        };
        if lease_id == 0 {
            return Ok(());
        }

        let leader = LeaderKey::new()
            .with_name(self.leader_key)
            .with_key(format!("{}/{}", self.leader_key, lease_id))
            .with_lease(lease_id);
        let resign_option = ResignOptions::new().with_leader(leader);

        let mut client = self.client.get_client().await?;
        client.resign(Some(resign_option)).await?;

        let lease = self.lease.read().await.clone();
        if let Some(l) = lease {
            l.close().await?;
        }
        Ok(())
    }
}

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

    use tokio::time;

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

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

    #[tokio::test]
    async fn test_campaign_and_resign() -> 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 leader_key = "test-ha-election-key";
        let leader_value = "test-ha-election-value";

        // Clean up any existing election data
        etcd_client.start().await?;
        let mut client = etcd_client.get_client().await?;
        let _ = client.delete(leader_key, None).await;

        let election_client = EtcdElectionClient::new(
            leader_key,
            leader_value.to_string(),
            etcd_client.clone(),
            TEST_KEEP_ALIVE_INTERVAL,
        );

        // First campaign should succeed
        let is_leader = election_client.campaign(Duration::from_secs(5)).await?;
        assert!(is_leader.is_some());

        // Should be leader now
        assert!(election_client.is_leader().await?);

        // Resign leadership
        election_client.resign().await?;

        // Should no longer be leader
        assert!(!election_client.is_leader().await?);

        Ok(())
    }

    #[tokio::test]
    async fn test_only_one_leader() -> 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 leader_key = "test-election-key-multi";
        let leader_value1 = "test-election-value-1";
        let leader_value2 = "test-election-value-2";

        // Clean up any existing election data
        etcd_client.start().await?;
        let mut client = etcd_client.get_client().await?;
        let _ = client.delete(leader_key, None).await;

        let election_client1 = EtcdElectionClient::new(
            leader_key,
            leader_value1.to_string(),
            etcd_client.clone(),
            TEST_KEEP_ALIVE_INTERVAL,
        );
        let election_client2 = EtcdElectionClient::new(
            leader_key,
            leader_value2.to_string(),
            etcd_client.clone(),
            TEST_KEEP_ALIVE_INTERVAL,
        );

        // First client campaigns
        let is_leader1 = election_client1.campaign(Duration::from_secs(5)).await?;
        assert!(is_leader1.is_some());
        assert!(election_client1.is_leader().await?);
        let start = Instant::now();

        // Second client campaigns - should block until client1 lease timeout
        let is_leader2 = election_client2.campaign(Duration::from_secs(5)).await?;
        assert!(is_leader2.is_some());
        let elapse = Instant::now() - start;
        assert!(elapse.as_secs() >= 5);
        // but client2 lease also timeout, election_client2 is not leader now.
        assert!(!election_client2.is_leader().await?);
        assert!(!election_client1.is_leader().await?);

        // Now first client campaigns - should return fast.
        let start = Instant::now();
        let is_leader1 = election_client1.campaign(Duration::from_secs(5)).await?;
        assert!(is_leader1.is_some());
        let elapse = Instant::now() - start;
        assert!(elapse.as_secs() < 5);
        assert!(election_client1.is_leader().await?);

        election_client1.resign().await?;

        Ok(())
    }

    #[tokio::test]
    async fn test_lease_expiration() -> 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 leader_key = "test-ha-election-key-expire";
        let leader_value = "test-ha-election-value-expire";

        // Clean up any existing election data
        etcd_client.start().await?;
        let mut client = etcd_client.get_client().await?;
        let _ = client.delete(leader_key, None).await;

        let election_client = EtcdElectionClient::new(
            leader_key,
            leader_value.to_string(),
            etcd_client.clone(),
            TEST_KEEP_ALIVE_INTERVAL,
        );

        // Campaign with very short lease
        let is_leader = election_client.campaign(Duration::from_secs(1)).await?;
        assert!(is_leader.is_some());
        assert!(election_client.is_leader().await?);

        // Wait for lease to expire
        time::sleep(Duration::from_secs(2)).await;

        // Should no longer be leader
        assert!(!election_client.is_leader().await?);

        Ok(())
    }

    #[tokio::test(flavor = "multi_thread")]
    async fn test_keep_alive() -> 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 leader_key = "test-election-key-keepalive";
        let leader_value = "test-election-value-keepalive";

        // Clean up any existing election data
        etcd_client.start().await?;
        let mut client = etcd_client.get_client().await?;
        let _ = client.delete(leader_key, None).await;

        let election_client = EtcdElectionClient::new(
            leader_key,
            leader_value.to_string(),
            etcd_client.clone(),
            TEST_KEEP_ALIVE_INTERVAL,
        );

        // Campaign with short lease but keep alive
        let is_leader = election_client.campaign(Duration::from_secs(2)).await?;
        assert!(is_leader.is_some());

        let client_clone = election_client.clone();
        let handle = tokio::spawn(async move {
            assert!(client_clone.is_leader().await.unwrap());
            let is_leader = client_clone.campaign(Duration::from_secs(2)).await;
            assert!(is_leader.is_ok());
            assert!(is_leader.unwrap().is_none());
        });

        // Should remain leader for longer than lease time
        for _ in 0..4 {
            assert!(election_client.is_leader().await?);
            time::sleep(Duration::from_secs(1)).await;
        }

        election_client.resign().await?;
        tokio::time::timeout(Duration::from_secs(5), handle)
            .await
            .map_err(|_| anyhow::anyhow!("keep_alive did not stop after resign"))?
            .map_err(|err| anyhow::anyhow!("keep_alive task panicked: {}", err))?;

        Ok(())
    }

    #[tokio::test(flavor = "multi_thread")]
    async fn test_keep_alive_return_clears_cached_lease() -> 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 leader_key = "test-election-key-clear-stale-lease";
        let leader_value = "test-election-value-clear-stale-lease";

        etcd_client.start().await?;
        let mut client = etcd_client.get_client().await?;
        let _ = client.delete(leader_key, None).await;

        let election_client = EtcdElectionClient::new(
            leader_key,
            leader_value.to_string(),
            etcd_client.clone(),
            TEST_KEEP_ALIVE_INTERVAL,
        );

        let lease_id = election_client
            .campaign(Duration::from_secs(2))
            .await?
            .ok_or_else(|| anyhow::anyhow!("initial campaign should become leader"))?;
        assert!(election_client.is_leader().await?);

        let keep_alive_client = election_client.clone();
        let keep_alive_handle =
            tokio::spawn(async move { keep_alive_client.campaign(Duration::from_secs(2)).await });

        time::sleep(Duration::from_millis(500)).await;
        election_client.resign().await?;

        let keep_alive_result = tokio::time::timeout(Duration::from_secs(10), keep_alive_handle)
            .await
            .map_err(|_| anyhow::anyhow!("keep_alive campaign did not return after lease revoke"))?
            .map_err(|err| anyhow::anyhow!("keep_alive campaign task panicked: {}", err))??;
        assert!(keep_alive_result.is_none());
        assert!(election_client.lease.read().await.is_none());

        let new_lease_id = election_client
            .campaign(Duration::from_secs(2))
            .await?
            .ok_or_else(|| anyhow::anyhow!("campaign should recover with a new lease"))?;
        assert_ne!(new_lease_id, lease_id);

        election_client.resign().await?;

        Ok(())
    }
}