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

use anyhow::Result;
use etcd_client::Client;
use log::{error, info, warn};
use scopeguard::defer;
use tokio::sync::RwLock;

const REVOKE_LEASE_TIMEOUT: Duration = Duration::from_secs(1);
const SLOW_REQUEST_TIME: Duration = Duration::from_secs(1);

// EtcdLease is used as the low-level mechanism for campaigning and renewing
// elected leadership. The way to gain and maintain leadership is to update and
// keep the lease alive continuously.
#[derive(Clone)]
pub struct EtcdLease {
    client: Arc<RwLock<Client>>,
    pub lease_id: Arc<AtomicI64>,
    lease_timeout: Arc<RwLock<Duration>>,
    expire_time: Arc<RwLock<SystemTime>>,
    keep_alive_interval: Duration,
}

impl EtcdLease {
    pub fn new(client: Arc<RwLock<Client>>, keep_alive_interval: Duration) -> Self {
        EtcdLease {
            client,
            lease_id: Arc::new(AtomicI64::new(0)),
            lease_timeout: Arc::new(RwLock::new(Duration::from_secs(0))),
            expire_time: Arc::new(RwLock::new(SystemTime::UNIX_EPOCH)),
            keep_alive_interval,
        }
    }

    // grant initializes the lease_timeout and expire_time.
    pub async fn grant(&self, lease_timeout: Duration) -> Result<()> {
        // TODO(P1) There is a gap here, because etcd can only handles seconds in
        // integer.
        let ttl_secs = i64::try_from(lease_timeout.as_millis().saturating_add(999) / 1_000)
            .unwrap_or(i64::MAX);
        let start = Instant::now();

        let resp = self
            .client
            .write()
            .await
            .lease_grant(ttl_secs, None)
            .await?;

        if start.elapsed() > SLOW_REQUEST_TIME {
            warn!("lease grants too slow, cost={:?}", start.elapsed(),);
        }

        info!(
            "lease granted, lease_id={}, lease_timeout={}",
            resp.id(),
            ttl_secs
        );

        self.lease_id.store(resp.id(), Ordering::SeqCst);
        *self.lease_timeout.write().await = lease_timeout;
        *self.expire_time.write().await = SystemTime::now() + lease_timeout;

        Ok(())
    }

    // close releases the lease.
    pub async fn close(&self) -> Result<()> {
        *self.expire_time.write().await = SystemTime::UNIX_EPOCH;

        let lease_id = self.lease_id.load(Ordering::SeqCst);
        if lease_id != 0 {
            if let Err(e) = tokio::time::timeout(
                REVOKE_LEASE_TIMEOUT,
                self.client.write().await.lease_revoke(lease_id),
            )
            .await
            {
                error!("revoke lease failed, err={}", e);
            }
        }

        Ok(())
    }

    // is_expired checks if the lease is expired. If it returns true,
    // current leader should step down and try to re-elect again.
    pub async fn is_expired(&self) -> bool {
        let expire_time = *self.expire_time.read().await;
        if expire_time == SystemTime::UNIX_EPOCH {
            return true;
        }
        SystemTime::now() > expire_time
    }

