use std::sync::{
    Arc,
    atomic::{AtomicI64, Ordering},
};

use anyhow::Error;
use log::{error, info};
use tokio::{sync::Notify, task::JoinHandle, time::Duration};

use super::client::ElectionClient;

const DEFAULT_CAMPAIGN_INTERVAL: Duration = Duration::from_secs(1);

#[async_trait::async_trait]
pub trait Leadership: Send + Sync {
    async fn stop_campaign(&self) -> Result<(), Error>;

    // Non-blocking function for get leader term
    // If not leader, return None
    // If leader, return Some(leader_term)
    fn try_get_leader_term(&self) -> Option<i64>;
}

pub type LeadershipObserver = Arc<dyn Fn(bool) + Send + Sync>;

#[derive(Clone)]
pub struct EtcdLeadership {
    client: Arc<dyn ElectionClient>,
    // leader lease id, 0 means not leader
    leader_lease_id: Arc<AtomicI64>,
    campaign_handle: Arc<JoinHandle<()>>,
    notify: Arc<Notify>,
    observer: Option<LeadershipObserver>,
}

impl EtcdLeadership {
    pub fn new(client: Arc<dyn ElectionClient>, lease_timeout: Duration) -> Self {
        Self::new_with_observer_and_interval(client, lease_timeout, DEFAULT_CAMPAIGN_INTERVAL, None)
    }

    pub fn new_with_observer(
        client: Arc<dyn ElectionClient>,
        lease_timeout: Duration,
        observer: Option<LeadershipObserver>,
    ) -> Self {
        Self::new_with_observer_and_interval(
            client,
            lease_timeout,
            DEFAULT_CAMPAIGN_INTERVAL,
            observer,
        )
    }

    pub fn new_with_observer_and_interval(
        client: Arc<dyn ElectionClient>,
        lease_timeout: Duration,
        campaign_interval: Duration,
        observer: Option<LeadershipObserver>,
    ) -> Self {
        let leader_lease_id = Arc::new(AtomicI64::new(0));
        let notify = Arc::new(Notify::new());
        let observer_for_loop = observer.clone();
        let handle = tokio::spawn(Self::run_campaign_loop(
            client.clone(),
            lease_timeout,
            campaign_interval,
            leader_lease_id.clone(),
            notify.clone(),
            observer_for_loop,
        ));

        Self {
            client,
            leader_lease_id,
            notify,
            campaign_handle: Arc::new(handle),
            observer,
        }
    }

    async fn run_campaign_loop(
        client: Arc<dyn ElectionClient>,
        lease_timeout: Duration,
        campaign_interval: Duration,
        leader_lease_id: Arc<AtomicI64>,
        notify: Arc<Notify>,
        observer: Option<LeadershipObserver>,
    ) {
        let mut interval = tokio::time::interval(campaign_interval);
        loop {
            interval.tick().await;
            // start campaign
            match client.campaign(lease_timeout).await {
                Ok(Some(lease_id)) => {
                    leader_lease_id.store(lease_id, Ordering::SeqCst);
                    info!("Become leader with lease_id={}", lease_id);
                }
                Ok(None) => {
                    if leader_lease_id.swap(0, Ordering::SeqCst) != 0 {
                        info!("Leadership lost, reset local leader lease id");
                    }
                    info!("Campaign returned none");
                }
                Err(e) => {
                    if leader_lease_id.swap(0, Ordering::SeqCst) != 0 {
                        info!("Leadership reset after campaign error");
                    }
                    error!("Campaign err={:?}", e);
                }
            }
            notify.notify_waiters();
            if let Some(observer) = observer.as_ref() {
                let is_leader = leader_lease_id.load(Ordering::SeqCst) != 0;
                observer(is_leader);
            }
        }
    }
}

#[async_trait::async_trait]
impl Leadership for EtcdLeadership {
    async fn stop_campaign(&self) -> Result<(), Error> {
        info!("leadership campaign is stopping...");
        // Abort the campaign task if it is running
        self.campaign_handle.abort();

        // Reset the leader_lease_id so that this node is not considered as leader
        // locally
        self.leader_lease_id.store(0, Ordering::SeqCst);
        self.notify.notify_waiters();
        if let Some(observer) = self.observer.as_ref() {
            observer(false);
        }

        // Resign from etcd election
        self.client.resign().await?;
        info!("leadership campaign stopped, resigned from election");
        Ok(())
    }

    fn try_get_leader_term(&self) -> Option<i64> {
        fail::fail_point!("mock_etcd_leader_term", |_| Some(1));
        let lease_id = self.leader_lease_id.load(Ordering::SeqCst);
        if lease_id == 0 { None } else { Some(lease_id) }
    }
}

pub async fn block_wait_leader_term(leadership: &EtcdLeadership, interval: Duration) -> i64 {
    let mut timer = tokio::time::interval(interval);
    loop {
        tokio::select! {
            _ = timer.tick() => {
                if let Some(lease_id) = leadership.try_get_leader_term() {
                    info!("Get leader term success, lease_id={}", lease_id);
                    return lease_id;
                }
                info!("Not leader after interval {:?}, retrying...", interval);
                // continue waiting
            }
            _ = leadership.notify.notified() => {
                if let Some(lease_id) = leadership.try_get_leader_term() {
                    info!("Get leader term success, lease_id={}", lease_id);
                    return lease_id;
                }
                info!("Not leader yet, continue waiting...");
                // continue waiting
            }
        }
    }
}

