use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};

use crate::replication::shard::ShardId;

/// Tracks the per-shard replication high-water-mark.
///
/// Each shard tracks the last WAL segment seqno that has been fully
/// applied to the local state machine. Segments with seqno <=
/// `high_water_mark` are idempotent no-ops on replay.
///
/// Thread-safe: atomics for lock-free reads/writes. Clone takes a
/// snapshot (clones the values, not the Arcs).
#[derive(Debug)]
pub struct ReplicationState {
    /// Per-shard high-water-mark seqno. 0 means "no segments applied yet".
    applied: HashMap<ShardId, Arc<AtomicU64>>,
}

impl ReplicationState {
    /// Create a new `ReplicationState` tracking the given shards.
    /// All high-water-marks start at 0.
    #[must_use]
    pub fn new(shards: &[ShardId]) -> Self {
        let applied = shards
            .iter()
            .map(|&s| (s, Arc::new(AtomicU64::new(0))))
            .collect();
        Self { applied }
    }

    /// Create a single-node `ReplicationState` (tracks only `ShardId::SINGLE`).
    #[must_use]
    pub fn single() -> Self {
        Self::new(&[ShardId::SINGLE])
    }

    /// Get the high-water-mark seqno for a shard.
    /// Returns `None` if the shard is unknown.
    #[must_use]
    pub fn applied_seqno(&self, shard_id: ShardId) -> Option<u64> {
        self.applied
            .get(&shard_id)
            .map(|a| a.load(Ordering::Acquire))
    }

    /// Advance the high-water-mark for a shard.
    ///
    /// Only moves forward -- smaller seqnos are silently ignored.
    /// Safe to call from multiple threads concurrently via a CAS loop
    /// that enforces monotonicity.
    pub fn advance(&self, shard_id: ShardId, seqno: u64) {
        let Some(atomic) = self.applied.get(&shard_id) else {
            tracing::warn!(
                shard = %shard_id,
                seqno,
                "advance called for unknown shard; ignored"
            );
            return;
        };
        // Acquire: see the latest value before deciding whether to advance.
        let mut current = atomic.load(Ordering::Acquire);
        loop {
            if seqno <= current {
                break;
            }
            // AcqRel on success: the new value is visible to subsequent Acquire loads.
            // Acquire on failure: reload the latest value for the next iteration.
            match atomic.compare_exchange_weak(current, seqno, Ordering::AcqRel, Ordering::Acquire)
            {
                Ok(_) => break,
                Err(actual) => current = actual,
            }
        }
    }

    /// Returns a snapshot of all tracked shards and their current seqnos.
    #[must_use]
    pub fn snapshot(&self) -> HashMap<ShardId, u64> {
        self.applied
            .iter()
            .map(|(&s, a)| (s, a.load(Ordering::Acquire)))
            .collect()
    }

    /// Serialize for checkpoint persistence.
    ///
    /// Uses `HashMap<u16, u64>` as the wire format because JSON object keys
    /// must be strings, and `serde_json` serializes integer keys correctly for
    /// `u16` but not for newtype wrappers like `ShardId`.
    ///
    /// # Panics
    ///
    /// Panics if `serde_json` fails to serialize `HashMap<u16, u64>`, which
    /// is unreachable because both key and value types are always representable.
    #[must_use]
    pub fn to_checkpoint_bytes(&self) -> Vec<u8> {
        let snap: HashMap<u16, u64> = self.snapshot().into_iter().map(|(s, v)| (s.0, v)).collect();
        // Serialization of HashMap<u16, u64> to JSON is infallible: both
        // key and value types are always representable.
        serde_json::to_vec(&snap).expect("ReplicationState serialization is infallible")
    }

    /// Restore from checkpoint bytes.
    ///
    /// Shards not present in the checkpoint start at 0. If the bytes are
    /// malformed, all shards start at 0 (safe: replaying from 0 is correct,
    /// just slower).
    #[must_use]
    pub fn from_checkpoint_bytes(bytes: &[u8], shards: &[ShardId]) -> Self {
        let snap: HashMap<u16, u64> = serde_json::from_slice(bytes).unwrap_or_default();
        let applied = shards
            .iter()
            .map(|&s| {
                let seqno = snap.get(&s.0).copied().unwrap_or(0);
                (s, Arc::new(AtomicU64::new(seqno)))
            })
            .collect();
        Self { applied }
    }
}

