//! Clock skew injection for testing HLC behavior under clock drift.
//!
//! The [`ClockController`] allows injecting clock offsets into individual nodes
//! to test how the cluster handles clock skew scenarios.

use dashmap::DashMap;
use std::sync::atomic::{AtomicI64, Ordering};
use std::sync::Arc;
use std::time::Duration;
use tracing::{debug, info, instrument};

/// Controller for injecting clock skew into test nodes.
///
/// Tracks per-node time offsets that can be applied when creating HLC timestamps.
/// Offsets can be positive (clock running ahead) or negative (clock running behind).
///
/// # Thread Safety
///
/// All methods are safe to call concurrently from multiple threads.
#[derive(Debug)]
pub struct ClockController {
    /// Per-node offset in milliseconds.
    /// Positive = clock ahead, negative = clock behind.
    node_offsets: DashMap<usize, i64>,

    /// Global offset applied to all nodes.
    global_offset_ms: AtomicI64,
}

impl Default for ClockController {
    fn default() -> Self {
        Self::new()
    }
}

impl ClockController {
    /// Creates a new clock controller with no skew configured.
    pub fn new() -> Self {
        Self { node_offsets: DashMap::new(), global_offset_ms: AtomicI64::new(0) }
    }

    /// Injects clock skew for a specific node.
    ///
    /// # Arguments
    ///
    /// * `node` - Node index to apply skew to
    /// * `offset_ms` - Offset in milliseconds (positive = ahead, negative = behind)
    ///
    /// # Example
    ///
    /// ```ignore
    /// // Node 1's clock is 5 seconds ahead
    /// clock.inject_skew(1, 5000);
    ///
    /// // Node 2's clock is 3 seconds behind
    /// clock.inject_skew(2, -3000);
    /// ```
    #[instrument(skip(self))]
    pub fn inject_skew(&self, node: usize, offset_ms: i64) {
        info!(node, offset_ms, "Injecting clock skew");
        self.node_offsets.insert(node, offset_ms);
    }

    /// Gets the clock offset for a specific node.
    ///
    /// Returns the node-specific offset plus the global offset.
    pub fn get_offset(&self, node: usize) -> i64 {
        let node_offset = self.node_offsets.get(&node).map(|v| *v).unwrap_or(0);
        let global_offset = self.global_offset_ms.load(Ordering::Relaxed);
        node_offset + global_offset
    }

    /// Gets only the node-specific offset (excludes global offset).
    pub fn get_node_offset(&self, node: usize) -> i64 {
        self.node_offsets.get(&node).map(|v| *v).unwrap_or(0)
    }

    /// Sets a global clock offset applied to all nodes.
    ///
    /// Useful for simulating system-wide time jumps.
    pub fn set_global_offset(&self, offset_ms: i64) {
        debug!(offset_ms, "Setting global clock offset");
        self.global_offset_ms.store(offset_ms, Ordering::Relaxed);
    }

    /// Gets the global clock offset.
    pub fn get_global_offset(&self) -> i64 {
        self.global_offset_ms.load(Ordering::Relaxed)
    }

    /// Clears clock skew for a specific node.
    pub fn clear_skew(&self, node: usize) {
        self.node_offsets.remove(&node);
    }

    /// Clears all clock skew (node-specific and global).
    #[instrument(skip(self))]
    pub fn clear_all_skew(&self) {
        info!("Clearing all clock skew");
        self.node_offsets.clear();
        self.global_offset_ms.store(0, Ordering::Relaxed);
    }

    /// Creates a `uhlc::HLC` clock with skew applied for a specific node.
    ///
    /// The returned HLC will generate timestamps that are offset by the
    /// configured skew amount.
    ///
    /// # Note
    ///
    /// The uhlc crate doesn't directly support custom time sources in a way
    /// that allows runtime skew injection. Instead, we apply the skew when
    /// creating `HlcTimestamp` values via `apply_skew_to_timestamp`.
    ///
    /// # Important
    ///
    /// This method clones the controller. If you need the SkewedHlc to share
    /// the same controller (so that `inject_skew` affects it), use
    /// `SkewedHlc::new` directly with an `Arc<ClockController>`.
    pub fn create_hlc(&self, node: usize) -> SkewedHlc {
        SkewedHlc::new(Arc::new(self.clone()), node)
    }

    /// Creates a `SkewedHlc` that shares this controller via Arc.
    ///
    /// Changes to skew via `inject_skew` will affect the returned HLC.
    pub fn create_hlc_shared(self: &Arc<Self>, node: usize) -> SkewedHlc {
        SkewedHlc::new(self.clone(), node)
    }

