//! Control plane: cluster topology management and shard health tracking.

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

use dashmap::DashMap;

use super::lag::ReplicationLagGauge;
use super::shard::{RegionId, ShardId};
use super::tenant::{ClusterTopology, ShardAssignment, TenantRouter};

/// Per-shard operational metrics reported via heartbeat.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct ShardStats {
    pub shard_id: ShardId,
    pub region_id: RegionId,
    pub entity_count: u64,
    pub signal_throughput_eps: f64,
    pub replication_lag: HashMap<RegionId, u64>,
    pub disk_bytes: u64,
    pub last_heartbeat_ns: u64,
}

/// Health classification for a region.
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub enum RegionHealth {
    Healthy,
    Degraded,
    Offline,
}

/// Snapshot of cluster-wide health.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct ClusterHealth {
    pub snapshot_ns: u64,
    pub shards: Vec<ShardStats>,
    pub regions: HashMap<RegionId, RegionHealth>,
    pub tenant_count: usize,
    pub total_entities: u64,
    pub total_signals_eps: f64,
}

/// Central coordinator for cluster topology and shard health.
pub struct ControlPlane {
    topology: Arc<RwLock<ClusterTopology>>,
    tenant_router: Arc<TenantRouter>,
    lag_gauge: Arc<ReplicationLagGauge>,
    shard_stats: DashMap<ShardId, ShardStats>,
    region_health: DashMap<RegionId, RegionHealth>,
}

impl ControlPlane {
    /// Replication lag (ns) above which a region is classified as `Degraded`.
    /// Matches spec task-03: "`replication_lag` > 5s → Degraded".
    const DEGRADED_LAG_NS: u64 = 5_000_000_000;
    /// Heartbeat age (ns) above which a region is classified as `Offline`.
    /// Matches spec task-03: "no heartbeat for > 30s → Offline".
    const OFFLINE_THRESHOLD_NS: u64 = 30_000_000_000;

    /// Create a new control plane.
    #[must_use]
    pub fn new(
        topology: Arc<RwLock<ClusterTopology>>,
        tenant_router: Arc<TenantRouter>,
        lag_gauge: Arc<ReplicationLagGauge>,
    ) -> Self {
        Self {
            topology,
            tenant_router,
            lag_gauge,
            shard_stats: DashMap::new(),
            region_health: DashMap::new(),
        }
    }

    /// Record a heartbeat from a shard and recompute its region's health.
    pub fn record_shard_heartbeat(&self, stats: ShardStats) {
        let region = stats.region_id;
        self.shard_stats.insert(stats.shard_id, stats);
        self.recompute_region_health(region);
    }

    /// Build a point-in-time cluster health snapshot.
    #[must_use]
    pub fn health(&self) -> ClusterHealth {
        let now = crate::replication::now_ns();
        let shards: Vec<ShardStats> = self.shard_stats.iter().map(|e| e.value().clone()).collect();
        let total_entities: u64 = shards.iter().map(|s| s.entity_count).sum();
        let total_signals_eps: f64 = shards.iter().map(|s| s.signal_throughput_eps).sum();
        let regions: HashMap<RegionId, RegionHealth> = self
            .region_health
            .iter()
            .map(|e| (*e.key(), *e.value()))
            .collect();
        ClusterHealth {
            snapshot_ns: now,
            shards,
            regions,
            tenant_count: self.tenant_router.tenant_count(),
            total_entities,
            total_signals_eps,
        }
    }

    /// Serialize cluster health as a pretty-printed JSON string.
    ///
    /// # Errors
    ///
    /// Returns an error if JSON serialization fails (should never happen for
    /// well-formed `ClusterHealth`, but propagated rather than swallowed).
    pub fn health_json(&self) -> crate::Result<String> {
        serde_json::to_string_pretty(&self.health()).map_err(|e| {
            crate::TidalError::internal("health_json", format!("serialization failed: {e}"))
        })
    }

