//! Rolling upgrade coordinator.

use std::collections::HashSet;
use std::sync::{Arc, RwLock};

use super::control::ControlPlane;
use super::shard::ShardId;
use crate::schema::TidalError;

/// Phases of a rolling upgrade.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum UpgradePhase {
    Ready,
    Draining { shard_id: ShardId },
}

/// Combined mutable state kept under one lock to prevent desync between
/// the current phase and the set of drained shards.
struct UpgradeState {
    phase: UpgradePhase,
    drained: HashSet<ShardId>,
}

impl UpgradeState {
    fn new() -> Self {
        Self {
            phase: UpgradePhase::Ready,
            drained: HashSet::new(),
        }
    }
}

/// Coordinates rolling upgrades by draining and rejoining shards one at a time.
///
/// Protocol:
/// 1. `drain(shard)` — stop routing new writes to the target; wait for lag = 0.
/// 2. Caller performs the upgrade (outside this coordinator's scope).
/// 3. `rejoin(shard)` — re-enable routing; verify the node can receive WAL.
///
/// Safety invariant: at least one shard must remain serving at all times.
/// `drain` enforces this by checking `topology_shard_count` before proceeding.
pub struct RollingUpgradeCoordinator {
    control_plane: Arc<ControlPlane>,
    state: RwLock<UpgradeState>,
}

impl RollingUpgradeCoordinator {
    /// Create a new coordinator in the `Ready` phase.
    #[must_use]
    pub fn new(control_plane: Arc<ControlPlane>) -> Self {
        Self {
            control_plane,
            state: RwLock::new(UpgradeState::new()),
        }
    }

    /// Begin draining a shard for upgrade.
    ///
    /// # Errors
    ///
    /// Returns `InvalidState` if another upgrade is already in progress, or if
    /// draining the requested shard would leave zero serving nodes.
    pub fn drain(&self, shard_id: ShardId) -> crate::Result<()> {
        let total_shards = self.control_plane.topology_shard_count();
        let mut state = self
            .state
            .write()
            .map_err(|_| TidalError::internal("upgrade", "state lock poisoned"))?;

        if state.phase != UpgradePhase::Ready {
            return Err(TidalError::InvalidState(
                "drain called while another upgrade is in progress".into(),
            ));
        }

        // Safety check: draining must leave at least 1 node serving.
        let would_leave = total_shards.saturating_sub(state.drained.len() + 1);
        if would_leave == 0 {
            return Err(TidalError::InvalidState(format!(
                "cannot drain shard {shard_id:?}: would leave 0 serving nodes \
                 (cluster has {total_shards} shard(s), {already} already drained)",
                already = state.drained.len(),
            )));
        }

        state.drained.insert(shard_id);
        state.phase = UpgradePhase::Draining { shard_id };
        drop(state);
        Ok(())
    }

    /// Rejoin a drained shard and reset to `Ready`.
    ///
    /// # Errors
    ///
    /// Returns `InvalidState` if the current phase is not `Draining` for the
    /// given shard, or if the state lock is poisoned.
    pub fn rejoin(&self, shard_id: ShardId) -> crate::Result<()> {
        let mut state = self
            .state
            .write()
            .map_err(|_| TidalError::internal("upgrade", "state lock poisoned"))?;
        match state.phase {
            UpgradePhase::Draining { shard_id: draining } if draining == shard_id => {}
            _ => {
                return Err(TidalError::InvalidState(format!(
                    "rejoin({shard_id:?}) called outside Draining phase for that shard"
                )));
            }
        }
        state.drained.remove(&shard_id);
        state.phase = UpgradePhase::Ready;
        drop(state);
        Ok(())
    }

    /// Check if a shard has been drained.
    #[must_use]
    pub fn is_drained(&self, shard_id: ShardId) -> bool {
        self.state
            .read()
            .map(|s| s.drained.contains(&shard_id))
            .unwrap_or(false)
    }

    /// Return the current upgrade phase.
    ///
    /// # Panics
    ///
    /// Panics if the internal state lock is poisoned.
    #[must_use]
    pub fn current_phase(&self) -> UpgradePhase {
        self.state
            .read()
            .expect("state lock poisoned")
            .phase
            .clone()
    }

