use std::sync::Mutex;
use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU64, Ordering};
use std::time::Instant;

use tokio::sync::watch;

use crate::types::now_millis;

/// Server lifecycle states.
#[repr(u8)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum ServerState {
    /// Normal operation. Accepting connections, processing all frames.
    Running = 0,
    /// Draining. No new connections accepted. Existing connections receive
    /// GoingAway (if capable) and are given time to finish. Publishes still
    /// accepted so in-flight work completes. New subscribes rejected.
    Draining = 1,
    /// Final shutdown. All connections force-closed. Engine and storage
    /// flushed and shut down.
    ShuttingDown = 2,
}

impl TryFrom<u8> for ServerState {
    type Error = ();

    fn try_from(value: u8) -> Result<Self, ()> {
        match value {
            0 => Ok(Self::Running),
            1 => Ok(Self::Draining),
            2 => Ok(Self::ShuttingDown),
            _ => Err(()),
        }
    }
}

/// Tracks server lifecycle state with lock-free reads on the hot path.
///
/// The `watch::Sender` broadcasts state transitions to all tasks that need
/// to react (TCP accept loop, HTTP server, connection handlers, background
/// tasks).
pub struct ServerLifecycle {
    state: AtomicU8,
    initialized: AtomicBool,
    active_connections: AtomicU64,
    notify: watch::Sender<ServerState>,
    drain_started_at: Mutex<Option<Instant>>,
    connections_at_drain_start: AtomicU64,
    /// Wall-clock timestamp (millis since epoch) when drain started.
    drain_started_wall: Mutex<Option<u64>>,
    /// Wall-clock timestamp (millis since epoch) when drain is expected to end.
    drain_deadline_wall: Mutex<Option<u64>>,
}

impl ServerLifecycle {
    pub fn new() -> (Self, watch::Receiver<ServerState>) {
        let (tx, rx) = watch::channel(ServerState::Running);
        let lifecycle = Self {
            state: AtomicU8::new(ServerState::Running as u8),
            initialized: AtomicBool::new(false),
            active_connections: AtomicU64::new(0),
            notify: tx,
            drain_started_at: Mutex::new(None),
            connections_at_drain_start: AtomicU64::new(0),
            drain_started_wall: Mutex::new(None),
            drain_deadline_wall: Mutex::new(None),
        };
        (lifecycle, rx)
    }

    /// Returns the current server state.
    pub fn state(&self) -> ServerState {
        ServerState::try_from(self.state.load(Ordering::Acquire)).unwrap_or(ServerState::Running)
    }

    /// Returns true when the server is fully initialized and in Running state.
    pub fn is_ready(&self) -> bool {
        self.initialized.load(Ordering::Acquire)
            && self.state.load(Ordering::Acquire) == ServerState::Running as u8
    }

    /// Returns true when the server is draining or shutting down.
    pub fn is_draining(&self) -> bool {
        self.state.load(Ordering::Acquire) >= ServerState::Draining as u8
    }

    /// Returns true when initialization is complete.
    pub fn is_initialized(&self) -> bool {
        self.initialized.load(Ordering::Acquire)
    }

    /// Mark initialization as complete.
    pub fn set_initialized(&self) {
        self.initialized.store(true, Ordering::Release);
    }

    /// Transition to a new state. Only forward transitions are allowed
    /// (Running → Draining → ShuttingDown).
    ///
    /// For transitioning to Draining, prefer [`begin_drain`] which also records
    /// wall-clock timestamps and drain deadline.
    pub fn transition_to(&self, new_state: ServerState) {
        let new_val = new_state as u8;
        let prev = self.state.fetch_max(new_val, Ordering::AcqRel);

        if prev < new_val {
            if new_state == ServerState::Draining {
                let conns = self.active_connections.load(Ordering::Acquire);
                self.connections_at_drain_start
                    .store(conns, Ordering::Release);
                *self.drain_started_at.lock().unwrap() = Some(Instant::now());
                *self.drain_started_wall.lock().unwrap() = Some(now_millis());
            }
            metrics::gauge!("fafb_server_state").set(new_val as f64);
            let _ = self.notify.send(new_state);
        }
    }

