//! Fencing tokens for split-brain prevention.
//!
//! Each partition ownership claim gets a monotonically increasing epoch number
//! stored in etcd. Every write operation on a partition must include the epoch;
//! if the epoch is stale (lower than the current epoch in etcd), the write is
//! rejected.
//!
//! This prevents a pod that was partitioned from the network (and doesn't know
//! it lost ownership) from writing stale data after a new owner has taken over.
//!
//! ## etcd key schema
//!
//! ```text
//! /ocg/partitions/{partition_id}/epoch → monotonically increasing u64
//! ```

#[cfg(feature = "distributed")]
use std::sync::atomic::{AtomicU64, Ordering};

#[cfg(feature = "distributed")]
use etcd_client::{Client, Compare, CompareOp, GetOptions, Txn, TxnOp};

#[cfg(feature = "distributed")]
use crate::error::{CypherError, CypherResult, ExecutionError};

// ── Fencing token ────────────────────────────────────────────────────────────

/// A fencing token that tracks the current epoch for a partition.
///
/// The epoch is a monotonically increasing counter stored in etcd. When a pod
/// claims ownership of a partition, it increments the epoch and records it
/// locally. All mutations check that the local epoch matches the token's epoch
/// before proceeding.
#[cfg(feature = "distributed")]
#[derive(Debug)]
pub struct FencingToken {
    partition_id: u32,
    epoch: AtomicU64,
}

#[cfg(feature = "distributed")]
impl FencingToken {
    /// Create a fencing token with the given initial epoch.
    pub fn new(partition_id: u32, epoch: u64) -> Self {
        Self {
            partition_id,
            epoch: AtomicU64::new(epoch),
        }
    }

    /// Get the partition ID this token belongs to.
    pub fn partition_id(&self) -> u32 {
        self.partition_id
    }

    /// Get the current epoch value.
    pub fn epoch(&self) -> u64 {
        self.epoch.load(Ordering::SeqCst)
    }

    /// Update the local epoch (after a successful CAS claim in etcd).
    pub fn set_epoch(&self, new_epoch: u64) {
        self.epoch.store(new_epoch, Ordering::SeqCst);
    }

    /// Validate that a given epoch matches the current one.
    ///
    /// Returns `Ok(())` if the epoch matches, or `Err(StaleEpoch)` if
    /// the caller's epoch is lower than the current one.
    pub fn validate(&self, caller_epoch: u64) -> CypherResult<()> {
        let current = self.epoch();
        if caller_epoch < current {
            Err(CypherError::Execution(ExecutionError::Internal(format!(
                "stale epoch for partition {}: caller has {}, current is {}",
                self.partition_id, caller_epoch, current
            ))))
        } else {
            Ok(())
        }
    }
}

#[cfg(feature = "distributed")]
impl Clone for FencingToken {
    fn clone(&self) -> Self {
        Self {
            partition_id: self.partition_id,
            epoch: AtomicU64::new(self.epoch.load(Ordering::SeqCst)),
        }
    }
}

// ── etcd epoch operations ────────────────────────────────────────────────────

/// Read the current epoch for a partition from etcd.
///
/// Returns `0` if no epoch key exists yet (first claim).
#[cfg(feature = "distributed")]
pub async fn read_epoch(client: &mut Client, partition_id: u32) -> CypherResult<u64> {
    let key = epoch_key(partition_id);
    let resp = client.get(key, None::<GetOptions>).await
        .map_err(|e| fence_err(format!("read epoch: {e}")))?;

    match resp.kvs().first() {
        Some(kv) => {
            let val = kv.value_str()
                .map_err(|e| fence_err(format!("epoch value utf-8: {e}")))?;
            val.parse::<u64>()
                .map_err(|e| fence_err(format!("epoch parse: {e}")))
        }
        None => Ok(0),
    }
}

