//! Network fault injection for simulating partitions, latency, and drops.
//!
//! The [`NetworkController`] intercepts communication between nodes and
//! applies configurable fault patterns:
//!
//! - **Partitions**: Completely block communication between node groups
//! - **Latency**: Add delays to specific routes
//! - **Drops**: Probabilistically drop messages

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

/// Controller for simulating network faults between nodes.
///
/// Uses DashMap for lock-free concurrent access from multiple test threads.
/// All operations are indexed by node index pairs (from, to).
///
/// # Thread Safety
///
/// All methods are safe to call concurrently from multiple threads.
/// State changes are immediately visible to other threads.
#[derive(Debug)]
pub struct NetworkController {
    /// Which (from, to) pairs are partitioned.
    /// True means communication is blocked.
    partitions: DashMap<(usize, usize), bool>,

    /// Latency per directional pair in milliseconds.
    latencies: DashMap<(usize, usize), Duration>,

    /// Drop probability per pair (0.0 to 1.0).
    drop_rates: DashMap<(usize, usize), f64>,

    /// Total messages that were dropped due to partition or drop rate.
    pub messages_dropped: AtomicU64,

    /// Total messages that were delayed due to latency injection.
    pub messages_delayed: AtomicU64,

    /// Total number of partitions created (for metrics).
    partition_events: AtomicU64,
}

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

impl NetworkController {
    /// Creates a new network controller with no faults configured.
    pub fn new() -> Self {
        Self {
            partitions: DashMap::new(),
            latencies: DashMap::new(),
            drop_rates: DashMap::new(),
            messages_dropped: AtomicU64::new(0),
            messages_delayed: AtomicU64::new(0),
            partition_events: AtomicU64::new(0),
        }
    }

    /// Creates a network partition between two groups of nodes.
    ///
    /// Nodes in `group_a` cannot communicate with nodes in `group_b` and vice versa.
    /// Communication within each group is unaffected.
    ///
    /// # Arguments
    ///
    /// * `group_a` - Node indices in the first partition
    /// * `group_b` - Node indices in the second partition
    ///
    /// # Example
    ///
    /// ```ignore
    /// // Partition [0, 1, 2] from [3, 4] in a 5-node cluster
    /// network.partition(&[0, 1, 2], &[3, 4]);
    /// ```
    #[instrument(skip(self))]
    pub fn partition(&self, group_a: &[usize], group_b: &[usize]) {
        info!(
            group_a = ?group_a,
            group_b = ?group_b,
            "Creating network partition"
        );

        // Block communication from A to B
        for &a in group_a {
            for &b in group_b {
                self.partitions.insert((a, b), true);
                self.partitions.insert((b, a), true);
            }
        }

        self.partition_events.fetch_add(1, Ordering::Relaxed);
    }

    /// Heals all network partitions, restoring full connectivity.
    ///
    /// Clears all partition state but preserves latency and drop rate settings.
    #[instrument(skip(self))]
    pub fn heal(&self) {
        info!("Healing all network partitions");
        self.partitions.clear();
    }

    /// Creates an asymmetric partition: blocks communication from `from` to `to`
    /// but not the reverse direction.
    ///
    /// # Arguments
    ///
    /// * `from` - Source node index that cannot send
    /// * `to` - Destination node index that cannot receive from source
    pub fn partition_one_way(&self, from: usize, to: usize) {
        debug!(from, to, "Creating one-way partition");
        self.partitions.insert((from, to), true);
        self.partition_events.fetch_add(1, Ordering::Relaxed);
    }

    /// Checks if communication from `from` to `to` is blocked by a partition.
    ///
    /// # Returns
    ///
    /// `true` if the nodes are partitioned, `false` if communication is allowed.
    pub fn is_partitioned(&self, from: usize, to: usize) -> bool {
        self.partitions.get(&(from, to)).map(|v| *v).unwrap_or(false)
    }

    /// Sets a latency delay for communication from `from` to `to`.
    ///
    /// Latency is unidirectional - set both directions for symmetric delay.
    ///
    /// # Arguments
    ///
    /// * `from` - Source node index
    /// * `to` - Destination node index
    /// * `delay` - Duration to delay messages
    #[instrument(skip(self))]
    pub fn set_latency(&self, from: usize, to: usize, delay: Duration) {
        debug!(from, to, delay_ms = delay.as_millis(), "Setting network latency");
        self.latencies.insert((from, to), delay);
    }

    /// Sets a bidirectional latency delay between two nodes.
    ///
    /// Convenience method that sets the same delay in both directions.
    pub fn set_latency_bidirectional(&self, node_a: usize, node_b: usize, delay: Duration) {
        self.set_latency(node_a, node_b, delay);
        self.set_latency(node_b, node_a, delay);
    }

    /// Clears latency for a specific route.
    pub fn clear_latency(&self, from: usize, to: usize) {
        self.latencies.remove(&(from, to));
    }

    /// Gets the latency configured for a route, if any.
    pub fn get_latency(&self, from: usize, to: usize) -> Option<Duration> {
        self.latencies.get(&(from, to)).map(|v| *v)
    }