// Compile-time assertion: ReplicationState is Send + Sync.
// Arc<AtomicU64> is Send + Sync, and HashMap<K, V> is Send + Sync when both
// K and V are. These functions are never called -- they exist only to produce
// a compile error if the trait bounds are not satisfied.
#[allow(dead_code)]
const _: () = {
    const fn assert_send<T: Send>() {}
    const fn assert_sync<T: Sync>() {}
    const fn _check() {
        assert_send::<ReplicationState>();
        assert_sync::<ReplicationState>();
    }
};

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

    #[test]
    fn single_starts_at_zero() {
        let state = ReplicationState::single();
        assert_eq!(state.applied_seqno(ShardId::SINGLE), Some(0));
    }

    #[test]
    fn unknown_shard_returns_none() {
        let state = ReplicationState::single();
        assert_eq!(state.applied_seqno(ShardId(99)), None);
    }

    #[test]
    fn advance_monotone() {
        let state = ReplicationState::single();
        state.advance(ShardId::SINGLE, 10);
        assert_eq!(state.applied_seqno(ShardId::SINGLE), Some(10));
        // Cannot go backwards
        state.advance(ShardId::SINGLE, 5);
        assert_eq!(state.applied_seqno(ShardId::SINGLE), Some(10));
        // Can go further forward
        state.advance(ShardId::SINGLE, 20);
        assert_eq!(state.applied_seqno(ShardId::SINGLE), Some(20));
    }

    #[test]
    fn advance_unknown_shard_is_noop() {
        let state = ReplicationState::single();
        state.advance(ShardId(99), 100); // Should not panic or affect other shards
        assert_eq!(state.applied_seqno(ShardId::SINGLE), Some(0));
    }

    #[test]
    fn snapshot_returns_all_shards() {
        let shards = [ShardId(0), ShardId(1), ShardId(2)];
        let state = ReplicationState::new(&shards);
        state.advance(ShardId(1), 42);
        let snap = state.snapshot();
        assert_eq!(snap[&ShardId(0)], 0);
        assert_eq!(snap[&ShardId(1)], 42);
        assert_eq!(snap[&ShardId(2)], 0);
    }

    #[test]
    fn checkpoint_roundtrip() {
        let shards = [ShardId(0), ShardId(1)];
        let state = ReplicationState::new(&shards);
        state.advance(ShardId(0), 100);
        state.advance(ShardId(1), 200);

        let bytes = state.to_checkpoint_bytes();
        let restored = ReplicationState::from_checkpoint_bytes(&bytes, &shards);
        assert_eq!(restored.applied_seqno(ShardId(0)), Some(100));
        assert_eq!(restored.applied_seqno(ShardId(1)), Some(200));
    }

    #[test]
    fn checkpoint_missing_shard_starts_at_zero() {
        let snap: HashMap<u16, u64> = HashMap::new();
        let bytes = serde_json::to_vec(&snap).unwrap();
        let state = ReplicationState::from_checkpoint_bytes(&bytes, &[ShardId(5)]);
        assert_eq!(state.applied_seqno(ShardId(5)), Some(0));
    }

    #[test]
    fn checkpoint_malformed_bytes_starts_at_zero() {
        let state = ReplicationState::from_checkpoint_bytes(b"not json at all", &[ShardId(0)]);
        assert_eq!(state.applied_seqno(ShardId(0)), Some(0));
    }

    #[test]
    fn concurrent_advance_is_monotone() {
        let state = Arc::new(ReplicationState::new(&[ShardId::SINGLE]));
        let mut handles = vec![];
        for i in 0u64..4 {
            let s = Arc::clone(&state);
            handles.push(std::thread::spawn(move || {
                for j in 0u64..250 {
                    s.advance(ShardId::SINGLE, i * 250 + j);
                }
            }));
        }
        for h in handles {
            h.join().expect("thread panicked");
        }
        // After 4 threads each advancing, the monotonic CAS guarantees convergence
        // to the maximum written value. Thread 3 writes i*250+j for i=3, j=0..249,
        // so max = 3*250 + 249 = 999.
        let final_seqno = state.applied_seqno(ShardId::SINGLE).unwrap();
        assert_eq!(final_seqno, 999);
    }
}