    /// Applies skew to an existing HLC timestamp.
    ///
    /// This modifies the NTP64 time component to simulate clock drift.
    pub fn apply_skew_to_ntp64(&self, node: usize, ntp64: u64) -> u64 {
        let offset = self.get_offset(node);

        if offset == 0 {
            return ntp64;
        }

        // NTP64: upper 32 bits = seconds, lower 32 bits = fractional seconds
        // Convert offset from milliseconds to NTP64 units
        let offset_ntp64 = ms_to_ntp64_offset(offset);

        if offset >= 0 {
            ntp64.saturating_add(offset_ntp64 as u64)
        } else {
            ntp64.saturating_sub((-offset_ntp64) as u64)
        }
    }

    /// Checks if any significant skew (>500ms) exists between two nodes.
    pub fn has_significant_skew(&self, node_a: usize, node_b: usize) -> bool {
        let offset_a = self.get_offset(node_a);
        let offset_b = self.get_offset(node_b);
        (offset_a - offset_b).unsigned_abs() > 500
    }

    /// Gets the maximum absolute skew across all configured nodes.
    pub fn max_skew_ms(&self) -> i64 {
        let mut max_skew = 0i64;
        for entry in self.node_offsets.iter() {
            max_skew = max_skew.max(entry.value().abs());
        }
        max_skew.max(self.global_offset_ms.load(Ordering::Relaxed).abs())
    }
}

impl Clone for ClockController {
    fn clone(&self) -> Self {
        let new = Self::new();
        for entry in self.node_offsets.iter() {
            new.node_offsets.insert(*entry.key(), *entry.value());
        }
        new.global_offset_ms
            .store(self.global_offset_ms.load(Ordering::Relaxed), Ordering::Relaxed);
        new
    }
}

/// Converts milliseconds to NTP64 offset.
///
/// NTP64 format: upper 32 bits are seconds, lower 32 bits are fractional seconds.
fn ms_to_ntp64_offset(ms: i64) -> i64 {
    // 1 second = 2^32 NTP64 units
    // 1 ms = 2^32 / 1000 NTP64 units ≈ 4,294,967 units
    let seconds = ms / 1000;
    let remaining_ms = ms % 1000;

    let seconds_ntp64 = seconds << 32;
    // remaining_ms * 2^32 / 1000
    let fraction_ntp64 = (remaining_ms.unsigned_abs() << 32) / 1000;

    if ms >= 0 {
        seconds_ntp64 + fraction_ntp64 as i64
    } else {
        seconds_ntp64 - fraction_ntp64 as i64
    }
}

/// A wrapper around `uhlc::HLC` that applies clock skew.
///
/// Since `uhlc` doesn't support custom time sources at runtime, this wrapper
/// applies skew by adjusting the timestamps after generation.
pub struct SkewedHlc {
    inner: uhlc::HLC,
    controller: Arc<ClockController>,
    node: usize,
}

impl SkewedHlc {
    /// Creates a new skewed HLC for the given node.
    pub fn new(controller: Arc<ClockController>, node: usize) -> Self {
        let inner = uhlc::HLCBuilder::new().build();
        Self { inner, controller, node }
    }

    /// Generates a new timestamp with skew applied.
    ///
    /// Returns an `HlcTimestamp` with the clock offset applied.
    pub fn new_timestamp(&self) -> stemedb_core::types::HlcTimestamp {
        let ts = self.inner.new_timestamp();
        let original_ntp64 = ts.get_time().as_u64();
        let skewed_ntp64 = self.controller.apply_skew_to_ntp64(self.node, original_ntp64);

        stemedb_core::types::HlcTimestamp::new(skewed_ntp64, ts.get_id().to_le_bytes())
    }

    /// Updates the clock with a received timestamp.
    ///
    /// This should be called when receiving timestamps from other nodes
    /// to maintain the HLC invariant.
    pub fn update_with_timestamp(&self, ts: &uhlc::Timestamp) -> Result<(), String> {
        self.inner.update_with_timestamp(ts).map_err(|e| e.to_string())
    }

    /// Gets the current offset for this node.
    pub fn current_offset_ms(&self) -> i64 {
        self.controller.get_offset(self.node)
    }
}

impl std::fmt::Debug for SkewedHlc {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("SkewedHlc")
            .field("node", &self.node)
            .field("offset_ms", &self.controller.get_offset(self.node))
            .finish()
    }
}

