//! Replication lag gauge: computes per-shard lag in WAL segments.
//!
//! The lag is defined as `leader_seqno - applied_seqno` for each shard.
//! A lag of 0 means the follower is fully caught up. The gauge is designed
//! for integration into the Prometheus metrics endpoint.

use std::sync::atomic::{AtomicU64, Ordering};

use super::shard::ShardId;
use super::state::ReplicationState;

/// Tracks the latest known leader seqno per shard and computes the lag
/// relative to the local `ReplicationState`.
///
/// # Thread Safety
///
/// Uses `AtomicU64` for the leader seqno. The `ReplicationState` it
/// references is also lock-free. Reading lag is wait-free.
#[derive(Debug)]
pub struct ReplicationLagGauge {
    /// The latest seqno observed from the leader for the primary shard.
    /// In single-shard deployments, this is `ShardId::SINGLE`.
    leader_seqno: AtomicU64,
    /// The local replication state (shared with the receiver).
    state: std::sync::Arc<ReplicationState>,
    /// The shard we are tracking.
    shard_id: ShardId,
}

impl ReplicationLagGauge {
    /// Create a new lag gauge for a single shard.
    #[must_use]
    pub const fn new(shard_id: ShardId, state: std::sync::Arc<ReplicationState>) -> Self {
        Self {
            leader_seqno: AtomicU64::new(0),
            state,
            shard_id,
        }
    }

    /// Update the leader's latest known seqno.
    ///
    /// Monotonic: smaller values are silently ignored.
    pub fn update_leader_seqno(&self, seqno: u64) {
        // CAS loop to enforce monotonicity (same pattern as ReplicationState::advance).
        let mut current = self.leader_seqno.load(Ordering::Acquire);
        loop {
            if seqno <= current {
                break;
            }
            match self.leader_seqno.compare_exchange_weak(
                current,
                seqno,
                Ordering::AcqRel,
                Ordering::Acquire,
            ) {
                Ok(_) => break,
                Err(actual) => current = actual,
            }
        }
    }

    /// Compute the current lag in WAL segments.
    ///
    /// Returns `leader_seqno - applied_seqno`. If the leader seqno is
    /// unknown (still 0), returns 0.
    #[must_use]
    pub fn lag_segments(&self) -> u64 {
        let leader = self.leader_seqno.load(Ordering::Acquire);
        let applied = self.state.applied_seqno(self.shard_id).unwrap_or(0);
        leader.saturating_sub(applied)
    }

    /// The shard this gauge tracks.
    #[must_use]
    pub const fn shard_id(&self) -> ShardId {
        self.shard_id
    }

    /// The latest leader seqno observed.
    #[must_use]
    pub fn leader_seqno(&self) -> u64 {
        self.leader_seqno.load(Ordering::Acquire)
    }

    /// The latest applied seqno on this node.
    #[must_use]
    pub fn applied_seqno(&self) -> u64 {
        self.state.applied_seqno(self.shard_id).unwrap_or(0)
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod tests {
    use std::sync::Arc;

    use super::*;

    #[test]
    fn lag_starts_at_zero() {
        let state = Arc::new(ReplicationState::single());
        let gauge = ReplicationLagGauge::new(ShardId::SINGLE, state);
        assert_eq!(gauge.lag_segments(), 0);
    }

    #[test]
    fn lag_reflects_leader_ahead() {
        let state = Arc::new(ReplicationState::single());
        let gauge = ReplicationLagGauge::new(ShardId::SINGLE, Arc::clone(&state));
        gauge.update_leader_seqno(10);
        assert_eq!(gauge.lag_segments(), 10);
        // Apply some segments
        state.advance(ShardId::SINGLE, 7);
        assert_eq!(gauge.lag_segments(), 3);
    }

    #[test]
    fn lag_zero_when_caught_up() {
        let state = Arc::new(ReplicationState::single());
        let gauge = ReplicationLagGauge::new(ShardId::SINGLE, Arc::clone(&state));
        gauge.update_leader_seqno(5);
        state.advance(ShardId::SINGLE, 5);
        assert_eq!(gauge.lag_segments(), 0);
    }

    #[test]
    fn leader_seqno_is_monotone() {
        let state = Arc::new(ReplicationState::single());
        let gauge = ReplicationLagGauge::new(ShardId::SINGLE, state);
        gauge.update_leader_seqno(10);
        gauge.update_leader_seqno(5); // should be ignored
        assert_eq!(gauge.leader_seqno(), 10);
        gauge.update_leader_seqno(15);
        assert_eq!(gauge.leader_seqno(), 15);
    }

    #[test]
    fn accessors_work() {
        let state = Arc::new(ReplicationState::single());
        state.advance(ShardId::SINGLE, 3);
        let gauge = ReplicationLagGauge::new(ShardId::SINGLE, Arc::clone(&state));
        gauge.update_leader_seqno(7);
        assert_eq!(gauge.shard_id(), ShardId::SINGLE);
        assert_eq!(gauge.leader_seqno(), 7);
        assert_eq!(gauge.applied_seqno(), 3);
        assert_eq!(gauge.lag_segments(), 4);
    }
}