    /// Transition to Draining state with wall-clock metadata.
    ///
    /// Records the drain start time (both monotonic and wall-clock) and the
    /// expected deadline. Returns `true` if the transition was performed
    /// (i.e., the server was in Running state).
    pub fn begin_drain(&self, timeout_secs: u64) -> bool {
        let new_val = ServerState::Draining as u8;
        let prev = self.state.fetch_max(new_val, Ordering::AcqRel);

        if prev < new_val {
            let conns = self.active_connections.load(Ordering::Acquire);
            self.connections_at_drain_start
                .store(conns, Ordering::Release);
            *self.drain_started_at.lock().unwrap() = Some(Instant::now());
            let wall_now = now_millis();
            *self.drain_started_wall.lock().unwrap() = Some(wall_now);
            *self.drain_deadline_wall.lock().unwrap() = Some(wall_now + timeout_secs * 1000);
            metrics::gauge!("fafb_server_state").set(new_val as f64);
            let _ = self.notify.send(ServerState::Draining);
            true
        } else {
            false
        }
    }

    /// Subscribe to state change notifications.
    pub fn subscribe(&self) -> watch::Receiver<ServerState> {
        self.notify.subscribe()
    }

    /// Increment the active connection counter. Returns the new count.
    pub fn inc_connections(&self) -> u64 {
        self.active_connections.fetch_add(1, Ordering::AcqRel) + 1
    }

    /// Decrement the active connection counter. Returns the new count.
    pub fn dec_connections(&self) -> u64 {
        self.active_connections.fetch_sub(1, Ordering::AcqRel) - 1
    }

    /// Current active connection count.
    pub fn active_connections(&self) -> u64 {
        self.active_connections.load(Ordering::Acquire)
    }

    /// How many connections were active when drain started.
    pub fn connections_at_drain_start(&self) -> u64 {
        self.connections_at_drain_start.load(Ordering::Acquire)
    }

    /// When drain started, if it has (monotonic clock).
    pub fn drain_started_at(&self) -> Option<Instant> {
        *self.drain_started_at.lock().unwrap()
    }

    /// Wall-clock timestamp (millis since epoch) when drain started.
    pub fn drain_started_wall(&self) -> Option<u64> {
        *self.drain_started_wall.lock().unwrap()
    }

    /// Wall-clock timestamp (millis since epoch) of expected drain deadline.
    pub fn drain_deadline_wall(&self) -> Option<u64> {
        *self.drain_deadline_wall.lock().unwrap()
    }

    /// How many connections have been drained since drain started.
    pub fn connections_drained(&self) -> u64 {
        self.connections_at_drain_start
            .load(Ordering::Acquire)
            .saturating_sub(self.active_connections.load(Ordering::Acquire))
    }

