use std::{
    future::Future,
    pin::Pin,
    sync::{
        Arc,
        atomic::{AtomicBool, Ordering},
    },
    time::Duration,
};

use anyhow::Result;
use log::{info, warn};
use tokio::{sync::RwLock, task::JoinHandle};
use tokio_util::sync::CancellationToken;

pub const MIN_WATCH_RETRY_INTERVAL: Duration = Duration::from_secs(1);
pub const MAX_WATCH_RETRY_INTERVAL: Duration = Duration::from_secs(10);

/// A boxed async function that runs one watch session.
///
/// Called repeatedly by the supervisor. Should return `Ok(())` when the
/// session ends normally (stream closed) or `Err` on failure. The
/// supervisor restarts in either case, with exponential backoff on errors.
pub type SessionFn = Box<
    dyn Fn(CancellationToken) -> Pin<Box<dyn Future<Output = Result<()>> + Send>> + Send + Sync,
>;

/// Supervises a watch session with automatic restart and exponential backoff.
///
/// Used by both `LeaderTerm` (election watch) and `EtcdDiscoveryClient`
/// (discovery watch) to maintain resilient etcd watch connections.
pub struct WatchSupervisor {
    name: &'static str,
    shutdown: CancellationToken,
    handle: RwLock<Option<JoinHandle<()>>>,
    /// True while the session_fn is executing; false during backoff and
    /// before the first session starts.
    session_active: Arc<AtomicBool>,
}

impl WatchSupervisor {
    pub fn new(name: &'static str) -> Self {
        Self {
            name,
            shutdown: CancellationToken::new(),
            handle: RwLock::new(None),
            session_active: Arc::new(AtomicBool::new(false)),
        }
    }

    pub fn is_stopped(&self) -> bool {
        self.shutdown.is_cancelled()
    }

    /// Returns true only while a watch session has established its watch
    /// and finished gap-refresh. False during setup, backoff, or when stopped.
    /// The session_fn is responsible for setting this via the flag from
    /// `session_active_flag()` after it has created the watch and refreshed.
    pub fn is_session_active(&self) -> bool {
        self.session_active.load(Ordering::Acquire)
    }

    /// Returns the shared flag that the session_fn should set to `true`
    /// once it has established the watch and completed any gap-refresh.
    /// The supervisor automatically resets it to `false` when the session
    /// returns (either Ok or Err).
    pub fn session_active_flag(&self) -> Arc<AtomicBool> {
        self.session_active.clone()
    }

    /// Start the supervisor with the given session factory.
    /// Stops any previous instance first.
    pub async fn start(&mut self, session_fn: SessionFn) {
        self.stop_inner().await;

        if self.shutdown.is_cancelled() {
            self.shutdown = CancellationToken::new();
        }

        self.session_active.store(false, Ordering::Release);

        let shutdown = self.shutdown.clone();
        let name = self.name;
        let session_active = self.session_active.clone();

        let handle = tokio::spawn(async move {
            let mut backoff = MIN_WATCH_RETRY_INTERVAL;

            loop {
                if shutdown.is_cancelled() {
                    return;
                }

                // Note: session_active is NOT set here. The session_fn
                // must set it via the flag after watch + gap-refresh are done.
                match session_fn(shutdown.clone()).await {
                    Ok(()) => {
                        session_active.store(false, Ordering::Release);
                        if shutdown.is_cancelled() {
                            return;
                        }
                        info!("[{}] Watch session ended, restarting...", name);
                        backoff = MIN_WATCH_RETRY_INTERVAL;
                    }
                    Err(e) => {
                        session_active.store(false, Ordering::Release);
                        if shutdown.is_cancelled() {
                            return;
                        }
                        warn!(
                            "[{}] Watch session failed: {}, retrying in {:?}",
                            name, e, backoff
                        );
                    }
                }

                tokio::select! {
                    _ = tokio::time::sleep(backoff) => {}
                    _ = shutdown.cancelled() => return,
                }
                backoff = std::cmp::min(backoff * 2, MAX_WATCH_RETRY_INTERVAL);
            }
        });

        *self.handle.write().await = Some(handle);
    }

    /// Wait until the first session becomes active, or return an error
    /// if the supervisor stops or the timeout expires.
    pub async fn wait_session_active(&self, timeout: Duration) -> Result<()> {
        let start = tokio::time::Instant::now();
        loop {
            if self.is_session_active() {
                return Ok(());
            }
            if self.is_stopped() {
                return Err(anyhow::anyhow!(
                    "Supervisor stopped before session became active"
                ));
            }
            if start.elapsed() >= timeout {
                return Err(anyhow::anyhow!(
                    "Timed out waiting for watch session to become active"
                ));
            }
            tokio::time::sleep(Duration::from_millis(50)).await;
        }
    }

    async fn stop_inner(&self) {
        if let Some(handle) = self.handle.write().await.take() {
            handle.abort();
            let _ = handle.await;
        }
    }

    pub async fn stop(&mut self) {
        self.shutdown.cancel();
        self.stop_inner().await;
        self.session_active.store(false, Ordering::Release);
    }
}