    /// Return a reference to the control plane.
    #[must_use]
    #[allow(clippy::missing_const_for_fn)] // Arc field prevents const in practice
    pub fn control_plane(&self) -> &Arc<ControlPlane> {
        &self.control_plane
    }
}

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

    use super::*;
    use crate::replication::control::ControlPlane;
    use crate::replication::lag::ReplicationLagGauge;
    use crate::replication::shard::RegionId;
    use crate::replication::state::ReplicationState;
    use crate::replication::tenant::{ClusterTopology, ShardAssignment, TenantRouter};

    fn make_coordinator_with_shards(n: usize) -> RollingUpgradeCoordinator {
        let shards = (0..n as u16)
            .map(|i| ShardAssignment {
                shard_id: ShardId(i),
                region_id: RegionId(i),
            })
            .collect();
        let topo = Arc::new(std::sync::RwLock::new(ClusterTopology { shards }));
        let router = Arc::new(TenantRouter::new(Arc::clone(&topo)));
        let state = Arc::new(ReplicationState::single());
        let lag = Arc::new(ReplicationLagGauge::new(ShardId::SINGLE, state));
        let cp = Arc::new(ControlPlane::new(topo, router, lag));
        RollingUpgradeCoordinator::new(cp)
    }

    #[test]
    fn initial_phase_is_ready() {
        let c = make_coordinator_with_shards(2);
        assert_eq!(c.current_phase(), UpgradePhase::Ready);
    }

    #[test]
    fn drain_marks_shard_as_drained() {
        let c = make_coordinator_with_shards(2);
        c.drain(ShardId(0)).unwrap();
        assert!(c.is_drained(ShardId(0)));
        assert!(matches!(
            c.current_phase(),
            UpgradePhase::Draining {
                shard_id: ShardId(0)
            }
        ));
    }

    #[test]
    fn drain_fails_when_not_ready() {
        let c = make_coordinator_with_shards(3);
        c.drain(ShardId(0)).unwrap();
        let err = c.drain(ShardId(1)).unwrap_err();
        assert!(matches!(err, TidalError::InvalidState(_)));
    }

    #[test]
    fn drain_fails_when_would_leave_zero_serving_nodes() {
        // Single-shard cluster: draining the only shard must fail.
        let c = make_coordinator_with_shards(1);
        let err = c.drain(ShardId::SINGLE).unwrap_err();
        assert!(matches!(err, TidalError::InvalidState(_)));
    }

    #[test]
    fn rejoin_clears_drain_and_resets_to_ready() {
        let c = make_coordinator_with_shards(2);
        c.drain(ShardId(0)).unwrap();
        c.rejoin(ShardId(0)).unwrap();
        assert!(!c.is_drained(ShardId(0)));
        assert_eq!(c.current_phase(), UpgradePhase::Ready);
    }

    #[test]
    fn can_drain_different_shard_after_rejoin() {
        let c = make_coordinator_with_shards(3);
        c.drain(ShardId(0)).unwrap();
        c.rejoin(ShardId(0)).unwrap();
        c.drain(ShardId(1)).unwrap();
        assert!(c.is_drained(ShardId(1)));
    }

    #[test]
    fn is_drained_reflects_consistent_state() {
        // Verify is_drained() and current_phase() never show contradictory state.
        let c = make_coordinator_with_shards(2);
        // Before drain: both agree on not-drained.
        assert!(!c.is_drained(ShardId(0)));
        assert_eq!(c.current_phase(), UpgradePhase::Ready);
        // After drain: both agree on drained.
        c.drain(ShardId(0)).unwrap();
        assert!(c.is_drained(ShardId(0)));
        assert!(matches!(c.current_phase(), UpgradePhase::Draining { .. }));
        // After rejoin: both agree on not-drained.
        c.rejoin(ShardId(0)).unwrap();
        assert!(!c.is_drained(ShardId(0)));
        assert_eq!(c.current_phase(), UpgradePhase::Ready);
    }
}