    /// Run the drain-wait loop: wait for all connections to close or timeout,
    /// log progress, record drain duration metric, then transition to ShuttingDown.
    ///
    /// Call `begin_drain()` and any pre-drain work (force_flush, set_inflight_timeout)
    /// before calling this method.
    pub async fn run_drain_wait(&self, timeout: std::time::Duration) {
        if timeout.as_secs() > 0 {
            let drain_start = tokio::time::Instant::now();
            let drain_deadline = drain_start + timeout;
            let mut progress_interval = tokio::time::interval(std::time::Duration::from_secs(5));
            progress_interval.tick().await; // skip first immediate tick

            loop {
                if self.active_connections() == 0 {
                    tracing::info!("all connections closed, skipping remaining drain wait");
                    break;
                }
                if tokio::time::Instant::now() >= drain_deadline {
                    tracing::info!(
                        remaining_connections = self.active_connections(),
                        "drain timeout expired"
                    );
                    break;
                }
                progress_interval.tick().await;
                tracing::info!(
                    remaining_connections = self.active_connections(),
                    elapsed_secs = drain_start.elapsed().as_secs(),
                    "draining"
                );
            }
        }

        // Record drain duration metric.
        if let Some(started) = self.drain_started_at() {
            metrics::histogram!("fafb_drain_duration_seconds")
                .record(started.elapsed().as_secs_f64());
        }

        // Transition to ShuttingDown — force-close remaining connections.
        self.transition_to(ServerState::ShuttingDown);
        tracing::info!("entering shutdown phase");
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn initial_state_is_running() {
        let (lifecycle, _rx) = ServerLifecycle::new();
        assert_eq!(lifecycle.state(), ServerState::Running);
        assert!(!lifecycle.is_draining());
        assert!(!lifecycle.is_ready()); // not initialized yet
    }

    #[test]
    fn is_ready_requires_initialized_and_running() {
        let (lifecycle, _rx) = ServerLifecycle::new();
        assert!(!lifecycle.is_ready());

        lifecycle.set_initialized();
        assert!(lifecycle.is_ready());

        lifecycle.transition_to(ServerState::Draining);
        assert!(!lifecycle.is_ready());
    }

    #[test]
    fn forward_transitions_only() {
        let (lifecycle, _rx) = ServerLifecycle::new();

        lifecycle.transition_to(ServerState::Draining);
        assert_eq!(lifecycle.state(), ServerState::Draining);
        assert!(lifecycle.is_draining());

        // Backward transition should be ignored
        lifecycle.transition_to(ServerState::Running);
        assert_eq!(lifecycle.state(), ServerState::Draining);

        lifecycle.transition_to(ServerState::ShuttingDown);
        assert_eq!(lifecycle.state(), ServerState::ShuttingDown);
    }

    #[test]
    fn connection_counting() {
        let (lifecycle, _rx) = ServerLifecycle::new();
        assert_eq!(lifecycle.active_connections(), 0);

        assert_eq!(lifecycle.inc_connections(), 1);
        assert_eq!(lifecycle.inc_connections(), 2);
        assert_eq!(lifecycle.active_connections(), 2);

        assert_eq!(lifecycle.dec_connections(), 1);
        assert_eq!(lifecycle.active_connections(), 1);
    }

    #[test]
    fn drain_metadata_recorded() {
        let (lifecycle, _rx) = ServerLifecycle::new();
        lifecycle.inc_connections();
        lifecycle.inc_connections();
        lifecycle.inc_connections();

        assert!(lifecycle.drain_started_at().is_none());

        lifecycle.transition_to(ServerState::Draining);
        assert!(lifecycle.drain_started_at().is_some());
        assert_eq!(lifecycle.connections_at_drain_start(), 3);
        assert!(lifecycle.drain_started_wall().is_some());
    }

    #[test]
    fn begin_drain_records_wall_clock_and_deadline() {
        let (lifecycle, _rx) = ServerLifecycle::new();
        lifecycle.inc_connections();
        lifecycle.inc_connections();

        assert!(lifecycle.drain_started_wall().is_none());
        assert!(lifecycle.drain_deadline_wall().is_none());
        assert_eq!(lifecycle.connections_drained(), 0);

        let transitioned = lifecycle.begin_drain(30);
        assert!(transitioned);
        assert_eq!(lifecycle.state(), ServerState::Draining);

        let started = lifecycle.drain_started_wall().unwrap();
        let deadline = lifecycle.drain_deadline_wall().unwrap();
        assert!(started > 0);
        assert_eq!(deadline, started + 30_000);
        assert_eq!(lifecycle.connections_at_drain_start(), 2);
        assert_eq!(lifecycle.connections_drained(), 0);

        // Disconnect one connection
        lifecycle.dec_connections();
        assert_eq!(lifecycle.connections_drained(), 1);

        // Second begin_drain call should be a no-op
        let transitioned2 = lifecycle.begin_drain(60);
        assert!(!transitioned2);
    }

    #[tokio::test]
    async fn watch_receiver_gets_updates() {
        let (lifecycle, mut rx) = ServerLifecycle::new();

        lifecycle.transition_to(ServerState::Draining);
        rx.changed().await.unwrap();
        assert_eq!(*rx.borrow(), ServerState::Draining);

        lifecycle.transition_to(ServerState::ShuttingDown);
        rx.changed().await.unwrap();
        assert_eq!(*rx.borrow(), ServerState::ShuttingDown);
    }
}
