use std::collections::{HashMap, HashSet, VecDeque};
use std::sync::{Arc, Mutex};

/// Identifier for a node in the virtual network.
pub type NodeId = String;

/// A pending message in the virtual network.
#[derive(Debug, Clone)]
struct PendingMessage {
    from: NodeId,
    to: NodeId,
    data: Vec<u8>,
}

/// Virtual network for simulation testing.
/// Supports partition, latency injection, and random packet drop.
pub struct VirtualNetwork {
    /// Delivered message queues: node_id -> received messages.
    inbox: Arc<Mutex<HashMap<NodeId, VecDeque<(NodeId, Vec<u8>)>>>>,
    /// Pending messages not yet delivered (used for latency simulation).
    pending: Arc<Mutex<VecDeque<PendingMessage>>>,
    /// Partitioned links: (from, to) pairs that are blocked.
    partitions: Arc<Mutex<HashSet<(NodeId, NodeId)>>>,
    /// Drop probability (0.0 to 1.0).
    drop_probability: Arc<Mutex<f64>>,
}

impl VirtualNetwork {
    pub fn new() -> Self {
        Self {
            inbox: Arc::new(Mutex::new(HashMap::new())),
            pending: Arc::new(Mutex::new(VecDeque::new())),
            partitions: Arc::new(Mutex::new(HashSet::new())),
            drop_probability: Arc::new(Mutex::new(0.0)),
        }
    }

    /// Partition the network between two nodes (bidirectional).
    pub fn partition(&self, a: &str, b: &str) {
        let mut parts = self.partitions.lock().unwrap();
        parts.insert((a.to_string(), b.to_string()));
        parts.insert((b.to_string(), a.to_string()));
    }

    /// Heal the partition between two nodes (bidirectional).
    pub fn heal(&self, a: &str, b: &str) {
        let mut parts = self.partitions.lock().unwrap();
        parts.remove(&(a.to_string(), b.to_string()));
        parts.remove(&(b.to_string(), a.to_string()));
    }

    /// Heal all partitions.
    pub fn heal_all(&self) {
        self.partitions.lock().unwrap().clear();
    }

    /// Set the probability that a message will be dropped (0.0 = no drops, 1.0 = all dropped).
    pub fn set_drop_probability(&self, prob: f64) {
        *self.drop_probability.lock().unwrap() = prob.clamp(0.0, 1.0);
    }

    /// Send a message from one node to another.
    /// If the link is partitioned, the message is silently dropped.
    pub fn send(&self, from: &str, to: &str, data: Vec<u8>) -> Result<(), NetworkError> {
        // Check for partition.
        {
            let parts = self.partitions.lock().unwrap();
            if parts.contains(&(from.to_string(), to.to_string())) {
                return Ok(()); // Silently dropped.
            }
        }

        // Check for random drop.
        {
            let prob = *self.drop_probability.lock().unwrap();
            if prob > 0.0 {
                let random: f64 = simple_random();
                if random < prob {
                    return Ok(()); // Randomly dropped.
                }
            }
        }

        // Queue the message for delivery.
        let mut pending = self.pending.lock().unwrap();
        pending.push_back(PendingMessage {
            from: from.to_string(),
            to: to.to_string(),
            data,
        });

        Ok(())
    }

    /// Deliver all pending messages to their inboxes.
    /// Call this to simulate message delivery (allows controlling when messages arrive).
    pub fn deliver_pending(&self) {
        let messages: Vec<PendingMessage> = {
            let mut pending = self.pending.lock().unwrap();
            pending.drain(..).collect()
        };

        let mut inbox = self.inbox.lock().unwrap();
        for msg in messages {
            inbox
                .entry(msg.to.clone())
                .or_default()
                .push_back((msg.from, msg.data));
        }
    }

    /// Receive a message for a given node. Returns None if no messages are available.
    pub fn recv(&self, node: &str) -> Option<(NodeId, Vec<u8>)> {
        let mut inbox = self.inbox.lock().unwrap();
        inbox.get_mut(node).and_then(|q| q.pop_front())
    }

    /// Get the number of pending (undelivered) messages.
    pub fn pending_count(&self) -> usize {
        self.pending.lock().unwrap().len()
    }

    /// Get the number of messages in a node's inbox.
    pub fn inbox_count(&self, node: &str) -> usize {
        self.inbox
            .lock()
            .unwrap()
            .get(node)
            .map(|q| q.len())
            .unwrap_or(0)
    }
}

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