/// Atomically increment the epoch for a partition using CAS.
///
/// Reads the current epoch, then uses a transaction to set `epoch + 1`
/// only if no other pod has changed it in the meantime.
///
/// Returns the new epoch on success, or an error if CAS failed.
#[cfg(feature = "distributed")]
pub async fn increment_epoch(
    client: &mut Client,
    partition_id: u32,
    expected_epoch: u64,
) -> CypherResult<u64> {
    let key = epoch_key(partition_id);
    let new_epoch = expected_epoch + 1;
    let new_value = new_epoch.to_string();

    let txn = if expected_epoch == 0 {
        // First claim: key should not exist (version == 0 means "key doesn't exist")
        let cmp = Compare::version(key.clone(), CompareOp::Equal, 0);
        let put = TxnOp::put(key, new_value, None);
        Txn::new().when([cmp]).and_then([put]).or_else([])
    } else {
        // Subsequent claim: current value must match expected epoch
        let cmp = Compare::value(key.clone(), CompareOp::Equal, expected_epoch.to_string());
        let put = TxnOp::put(key, new_value, None);
        Txn::new().when([cmp]).and_then([put]).or_else([])
    };

    let resp = client.txn(txn).await
        .map_err(|e| fence_err(format!("increment_epoch txn: {e}")))?;

    if resp.succeeded() {
        Ok(new_epoch)
    } else {
        Err(fence_err(format!(
            "epoch CAS failed for partition {}: expected {}, someone else claimed it",
            partition_id, expected_epoch
        )))
    }
}

// ── Helpers ──────────────────────────────────────────────────────────────────

#[cfg(feature = "distributed")]
fn epoch_key(partition_id: u32) -> String {
    format!("/ocg/partitions/{partition_id}/epoch")
}

#[cfg(feature = "distributed")]
fn fence_err(msg: impl Into<String>) -> CypherError {
    CypherError::Execution(ExecutionError::Internal(msg.into()))
}

// ── Tests ────────────────────────────────────────────────────────────────────

#[cfg(all(test, feature = "distributed"))]
mod tests {
    use super::*;

    #[test]
    fn test_fencing_token_validate() {
        let token = FencingToken::new(0, 5);
        assert!(token.validate(5).is_ok());
        assert!(token.validate(6).is_ok());
        assert!(token.validate(4).is_err());
        assert!(token.validate(0).is_err());
    }

    #[test]
    fn test_fencing_token_set_epoch() {
        let token = FencingToken::new(0, 1);
        assert_eq!(token.epoch(), 1);
        token.set_epoch(42);
        assert_eq!(token.epoch(), 42);
    }

    #[test]
    fn test_fencing_token_clone() {
        let token = FencingToken::new(3, 10);
        let cloned = token.clone();
        assert_eq!(cloned.partition_id(), 3);
        assert_eq!(cloned.epoch(), 10);

        // Mutation of original doesn't affect clone
        token.set_epoch(11);
        assert_eq!(cloned.epoch(), 10);
    }

    #[test]
    fn test_epoch_key_format() {
        assert_eq!(epoch_key(0), "/ocg/partitions/0/epoch");
        assert_eq!(epoch_key(42), "/ocg/partitions/42/epoch");
    }

    #[tokio::test]
    #[ignore = "requires a running etcd server on localhost:2379"]
    async fn test_read_and_increment_epoch_live() {
        let mut client = etcd_client::Client::connect(
            ["http://127.0.0.1:2379"], None
        ).await.unwrap();

        // Clean up any pre-existing key
        let _ = client.delete(epoch_key(999), None).await;

        // First read should return 0
        let epoch = read_epoch(&mut client, 999).await.unwrap();
        assert_eq!(epoch, 0);

        // Increment from 0 → 1
        let new_epoch = increment_epoch(&mut client, 999, 0).await.unwrap();
        assert_eq!(new_epoch, 1);

        // Read should now return 1
        let epoch = read_epoch(&mut client, 999).await.unwrap();
        assert_eq!(epoch, 1);

        // Increment from 1 → 2
        let new_epoch = increment_epoch(&mut client, 999, 1).await.unwrap();
        assert_eq!(new_epoch, 2);

        // Stale increment (expected 1, but actual is 2) should fail
        let result = increment_epoch(&mut client, 999, 1).await;
        assert!(result.is_err());

        // Clean up
        let _ = client.delete(epoch_key(999), None).await;
    }
}