    /// Add or update a shard assignment in the cluster topology.
    pub fn update_topology(&self, assignment: ShardAssignment) {
        let mut topo = self
            .topology
            .write()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        if let Some(existing) = topo
            .shards
            .iter_mut()
            .find(|s| s.shard_id == assignment.shard_id)
        {
            *existing = assignment;
        } else {
            topo.shards.push(assignment);
        }
    }

    /// Return the replication lag in WAL segments for the given shard.
    ///
    /// **Single-shard limitation:** this implementation delegates to a single
    /// `ReplicationLagGauge`. Multi-shard deployments need a
    /// `HashMap<ShardId, ReplicationLagGauge>` (planned for M8p6+). In debug
    /// builds, passing a shard other than `ShardId::SINGLE` triggers an
    /// assertion to catch accidental misuse before multi-shard support lands.
    #[must_use]
    pub fn lag_for(&self, shard_id: ShardId) -> u64 {
        debug_assert_eq!(
            shard_id,
            ShardId::SINGLE,
            "lag_for: multi-shard lag tracking not yet implemented; got {shard_id:?}"
        );
        self.lag_gauge.lag_segments()
    }

    /// Return the number of shards currently registered in the cluster topology.
    ///
    /// Used by `RollingUpgradeCoordinator` to enforce the minimum-node safety
    /// constraint: at least one shard must remain serving during an upgrade.
    #[must_use]
    pub fn topology_shard_count(&self) -> usize {
        self.topology
            .read()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .shards
            .len()
    }