/// Simple deterministic pseudo-random based on thread-local state.
fn simple_random() -> f64 {
    use std::cell::Cell;
    thread_local! {
        static STATE: Cell<u64> = const { Cell::new(12345) };
    }
    STATE.with(|s| {
        let mut state = s.get();
        state ^= state << 13;
        state ^= state >> 7;
        state ^= state << 17;
        s.set(state);
        (state % 10000) as f64 / 10000.0
    })
}

#[derive(Debug, thiserror::Error)]
pub enum NetworkError {
    #[error("node '{0}' not reachable")]
    Unreachable(String),
}

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

    #[test]
    fn test_send_and_receive() {
        let net = VirtualNetwork::new();

        net.send("node-1", "node-2", b"hello".to_vec()).unwrap();
        net.deliver_pending();

        let (from, data) = net.recv("node-2").unwrap();
        assert_eq!(from, "node-1");
        assert_eq!(data, b"hello");
    }

    #[test]
    fn test_no_messages_returns_none() {
        let net = VirtualNetwork::new();
        assert!(net.recv("node-1").is_none());
    }

    #[test]
    fn test_partition_drops_messages() {
        let net = VirtualNetwork::new();

        net.partition("node-1", "node-2");

        net.send("node-1", "node-2", b"hello".to_vec()).unwrap();
        net.deliver_pending();

        assert!(net.recv("node-2").is_none());
    }

    #[test]
    fn test_partition_is_bidirectional() {
        let net = VirtualNetwork::new();

        net.partition("node-1", "node-2");

        net.send("node-1", "node-2", b"a->b".to_vec()).unwrap();
        net.send("node-2", "node-1", b"b->a".to_vec()).unwrap();
        net.deliver_pending();

        assert!(net.recv("node-2").is_none());
        assert!(net.recv("node-1").is_none());
    }

    #[test]
    fn test_heal_restores_communication() {
        let net = VirtualNetwork::new();

        net.partition("node-1", "node-2");
        net.send("node-1", "node-2", b"before".to_vec()).unwrap();
        net.deliver_pending();
        assert!(net.recv("node-2").is_none());

        net.heal("node-1", "node-2");
        net.send("node-1", "node-2", b"after".to_vec()).unwrap();
        net.deliver_pending();

        let (_, data) = net.recv("node-2").unwrap();
        assert_eq!(data, b"after");
    }

    #[test]
    fn test_heal_all() {
        let net = VirtualNetwork::new();

        net.partition("a", "b");
        net.partition("a", "c");
        net.heal_all();

        net.send("a", "b", b"msg".to_vec()).unwrap();
        net.send("a", "c", b"msg".to_vec()).unwrap();
        net.deliver_pending();

        assert!(net.recv("b").is_some());
        assert!(net.recv("c").is_some());
    }

    #[test]
    fn test_multiple_messages_ordered() {
        let net = VirtualNetwork::new();

        for i in 0..5 {
            net.send("a", "b", format!("msg-{i}").into_bytes())
                .unwrap();
        }
        net.deliver_pending();

        for i in 0..5 {
            let (_, data) = net.recv("b").unwrap();
            assert_eq!(data, format!("msg-{i}").as_bytes());
        }
        assert!(net.recv("b").is_none());
    }

    #[test]
    fn test_pending_and_inbox_counts() {
        let net = VirtualNetwork::new();

        net.send("a", "b", b"1".to_vec()).unwrap();
        net.send("a", "b", b"2".to_vec()).unwrap();

        assert_eq!(net.pending_count(), 2);
        assert_eq!(net.inbox_count("b"), 0);

        net.deliver_pending();

        assert_eq!(net.pending_count(), 0);
        assert_eq!(net.inbox_count("b"), 2);
    }

    #[test]
    fn test_partition_does_not_affect_other_links() {
        let net = VirtualNetwork::new();

        net.partition("a", "b");

        // a -> c should still work.
        net.send("a", "c", b"hello".to_vec()).unwrap();
        net.deliver_pending();

        assert!(net.recv("c").is_some());
    }

    #[test]
    fn test_drop_probability_all() {
        let net = VirtualNetwork::new();
        net.set_drop_probability(1.0);

        for _ in 0..10 {
            net.send("a", "b", b"msg".to_vec()).unwrap();
        }
        net.deliver_pending();

        // All messages should be dropped.
        assert_eq!(net.inbox_count("b"), 0);
    }

    #[test]
    fn test_drop_probability_none() {
        let net = VirtualNetwork::new();
        net.set_drop_probability(0.0);

        for _ in 0..10 {
            net.send("a", "b", b"msg".to_vec()).unwrap();
        }
        net.deliver_pending();

        // No messages should be dropped.
        assert_eq!(net.inbox_count("b"), 10);
    }
}