    /// Sets a drop probability for communication from `from` to `to`.
    ///
    /// # Arguments
    ///
    /// * `from` - Source node index
    /// * `to` - Destination node index
    /// * `rate` - Probability of dropping a message (0.0 to 1.0)
    ///
    /// # Panics
    ///
    /// Debug asserts that rate is in [0.0, 1.0].
    #[instrument(skip(self))]
    pub fn set_drop_rate(&self, from: usize, to: usize, rate: f64) {
        debug_assert!((0.0..=1.0).contains(&rate), "Drop rate must be in [0.0, 1.0]");
        debug!(from, to, rate, "Setting drop rate");
        self.drop_rates.insert((from, to), rate);
    }

    /// Clears the drop rate for a specific route.
    pub fn clear_drop_rate(&self, from: usize, to: usize) {
        self.drop_rates.remove(&(from, to));
    }

    /// Checks if a message from `from` to `to` should be dropped.
    ///
    /// Takes into account both partitions and probabilistic drops.
    /// Updates the `messages_dropped` counter if the message is dropped.
    ///
    /// # Returns
    ///
    /// `true` if the message should be dropped, `false` if it should be delivered.
    pub fn should_drop(&self, from: usize, to: usize) -> bool {
        // Check partition first
        if self.is_partitioned(from, to) {
            self.messages_dropped.fetch_add(1, Ordering::Relaxed);
            return true;
        }

        // Check drop rate
        if let Some(rate) = self.drop_rates.get(&(from, to)) {
            if rand::random::<f64>() < *rate {
                self.messages_dropped.fetch_add(1, Ordering::Relaxed);
                return true;
            }
        }

        false
    }

    /// Applies network effects (latency) for a message from `from` to `to`.
    ///
    /// This is an async method that will delay based on configured latency.
    /// Call `should_drop` first to check if the message should be sent at all.
    ///
    /// # Returns
    ///
    /// Returns `true` if latency was applied.
    pub async fn apply_effects(&self, from: usize, to: usize) -> bool {
        if let Some(delay) = self.get_latency(from, to) {
            self.messages_delayed.fetch_add(1, Ordering::Relaxed);
            tokio::time::sleep(delay).await;
            true
        } else {
            false
        }
    }

    /// Returns the total number of messages dropped.
    pub fn total_dropped(&self) -> u64 {
        self.messages_dropped.load(Ordering::Relaxed)
    }

    /// Returns the total number of messages delayed.
    pub fn total_delayed(&self) -> u64 {
        self.messages_delayed.load(Ordering::Relaxed)
    }

    /// Returns the total number of partition events.
    pub fn partition_count(&self) -> u64 {
        self.partition_events.load(Ordering::Relaxed)
    }

    /// Resets all network fault configurations and counters.
    pub fn reset(&self) {
        self.partitions.clear();
        self.latencies.clear();
        self.drop_rates.clear();
        self.messages_dropped.store(0, Ordering::Relaxed);
        self.messages_delayed.store(0, Ordering::Relaxed);
        self.partition_events.store(0, Ordering::Relaxed);
    }
}

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

    #[test]
    fn test_partition_and_heal() {
        let controller = NetworkController::new();

        // Initially not partitioned
        assert!(!controller.is_partitioned(0, 1));
        assert!(!controller.is_partitioned(1, 0));

        // Create partition
        controller.partition(&[0, 1], &[2, 3]);

        // Check partition state
        assert!(controller.is_partitioned(0, 2));
        assert!(controller.is_partitioned(0, 3));
        assert!(controller.is_partitioned(1, 2));
        assert!(controller.is_partitioned(1, 3));
        assert!(controller.is_partitioned(2, 0));
        assert!(controller.is_partitioned(3, 1));

        // Within-group communication not affected
        assert!(!controller.is_partitioned(0, 1));
        assert!(!controller.is_partitioned(2, 3));

        // Heal
        controller.heal();

        // All communication restored
        assert!(!controller.is_partitioned(0, 2));
        assert!(!controller.is_partitioned(2, 0));
    }

    #[test]
    fn test_should_drop_partition() {
        let controller = NetworkController::new();
        controller.partition(&[0], &[1]);

        assert!(controller.should_drop(0, 1));
        assert!(controller.should_drop(1, 0));
        assert!(!controller.should_drop(0, 2));

        assert_eq!(controller.total_dropped(), 2);
    }

    #[test]
    fn test_latency() {
        let controller = NetworkController::new();

        assert!(controller.get_latency(0, 1).is_none());

        controller.set_latency(0, 1, Duration::from_millis(100));
        assert_eq!(controller.get_latency(0, 1), Some(Duration::from_millis(100)));

        // Unidirectional by default
        assert!(controller.get_latency(1, 0).is_none());

        controller.clear_latency(0, 1);
        assert!(controller.get_latency(0, 1).is_none());
    }

    #[test]
    fn test_latency_bidirectional() {
        let controller = NetworkController::new();

        controller.set_latency_bidirectional(0, 1, Duration::from_millis(50));

        assert_eq!(controller.get_latency(0, 1), Some(Duration::from_millis(50)));
        assert_eq!(controller.get_latency(1, 0), Some(Duration::from_millis(50)));
    }

    #[test]
    fn test_reset() {
        let controller = NetworkController::new();

        controller.partition(&[0], &[1]);
        controller.set_latency(0, 1, Duration::from_millis(100));
        controller.set_drop_rate(0, 1, 0.5);
        let _ = controller.should_drop(0, 1); // Increment counter

        controller.reset();

        assert!(!controller.is_partitioned(0, 1));
        assert!(controller.get_latency(0, 1).is_none());
        assert_eq!(controller.total_dropped(), 0);
    }
}