#[derive(Default)]
pub struct MockLeadership {}

#[async_trait::async_trait]
impl Leadership for MockLeadership {
    async fn stop_campaign(&self) -> Result<(), Error> {
        Ok(())
    }

    fn try_get_leader_term(&self) -> Option<i64> {
        Some(1)
    }
}

#[cfg(test)]
mod tests {
    use std::{collections::VecDeque, sync::Arc, time::Duration};

    use anyhow::anyhow;
    use async_trait::async_trait;
    use tokio::sync::Mutex;

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

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

    #[derive(Default)]
    struct ScriptedElectionClient {
        campaigns: Mutex<VecDeque<Result<Option<i64>, Error>>>,
    }

    impl ScriptedElectionClient {
        fn new(campaigns: Vec<Result<Option<i64>, Error>>) -> Self {
            Self {
                campaigns: Mutex::new(campaigns.into()),
            }
        }
    }

    #[async_trait]
    impl ElectionClient for ScriptedElectionClient {
        async fn campaign(&self, _lease_timeout: Duration) -> Result<Option<i64>, Error> {
            self.campaigns.lock().await.pop_front().unwrap_or(Ok(None))
        }

        async fn is_leader(&self) -> Result<bool, Error> {
            Ok(false)
        }

        async fn resign(&self) -> Result<(), Error> {
            Ok(())
        }
    }

    async fn wait_for_term(leadership: &EtcdLeadership, expected: Option<i64>) {
        let deadline = tokio::time::Instant::now() + Duration::from_secs(2);
        loop {
            if leadership.try_get_leader_term() == expected {
                return;
            }
            assert!(
                tokio::time::Instant::now() < deadline,
                "timed out waiting for leader term {expected:?}, got {:?}",
                leadership.try_get_leader_term()
            );
            tokio::time::sleep(Duration::from_millis(10)).await;
        }
    }

    async fn wait_for_any_term(leadership: &EtcdLeadership) -> i64 {
        let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
        loop {
            if let Some(term) = leadership.try_get_leader_term() {
                return term;
            }
            assert!(
                tokio::time::Instant::now() < deadline,
                "timed out waiting for leadership to become leader"
            );
            tokio::time::sleep(Duration::from_millis(10)).await;
        }
    }

    async fn wait_for_new_term(leadership: &EtcdLeadership, old_term: i64) -> i64 {
        let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
        loop {
            if let Some(term) = leadership.try_get_leader_term() {
                if term != old_term {
                    return term;
                }
            }
            assert!(
                tokio::time::Instant::now() < deadline,
                "timed out waiting for new leader term after {old_term}, got {:?}",
                leadership.try_get_leader_term()
            );
            tokio::time::sleep(Duration::from_millis(10)).await;
        }
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn test_campaign_loop_resets_local_term_on_none() {
        let client = Arc::new(ScriptedElectionClient::new(vec![Ok(Some(7)), Ok(None)]));
        let leadership = EtcdLeadership::new(client, Duration::from_secs(30));

        wait_for_term(&leadership, Some(7)).await;

        wait_for_term(&leadership, None).await;

        leadership.stop_campaign().await.unwrap();
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn test_campaign_loop_resets_local_term_on_error() {
        let client = Arc::new(ScriptedElectionClient::new(vec![
            Ok(Some(11)),
            Err(anyhow!("boom")),
        ]));
        let leadership = EtcdLeadership::new(client, Duration::from_secs(30));

        wait_for_term(&leadership, Some(11)).await;

        wait_for_term(&leadership, None).await;

        leadership.stop_campaign().await.unwrap();
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn test_real_etcd_leadership_resets_local_term_after_keep_alive_failure() {
        let endpoint = match setup_etcd().await {
            Ok(endpoint) => endpoint,
            Err(e) => {
                eprintln!("Failed to set up etcd, err={}", e);
                return;
            }
        };

        let etcd_client = Arc::new(EtcdClient::new_for_test(vec![endpoint]));
        let leader_key = "test-leadership-reset-after-keepalive-failure";
        let leader_value = "127.0.0.1:19008";

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

        let election_client = Arc::new(EtcdElectionClient::new(
            leader_key,
            leader_value.to_string(),
            etcd_client.clone(),
            TEST_KEEP_ALIVE_INTERVAL,
        ));
        let leadership = EtcdLeadership::new(election_client.clone(), Duration::from_secs(2));

        let lease_id = wait_for_any_term(&leadership).await;

        // Wait for the next campaign tick so the leader is in the keep-alive path,
        // then close the current lease to force leadership loss.
        tokio::time::sleep(Duration::from_millis(1500)).await;
        election_client.resign().await.unwrap();

        let recovered_lease_id = wait_for_new_term(&leadership, lease_id).await;
        assert_ne!(recovered_lease_id, lease_id);

        leadership.stop_campaign().await.unwrap();
    }
}