    // keep_alive auto renews the lease and update expire_time.
    pub async fn keep_alive(&self) {
        let lease_timeout = *self.lease_timeout.read().await;
        let interval = self.keep_alive_interval;
        let mut timer = tokio::time::interval(interval);

        let lease_id = self.lease_id.load(Ordering::SeqCst);
        info!(
            "keep_alive started, lease_id={} interval={:?}",
            lease_id, interval
        );
        defer!({
            info!("keep_alive stopped, lease_id={}", lease_id);
        });

        let (mut keeper, mut stream) =
            match self.client.write().await.lease_keep_alive(lease_id).await {
                Ok(res) => res,
                Err(err) => {
                    *self.expire_time.write().await = SystemTime::UNIX_EPOCH;
                    warn!("lease keep alive failed, err={}", err);
                    return;
                }
            };
        loop {
            if self.is_expired().await {
                info!("lease expired, lease_id={}", lease_id);
                return;
            }

            let start = SystemTime::now();
            if let Err(err) = keeper.keep_alive().await {
                *self.expire_time.write().await = SystemTime::UNIX_EPOCH;
                warn!(
                    "lease keep alive failed, lease_id={} start={:?}, err={}",
                    lease_id, start, err
                );
                return;
            }
            match tokio::time::timeout(lease_timeout, stream.message()).await {
                Ok(Ok(Some(resp))) => {
                    if resp.ttl() > 0 {
                        let mut expire_time = self.expire_time.write().await;
                        if *expire_time == SystemTime::UNIX_EPOCH {
                            return;
                        }
                        let expire = start + Duration::from_secs(resp.ttl() as u64);
                        *expire_time = expire;
                    }
                }
                Ok(Ok(None)) => {
                    *self.expire_time.write().await = SystemTime::UNIX_EPOCH;
                    warn!(
                        "lease keep alive stream closed, lease_id={} start={:?}",
                        lease_id, start
                    );
                    return;
                }
                Ok(Err(err)) => {
                    *self.expire_time.write().await = SystemTime::UNIX_EPOCH;
                    error!(
                        "lease keep alive stream recv failed, lease_id={} start={:?}, err={}",
                        lease_id, start, err
                    );
                    return;
                }
                Err(_) => {
                    *self.expire_time.write().await = SystemTime::UNIX_EPOCH;
                    warn!(
                        "lease keep alive stream timed out, lease_id={} start={:?}, timeout={:?}",
                        lease_id, start, lease_timeout
                    );
                    return;
                }
            }

            timer.tick().await;
        }
    }
}

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

    use super::*;
    use crate::test_utils::setup_etcd_client;

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

    #[tokio::test]
    async fn test_lease_grant() {
        let client = match setup_etcd_client().await {
            Some(client) => client,
            None => return,
        };
        let lease = EtcdLease::new(Arc::new(RwLock::new(client)), TEST_KEEP_ALIVE_INTERVAL);

        const TTL: u64 = 10;

        // Test granting a lease
        lease
            .grant(Duration::from_secs(TTL))
            .await
            .expect("grant should succeed");
        assert_ne!(lease.lease_id.load(Ordering::SeqCst), 0);
        assert!(!lease.is_expired().await);

        // Check lease timeout was set correctly
        let timeout = *lease.lease_timeout.read().await;
        assert_eq!(timeout, Duration::from_secs(TTL));

        // Check expire time is in the future
        let expire_time = *lease.expire_time.read().await;
        assert!(expire_time > SystemTime::now());
    }

    #[tokio::test]
    async fn test_lease_close() {
        let client = match setup_etcd_client().await {
            Some(client) => client,
            None => return,
        };
        let lease = EtcdLease::new(Arc::new(RwLock::new(client)), TEST_KEEP_ALIVE_INTERVAL);

        lease
            .grant(Duration::from_secs(10))
            .await
            .expect("grant should succeed");
        let lease_id = lease.lease_id.load(Ordering::SeqCst);
        assert_ne!(lease_id, 0);

        // Close the lease
        lease.close().await.expect("close should succeed");
        assert!(lease.is_expired().await);
        assert_eq!(lease.lease_id.load(Ordering::SeqCst), lease_id); // ID remains but lease is expired
    }

    #[tokio::test]
    async fn test_lease_is_expired() {
        let client = match setup_etcd_client().await {
            Some(client) => client,
            None => return,
        };
        let lease = EtcdLease::new(Arc::new(RwLock::new(client)), TEST_KEEP_ALIVE_INTERVAL);

        // New lease should be considered expired (no lease granted)
        assert!(lease.is_expired().await);

        // After grant, should not be expired
        lease
            .grant(Duration::from_secs(2))
            .await
            .expect("grant should succeed");
        assert!(!lease.is_expired().await);

        // After timeout, should be expired
        tokio::time::sleep(Duration::from_secs(3)).await;
        assert!(lease.is_expired().await);
    }

    #[tokio::test(flavor = "multi_thread")]
    async fn test_keep_alive() {
        let client = match setup_etcd_client().await {
            Some(client) => client,
            None => return,
        };
        let lease = EtcdLease::new(Arc::new(RwLock::new(client)), TEST_KEEP_ALIVE_INTERVAL);

        const TTL: u64 = 5;

        lease
            .grant(Duration::from_secs(TTL))
            .await
            .expect("grant should succeed");
        assert!(!lease.is_expired().await);

        // Start keep alive worker
        let lease_clone = lease.clone();
        let handle = tokio::spawn(async move {
            lease_clone.keep_alive().await;
        });

        // Wait for a timeout
        assert!(!lease.is_expired().await);
        tokio::time::sleep(Duration::from_secs(TTL + 1)).await;
        assert!(!lease.is_expired().await);

        // Close and wait for a timeout.
        lease.close().await.expect("close should succeed");
        // tokio::time::sleep(Duration::from_secs(TTL + 1)).await;
        assert!(lease.is_expired().await);

        handle.await.expect("worker should stop cleanly");
    }

    #[tokio::test]
    async fn test_keep_alive_marks_expired_on_initial_error() {
        let client = match setup_etcd_client().await {
            Some(client) => client,
            None => return,
        };
        let lease = EtcdLease::new(Arc::new(RwLock::new(client)), TEST_KEEP_ALIVE_INTERVAL);

        lease
            .grant(Duration::from_secs(10))
            .await
            .expect("grant should succeed");
        let lease_id = lease.lease_id.load(Ordering::SeqCst);
        assert!(!lease.is_expired().await);

        lease
            .client
            .write()
            .await
            .lease_revoke(lease_id)
            .await
            .expect("revoke should succeed");

        lease.keep_alive().await;
        assert!(lease.is_expired().await);
    }
}