#![no_main]

use arbitrary::Arbitrary;
use bytes::Bytes;
use commonware_codec::codec::FixedSize;
use commonware_cryptography::{ed25519, PrivateKeyExt, Signer};
use commonware_p2p::{
    simulated, Channel, Receiver as ReceiverTrait, Recipients, Sender as SenderTrait,
};
use commonware_runtime::{deterministic, Clock, Metrics, Runner};
use libfuzzer_sys::fuzz_target;
use rand::Rng;
use std::{
    collections::{hash_map, HashMap, HashSet, VecDeque},
    time::Duration,
};

const MAX_OPERATIONS: usize = 50;
const MAX_PEERS: usize = 16;
const MIN_PEERS: usize = 2;
const MAX_MSG_SIZE: usize = 1024 * 1024; // 1MB
const MAX_SLEEP_DURATION: u64 = 1000; // milliseconds

/// Operations that can be performed on the simulated p2p network during fuzzing.
#[derive(Debug, Arbitrary)]
enum Operation {
    /// Register a new channel for a peer to send/receive messages.
    RegisterChannel {
        /// Index of the peer registering the channel.
        peer_idx: u8,
        /// Channel identifier
        channel_id: u8,
    },
    /// Send a message from one peer to another on a specific channel.
    SendMessage {
        /// Index of the peer sending the message.
        peer_idx: u8,
        /// Channel to send the message on.
        channel_id: u8,
        /// Index of the recipient peer.
        to_idx: u8,
        /// Size of the message payload in bytes.
        msg_size: usize,
    },
    /// Attempt to receive pending messages from all peers.
    ReceiveMessages,
    /// Add or update a network link between two peers with specific characteristics.
    AddLink {
        /// Index of the sender peer.
        from_idx: u8,
        /// Index of the receiver peer.
        to_idx: u8,
        /// Base latency in milliseconds.
        latency_ms: u16,
        /// Latency jitter in milliseconds.
        jitter: u16,
        /// Success rate (0-255 maps to 0.0-1.0).
        success_rate: u8,
    },
    /// Remove a network link between two peers.
    RemoveLink {
        /// Index of the sender peer.
        from_idx: u8,
        /// Index of the receiver peer.
        to_idx: u8,
    },
}

/// Input generated by the fuzzer for testing the simulated p2p network.
#[derive(Debug)]
struct FuzzInput {
    /// Random seed for deterministic execution.
    seed: u64,
    /// Sequence of operations to execute on the network.
    ///
    /// Length is in the range [1, MAX_OPERATIONS].
    operations: Vec<Operation>,
    /// Number of peers to create in the network.
    ///
    /// Must be in the range [MIN_PEERS, MAX_PEERS].
    num_peers: u8,
}

impl<'a> Arbitrary<'a> for FuzzInput {
    fn arbitrary(u: &mut arbitrary::Unstructured<'a>) -> arbitrary::Result<Self> {
        let seed = u.arbitrary()?;
        let num_operations = u.int_in_range(1..=MAX_OPERATIONS)?;
        let mut operations = Vec::with_capacity(num_operations);
        for _ in 0..num_operations {
            operations.push(u.arbitrary()?);
        }
        let num_peers = u.int_in_range(MIN_PEERS..=MAX_PEERS)? as u8;
        Ok(FuzzInput {
            seed,
            operations,
            num_peers,
        })
    }
}