    fn recompute_region_health(&self, region: RegionId) {
        let now = crate::replication::now_ns();
        let health = self
            .shard_stats
            .iter()
            .filter(|e| e.value().region_id == region)
            .map(|e| {
                let s = e.value();
                let age = now.saturating_sub(s.last_heartbeat_ns);
                if age >= Self::OFFLINE_THRESHOLD_NS {
                    // No heartbeat for > 30s → assume node is unreachable.
                    RegionHealth::Offline
                } else if s
                    .replication_lag
                    .values()
                    .any(|&lag| lag > Self::DEGRADED_LAG_NS)
                {
                    // Follower > 5s behind leader → replication is struggling.
                    RegionHealth::Degraded
                } else {
                    RegionHealth::Healthy
                }
            })
            .fold(RegionHealth::Healthy, |worst, h| {
                if h == RegionHealth::Offline || worst == RegionHealth::Offline {
                    RegionHealth::Offline
                } else if h == RegionHealth::Degraded {
                    RegionHealth::Degraded
                } else {
                    worst
                }
            });
        self.region_health.insert(region, health);
    }
}

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

    use super::*;
    use crate::replication::state::ReplicationState;
    use crate::replication::tenant::ClusterTopology;

    fn make_control_plane() -> ControlPlane {
        let topo = Arc::new(RwLock::new(ClusterTopology::single()));
        let router = Arc::new(crate::replication::tenant::TenantRouter::new(Arc::clone(
            &topo,
        )));
        let state = Arc::new(ReplicationState::single());
        let lag = Arc::new(ReplicationLagGauge::new(ShardId::SINGLE, state));
        ControlPlane::new(topo, router, lag)
    }

    fn make_stats(shard_id: ShardId, region_id: RegionId, age_ns: u64) -> ShardStats {
        let now = crate::replication::now_ns();
        ShardStats {
            shard_id,
            region_id,
            entity_count: 100,
            signal_throughput_eps: 50.0,
            replication_lag: HashMap::new(),
            disk_bytes: 1024,
            last_heartbeat_ns: now.saturating_sub(age_ns),
        }
    }

    #[test]
    fn health_starts_empty() {
        let cp = make_control_plane();
        let h = cp.health();
        assert!(h.shards.is_empty());
        assert_eq!(h.total_entities, 0);
    }

    #[test]
    fn record_heartbeat_updates_health() {
        let cp = make_control_plane();
        cp.record_shard_heartbeat(make_stats(ShardId::SINGLE, RegionId::SINGLE, 0));
        let h = cp.health();
        assert_eq!(h.shards.len(), 1);
        assert_eq!(h.total_entities, 100);
    }

    #[test]
    fn region_health_healthy_for_fresh_heartbeat() {
        let cp = make_control_plane();
        cp.record_shard_heartbeat(make_stats(ShardId::SINGLE, RegionId::SINGLE, 0));
        let h = cp.health();
        assert_eq!(h.regions[&RegionId::SINGLE], RegionHealth::Healthy);
    }

    fn make_stats_with_lag(shard_id: ShardId, region_id: RegionId, lag_ns: u64) -> ShardStats {
        let now = crate::replication::now_ns();
        let mut replication_lag = HashMap::new();
        // Simulate lag reported to a peer region.
        replication_lag.insert(RegionId(99), lag_ns);
        ShardStats {
            shard_id,
            region_id,
            entity_count: 100,
            signal_throughput_eps: 50.0,
            replication_lag,
            disk_bytes: 1024,
            last_heartbeat_ns: now, // fresh heartbeat; degraded is driven by lag, not age
        }
    }

    #[test]
    fn region_health_degraded_when_replication_lag_exceeds_5s() {
        let cp = make_control_plane();
        // Lag of 6s triggers Degraded; heartbeat is fresh so it is not Offline.
        cp.record_shard_heartbeat(make_stats_with_lag(
            ShardId::SINGLE,
            RegionId::SINGLE,
            6_000_000_000,
        ));
        let h = cp.health();
        assert_eq!(h.regions[&RegionId::SINGLE], RegionHealth::Degraded);
    }

    #[test]
    fn region_health_healthy_when_replication_lag_under_5s() {
        let cp = make_control_plane();
        // Lag of 2s is below the 5s threshold → Healthy.
        cp.record_shard_heartbeat(make_stats_with_lag(
            ShardId::SINGLE,
            RegionId::SINGLE,
            2_000_000_000,
        ));
        let h = cp.health();
        assert_eq!(h.regions[&RegionId::SINGLE], RegionHealth::Healthy);
    }

    #[test]
    fn region_health_offline_for_very_stale_heartbeat() {
        let cp = make_control_plane();
        // Heartbeat age of 35s exceeds the 30s offline threshold.
        cp.record_shard_heartbeat(make_stats(
            ShardId::SINGLE,
            RegionId::SINGLE,
            35_000_000_000,
        ));
        let h = cp.health();
        assert_eq!(h.regions[&RegionId::SINGLE], RegionHealth::Offline);
    }

    #[test]
    fn health_json_is_valid_json() {
        let cp = make_control_plane();
        cp.record_shard_heartbeat(make_stats(ShardId::SINGLE, RegionId::SINGLE, 0));
        let json = cp.health_json().unwrap();
        let v: serde_json::Value = serde_json::from_str(&json).unwrap();
        assert!(v.get("total_entities").is_some());
    }

    #[test]
    fn update_topology_adds_shard() {
        let cp = make_control_plane();
        cp.update_topology(ShardAssignment {
            shard_id: ShardId(2),
            region_id: RegionId(1),
        });
        let topo = cp.topology.read().unwrap();
        assert_eq!(topo.shards.len(), 2);
    }

    #[test]
    fn lag_for_returns_gauge_value() {
        let topo = Arc::new(RwLock::new(ClusterTopology::single()));
        let router = Arc::new(crate::replication::tenant::TenantRouter::new(Arc::clone(
            &topo,
        )));
        let state = Arc::new(ReplicationState::single());
        let lag = Arc::new(ReplicationLagGauge::new(
            ShardId::SINGLE,
            Arc::clone(&state),
        ));
        lag.update_leader_seqno(5);
        let cp = ControlPlane::new(topo, router, Arc::clone(&lag));
        assert_eq!(cp.lag_for(ShardId::SINGLE), 5);
    }
}