/// Calculates the simulated wall clock time with skew applied.
///
/// Returns Duration since Unix epoch with the skew offset.
#[allow(dead_code)]
pub fn skewed_now(controller: &ClockController, node: usize) -> Duration {
    let now = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap_or(Duration::ZERO);

    let offset_ms = controller.get_offset(node);

    if offset_ms >= 0 {
        now + Duration::from_millis(offset_ms as u64)
    } else {
        now.saturating_sub(Duration::from_millis((-offset_ms) as u64))
    }
}

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

    #[test]
    fn test_inject_and_get_skew() {
        let controller = ClockController::new();

        // Initially no skew
        assert_eq!(controller.get_offset(0), 0);
        assert_eq!(controller.get_offset(1), 0);

        // Inject skew
        controller.inject_skew(0, 5000);
        controller.inject_skew(1, -3000);

        assert_eq!(controller.get_offset(0), 5000);
        assert_eq!(controller.get_offset(1), -3000);
        assert_eq!(controller.get_offset(2), 0); // Unset node
    }

    #[test]
    fn test_global_offset() {
        let controller = ClockController::new();

        controller.inject_skew(0, 1000);
        controller.set_global_offset(500);

        // Node 0 gets both offsets
        assert_eq!(controller.get_offset(0), 1500);

        // Other nodes only get global
        assert_eq!(controller.get_offset(1), 500);
    }

    #[test]
    fn test_clear_skew() {
        let controller = ClockController::new();

        controller.inject_skew(0, 5000);
        controller.inject_skew(1, 3000);
        controller.set_global_offset(1000);

        controller.clear_skew(0);
        assert_eq!(controller.get_node_offset(0), 0);
        assert_eq!(controller.get_offset(0), 1000); // Still has global

        controller.clear_all_skew();
        assert_eq!(controller.get_offset(0), 0);
        assert_eq!(controller.get_offset(1), 0);
        assert_eq!(controller.get_global_offset(), 0);
    }

    #[test]
    fn test_has_significant_skew() {
        let controller = ClockController::new();

        // No skew
        assert!(!controller.has_significant_skew(0, 1));

        // 400ms difference - not significant
        controller.inject_skew(0, 200);
        controller.inject_skew(1, -200);
        assert!(!controller.has_significant_skew(0, 1));

        // 600ms difference - significant
        controller.inject_skew(0, 500);
        controller.inject_skew(1, -200);
        assert!(controller.has_significant_skew(0, 1));
    }

    #[test]
    fn test_max_skew() {
        let controller = ClockController::new();

        assert_eq!(controller.max_skew_ms(), 0);

        controller.inject_skew(0, 1000);
        controller.inject_skew(1, -2000);
        assert_eq!(controller.max_skew_ms(), 2000);

        controller.set_global_offset(-3000);
        assert_eq!(controller.max_skew_ms(), 3000);
    }

    #[test]
    fn test_apply_skew_to_ntp64() {
        let controller = ClockController::new();

        let original: u64 = 100 << 32; // 100 seconds in NTP64

        // No skew
        assert_eq!(controller.apply_skew_to_ntp64(0, original), original);

        // +1 second skew
        controller.inject_skew(0, 1000);
        let skewed = controller.apply_skew_to_ntp64(0, original);
        assert!(skewed > original);
        // Should be approximately 1 second more
        let diff_seconds = (skewed >> 32) - (original >> 32);
        assert_eq!(diff_seconds, 1);

        // -1 second skew
        controller.clear_all_skew();
        controller.inject_skew(1, -1000);
        let skewed = controller.apply_skew_to_ntp64(1, original);
        assert!(skewed < original);
        let diff_seconds = (original >> 32) - (skewed >> 32);
        assert_eq!(diff_seconds, 1);
    }

    #[test]
    fn test_skewed_hlc() {
        let controller = Arc::new(ClockController::new());
        controller.inject_skew(0, 5000);

        let hlc = SkewedHlc::new(controller.clone(), 0);
        let ts1 = hlc.new_timestamp();
        let ts2 = hlc.new_timestamp();

        // Timestamps should be monotonically increasing
        assert!(ts2 >= ts1);

        // Offset should be reported
        assert_eq!(hlc.current_offset_ms(), 5000);
    }

    #[test]
    fn test_clone() {
        let controller = ClockController::new();
        controller.inject_skew(0, 1000);
        controller.set_global_offset(500);

        let cloned = controller.clone();
        assert_eq!(cloned.get_offset(0), 1500);
        assert_eq!(cloned.get_global_offset(), 500);

        // Modifications to original don't affect clone
        controller.inject_skew(0, 2000);
        assert_eq!(cloned.get_node_offset(0), 1000);
    }
}