/// Main fuzzing entry point for testing the simulated p2p network.
///
/// Creates a deterministic network and executes a sequence of random operations,
/// verifying that received messages match expectations.
fn fuzz(input: FuzzInput) {
    let num_peers = input.num_peers;

    let p2p_cfg = simulated::Config {
        max_size: MAX_MSG_SIZE,
        disconnect_on_block: false,
        tracked_peer_sets: None,
    };

    let executor = deterministic::Runner::seeded(input.seed);
    executor.start(|mut context| async move {
        // Generate peer identities
        let mut peer_pks = Vec::new();
        for _ in 0..num_peers {
            let private_key = ed25519::PrivateKey::from_seed(context.gen());
            peer_pks.push(private_key.public_key());
        }

        // Create the simulated network and oracle for controlling it
        let (network, mut oracle) = simulated::Network::new(context.with_label("network"), p2p_cfg);
        let _network_handle = network.start();

        // Track registered channels: (peer_idx, channel_id) -> (sender, receiver)
        // Each peer can register multiple channels for message segregation
        // The receiver gets messages from ALL senders on that channel, not per-sender streams
        let mut channels: HashMap<
            (usize, u8),
            (
                commonware_p2p::simulated::Sender<ed25519::PublicKey>,
                commonware_p2p::simulated::Receiver<ed25519::PublicKey>,
            ),
        > = HashMap::new();

        // Track expected messages: (to_idx, sender_pk, channel_id) -> queue of messages
        // Messages may be dropped (unreliable links) but those delivered must match expectations
        let mut expected_msgs: HashMap<(usize, ed25519::PublicKey, u8), VecDeque<Bytes>> = HashMap::new();

        for op in input.operations.into_iter() {
            match op {
                Operation::RegisterChannel {
                    peer_idx,
                    channel_id,
                } => {
                    // Normalize peer index to valid range
                    let idx = (peer_idx as usize) % peer_pks.len();

                    // Only register if not already registered
                    if let hash_map::Entry::Vacant(e) = channels.entry((idx, channel_id)) {
                        if let Ok((sender, receiver)) =
                            oracle.control(peer_pks[idx].clone()).register(channel_id as u64).await
                        {
                            e.insert((sender, receiver));
                        }
                    }
                }

                Operation::SendMessage {
                    peer_idx,
                    channel_id,
                    to_idx,
                    msg_size,
                } => {
                    // Normalize sender and receiver indices to valid ranges
                    let from_idx = (peer_idx as usize) % peer_pks.len();
                    let to_idx = (to_idx as usize) % peer_pks.len();

                    // Clamp message size to not exceed max (accounting for channel overhead)
                    let msg_size = msg_size.clamp(0, MAX_MSG_SIZE - Channel::SIZE);

                    // Skip if receiver hasn't registered this channel - they won't be able to receive
                    if !channels.contains_key(&(to_idx, channel_id)) {
                        continue;
                    }

                    // Skip if sender channel not registered
                    let Some((sender, _)) = channels.get_mut(&(from_idx, channel_id)) else {
                        continue;
                    };

                    // Generate random message payload
                    let mut bytes = vec![0u8; msg_size];
                    context.fill(&mut bytes[..]);
                    let message = Bytes::from(bytes);

                    // Attempt to send the message
                    // Note: Success only means accepted for transmission, not guaranteed delivery
                    let sent = sender
                        .send(
                            Recipients::One(peer_pks[to_idx].clone()),
                            message.clone(),
                            true,
                        )
                        .await
                        .is_ok();

                    // Track message as expected only if send was accepted
                    // Note: Message may still be dropped by unreliable link
                    if sent {
                        expected_msgs
                            .entry((to_idx, peer_pks[from_idx].clone(), channel_id))
                            .or_default()
                            .push_back(message);
                    }
                }

                Operation::ReceiveMessages => {
                    // Attempt to receive one message from each receiver channel with pending messages
                    let receiver_channels: HashSet<(usize, u8)> = expected_msgs
                        .keys()
                        .map(|(to_idx, _sender_pk, channel_id)| (*to_idx, *channel_id))
                        .collect();

                    // Each (to_idx, channel_id) is a distinct mailbox receiving from all senders
                    for (to_idx, channel_id) in receiver_channels {
                        // Skip if receiver hasn't registered this channel
                        let Some((_, receiver)) = channels.get_mut(&(to_idx, channel_id)) else {
                            continue;
                        };

                        // Try to receive one message with timeout
                        commonware_macros::select! {
                            result = receiver.recv() => {
                                let Ok((sender_pk, message)) = result else {
                                    continue; // Receive error
                                };

                                // Find the expected queue for this sender-receiver-channel tuple
                                let expected_msgs_key = (to_idx, sender_pk.clone(), channel_id);
                                let queue = expected_msgs.get_mut(&expected_msgs_key).expect("Expected queue not found");

                                // Find message in expected queue
                                // Messages can be dropped, but if received they must be in order
                                if let Some(pos) = queue.iter().position(|m| m == &message) {
                                    // Remove all messages up to and including this one
                                    // Messages before it were implicitly dropped, this one is received
                                    for _ in 0..=pos {
                                        queue.pop_front();
                                    }

                                    // Clean up empty queue
                                    if queue.is_empty() {
                                        expected_msgs.remove(&expected_msgs_key);
                                    }
                                } else {
                                    panic!(
                                        "Unexpected message from sender {} to receiver {} on channel {}. Message len: {}",
                                        sender_pk, to_idx, channel_id, message.len()
                                    );
                                }
                            },
                            _ = context.sleep(Duration::from_millis(MAX_SLEEP_DURATION)) => {
                                continue; // Timeout - message may not have arrived yet
                            }
                        }
                    }
                }

                Operation::AddLink {
                    from_idx,
                    to_idx,
                    latency_ms,
                    jitter,
                    success_rate,
                } => {
                    // Normalize sender and receiver indices to valid ranges
                    let from_idx = (from_idx as usize) % peer_pks.len();
                    let to_idx = (to_idx as usize) % peer_pks.len();

                    // Create link with specified characteristics
                    // success_rate is normalized from u8 (0-255) to f64 (0.0-1.0)
                    let link = simulated::Link {
                        latency: Duration::from_millis(latency_ms as u64),
                        jitter: Duration::from_millis(jitter as u64),
                        success_rate: (success_rate as f64) / 255.0,
                    };
                    let _ = oracle
                        .add_link(peer_pks[from_idx].clone(), peer_pks[to_idx].clone(), link)
                        .await;
                }

                Operation::RemoveLink { from_idx, to_idx } => {
                    // Normalize sender and receiver indices to valid ranges
                    let from_idx = (from_idx as usize) % peer_pks.len();
                    let to_idx = (to_idx as usize) % peer_pks.len();

                    // Remove link to simulate network partition
                    let _ = oracle
                        .remove_link(peer_pks[from_idx].clone(), peer_pks[to_idx].clone())
                        .await;
                }
            }
        }
    });
}

fuzz_target!(|input: FuzzInput| {
    fuzz(input);
});