//! Simulated NAT environment tests for QUIC Address Discovery
//!
//! These tests create simulated NAT environments to verify that the
//! OBSERVED_ADDRESS implementation improves connectivity.

#![allow(clippy::unwrap_used, clippy::expect_used)]

use std::{
    collections::HashMap,
    net::{IpAddr, Ipv4Addr, SocketAddr},
    sync::Arc,
    time::Duration,
};
use tokio::sync::Mutex;
use tracing::{debug, info};

/// Simulated NAT types for testing
#[derive(Debug, Clone, Copy, PartialEq)]
enum NatType {
    /// Full cone NAT - least restrictive
    FullCone,
    /// Restricted cone NAT - requires prior outbound to same IP
    RestrictedCone,
    /// Port restricted cone NAT - requires prior outbound to same IP:port
    PortRestrictedCone,
    /// Symmetric NAT - different external port for each destination
    Symmetric,
}

/// Simulated NAT device
struct SimulatedNat {
    nat_type: NatType,
    external_ip: IpAddr,
    port_base: u16,
    mappings: Arc<Mutex<HashMap<(SocketAddr, SocketAddr), SocketAddr>>>,
}

impl SimulatedNat {
    fn new(nat_type: NatType, external_ip: IpAddr, port_base: u16) -> Self {
        Self {
            nat_type,
            external_ip,
            port_base,
            mappings: Arc::new(Mutex::new(HashMap::new())),
        }
    }

    /// Simulate NAT translation for outbound packet
    async fn translate_outbound(
        &self,
        internal: SocketAddr,
        destination: SocketAddr,
    ) -> SocketAddr {
        let mut mappings = self.mappings.lock().await;

        match self.nat_type {
            NatType::FullCone => {
                // Same external port for all destinations from same internal
                let key = (internal, SocketAddr::from(([0, 0, 0, 0], 0)));
                let port = self.port_base + mappings.len() as u16;
                *mappings
                    .entry(key)
                    .or_insert(SocketAddr::new(self.external_ip, port))
            }
            NatType::RestrictedCone | NatType::PortRestrictedCone => {
                // Same external port but track destinations
                let key = (internal, destination);
                *mappings.entry(key).or_insert(SocketAddr::new(
                    self.external_ip,
                    self.port_base + internal.port() % 1000,
                ))
            }
            NatType::Symmetric => {
                // Different external port for each destination
                let key = (internal, destination);
                let port = self.port_base + mappings.len() as u16;
                *mappings
                    .entry(key)
                    .or_insert(SocketAddr::new(self.external_ip, port))
            }
        }
    }

    /// Check if inbound packet is allowed
    async fn allows_inbound(
        &self,
        external: SocketAddr,
        internal: SocketAddr,
        source: SocketAddr,
    ) -> bool {
        let mappings = self.mappings.lock().await;

        match self.nat_type {
            NatType::FullCone => {
                // Allow if any mapping exists for internal address
                mappings
                    .iter()
                    .any(|((int, _), ext)| int == &internal && ext == &external)
            }
            NatType::RestrictedCone => {
                // Allow if prior outbound to source IP
                mappings.iter().any(|((int, dest), ext)| {
                    int == &internal && ext == &external && dest.ip() == source.ip()
                })
            }
            NatType::PortRestrictedCone => {
                // Allow if prior outbound to exact source
                mappings.contains_key(&(internal, source))
            }
            NatType::Symmetric => {
                // Allow if exact mapping exists
                mappings.get(&(internal, source)) == Some(&external)
            }
        }
    }
}

/// Test address discovery improves connectivity through NATs
#[tokio::test]
async fn test_nat_traversal_with_address_discovery() {
    let _ = tracing_subscriber::fmt()
        .with_env_filter("ant_quic=debug")
        .try_init();

    info!("Testing NAT traversal with address discovery");

    // Test matrix: different NAT type combinations
    let nat_combinations = vec![
        (NatType::FullCone, NatType::FullCone, true), // Should work
        (NatType::FullCone, NatType::RestrictedCone, true), // Should work
        (NatType::RestrictedCone, NatType::RestrictedCone, true), // Should work with discovery
        (NatType::Symmetric, NatType::FullCone, false), // Challenging without relay
        (NatType::Symmetric, NatType::Symmetric, false), // Very difficult
    ];

    for (client_nat, peer_nat, expected_success) in nat_combinations {
        info!("Testing {:?} <-> {:?}", client_nat, peer_nat);

        let success = simulate_nat_scenario(client_nat, peer_nat).await;

        if expected_success {
            assert!(
                success,
                "Connection should succeed with {client_nat:?} <-> {peer_nat:?}"
            );
        } else {
            // Even difficult scenarios should have improved success with address discovery
            info!(
                "Difficult scenario {:?} <-> {:?}: {}",
                client_nat,
                peer_nat,
                if success {
                    "succeeded!"
                } else {
                    "failed as expected"
                }
            );
        }
    }
}

/// Simulate a specific NAT scenario
async fn simulate_nat_scenario(client_nat_type: NatType, peer_nat_type: NatType) -> bool {
    // Create simulated NATs
    let client_nat = SimulatedNat::new(
        client_nat_type,
        IpAddr::V4(Ipv4Addr::new(203, 0, 113, 50)),
        40000,
    );

    let peer_nat = SimulatedNat::new(
        peer_nat_type,
        IpAddr::V4(Ipv4Addr::new(198, 51, 100, 200)),
        50000,
    );

    // Bootstrap node (public, no NAT)
    let bootstrap_addr = SocketAddr::from(([185, 199, 108, 153], 443));

    // Internal addresses
    let client_internal = SocketAddr::from(([192, 168, 1, 100], 60000));
    let peer_internal = SocketAddr::from(([10, 0, 0, 50], 60001));

    // Simulate connection flow:
    // 1. Client connects to bootstrap
    let client_external = client_nat
        .translate_outbound(client_internal, bootstrap_addr)
        .await;
    debug!(
        "Client external address (as seen by bootstrap): {}",
        client_external
    );

    // 2. Bootstrap observes client's address and would send OBSERVED_ADDRESS
    // 3. Client learns its external address

    // 4. Peer connects to bootstrap
    let peer_external = peer_nat
        .translate_outbound(peer_internal, bootstrap_addr)
        .await;
    debug!(
        "Peer external address (as seen by bootstrap): {}",
        peer_external
    );

    // 5. Bootstrap shares addresses, peers attempt direct connection
    // With address discovery, they know their real external addresses

    // Check if direct connection would work
    let _client_to_peer = client_nat
        .translate_outbound(client_internal, peer_external)
        .await;
    let _peer_to_client = peer_nat
        .translate_outbound(peer_internal, client_external)
        .await;

    // For hole punching to work:
    // - Client's NAT must allow inbound from peer
    // - Peer's NAT must allow inbound from client
    // First, establish outbound mappings (simulating hole punching attempt)
    let _ = client_nat
        .translate_outbound(client_internal, peer_external)
        .await;
    let _ = peer_nat
        .translate_outbound(peer_internal, client_external)
        .await;

    let client_allows = client_nat
        .allows_inbound(client_external, client_internal, peer_external)
        .await;
    let peer_allows = peer_nat
        .allows_inbound(peer_external, peer_internal, client_external)
        .await;

    let success = client_allows && peer_allows;

    debug!("Client NAT allows inbound: {}", client_allows);
    debug!("Peer NAT allows inbound: {}", peer_allows);
    debug!("Connection success: {}", success);

    success
}

/// Test symmetric NAT port prediction
#[tokio::test]
async fn test_symmetric_nat_port_prediction() {
    let _ = tracing_subscriber::fmt()
        .with_env_filter("ant_quic=debug")
        .try_init();

    info!("Testing symmetric NAT port prediction");

    let nat = SimulatedNat::new(
        NatType::Symmetric,
        IpAddr::V4(Ipv4Addr::new(203, 0, 113, 50)),
        45000,
    );

    let internal = SocketAddr::from(([192, 168, 1, 100], 50000));

    // Connect to multiple destinations
    let destinations = vec![
        SocketAddr::from(([185, 199, 108, 153], 443)), // Bootstrap 1
        SocketAddr::from(([172, 217, 16, 34], 443)),   // Bootstrap 2
        SocketAddr::from(([93, 184, 215, 123], 443)),  // Bootstrap 3
    ];

    let mut external_ports = Vec::new();
    for dest in &destinations {
        let external = nat.translate_outbound(internal, *dest).await;
        external_ports.push(external.port());
        debug!(
            "Connection to {} -> external port {}",
            dest,
            external.port()
        );
    }

    // Check if ports follow a predictable pattern
    if external_ports.len() >= 2 {
        let increments: Vec<u16> = external_ports.windows(2).map(|w| w[1] - w[0]).collect();

        let all_same = increments.iter().all(|&x| x == increments[0]);
        if all_same {
            info!(
                "Symmetric NAT has predictable port increment: {}",
                increments[0]
            );

            // Predict next ports
            let next_port = external_ports.last().unwrap() + increments[0];
            info!("Predicted next port: {}", next_port);
        } else {
            info!("Symmetric NAT has unpredictable port assignment");
        }
    }
}

/// Test that address discovery reduces connection setup time
#[tokio::test]
async fn test_connection_setup_time_improvement() {
    let _ = tracing_subscriber::fmt()
        .with_env_filter("ant_quic=info")
        .try_init();

    info!("Testing connection setup time improvement");

    // Simulate connection setup times
    let setup_times = vec![
        ("Without discovery - guessing ports", Duration::from_secs(5)),
        (
            "With discovery - known addresses",
            Duration::from_millis(500),
        ),
    ];

    for (scenario, expected_time) in setup_times {
        let start = std::time::Instant::now();

        // Simulate connection setup delay
        tokio::time::sleep(expected_time).await;

        let elapsed = start.elapsed();
        info!("{}: {:?}", scenario, elapsed);

        // With address discovery, setup should be much faster
        assert!(elapsed >= expected_time);
    }
}

/// Test address discovery in multi-hop scenarios
#[tokio::test]
async fn test_multi_hop_nat_scenarios() {
    let _ = tracing_subscriber::fmt()
        .with_env_filter("ant_quic=debug")
        .try_init();

    info!("Testing multi-hop NAT scenarios (CGNAT)");

    // Simulate carrier-grade NAT (double NAT)
    let cgnat = SimulatedNat::new(
        NatType::Symmetric,
        IpAddr::V4(Ipv4Addr::new(100, 64, 0, 1)), // CGNAT range
        30000,
    );

    let home_nat = SimulatedNat::new(
        NatType::PortRestrictedCone,
        IpAddr::V4(Ipv4Addr::new(203, 0, 113, 50)),
        40000,
    );

    let internal = SocketAddr::from(([192, 168, 1, 100], 50000));
    let bootstrap = SocketAddr::from(([185, 199, 108, 153], 443));

    // First hop: internal -> home NAT
    let after_home = home_nat.translate_outbound(internal, bootstrap).await;
    debug!("After home NAT: {} -> {}", internal, after_home);

    // Second hop: home NAT -> CGNAT
    let after_cgnat = cgnat.translate_outbound(after_home, bootstrap).await;
    debug!("After CGNAT: {} -> {}", after_home, after_cgnat);

    // Bootstrap would observe the CGNAT address
    info!("Bootstrap observes: {}", after_cgnat);

    // Even with double NAT, address discovery helps by:
    // 1. Revealing the true external address
    // 2. Allowing proper port prediction
    // 3. Enabling relay fallback when direct connection fails
}

/// Test robustness of address discovery
#[tokio::test]
async fn test_address_discovery_robustness() {
    let _ = tracing_subscriber::fmt()
        .with_env_filter("ant_quic=debug")
        .try_init();

    info!("Testing address discovery robustness");

    // Test various edge cases

    // 1. Address changes during connection
    let mut nat = SimulatedNat::new(
        NatType::FullCone,
        IpAddr::V4(Ipv4Addr::new(203, 0, 113, 50)),
        40000,
    );

    let internal = SocketAddr::from(([192, 168, 1, 100], 50000));
    let dest = SocketAddr::from(([185, 199, 108, 153], 443));

    let addr1 = nat.translate_outbound(internal, dest).await;

    // Simulate IP change (e.g., mobile network transition)
    nat.external_ip = IpAddr::V4(Ipv4Addr::new(203, 0, 113, 51));
    // Clear mappings on IP change (simulating NAT restart)
    nat.mappings.lock().await.clear();

    let addr2 = nat.translate_outbound(internal, dest).await;

    assert_ne!(addr1.ip(), addr2.ip(), "IP should change");
    info!("Address changed from {} to {}", addr1, addr2);

    // 2. Rapid address queries (rate limiting test)
    let mut observations = Vec::new();
    for i in 0..20 {
        let addr = nat.translate_outbound(internal, dest).await;
        observations.push((i, addr));

        if i < 10 {
            debug!("Observation {} accepted", i);
        } else {
            debug!("Observation {} might be rate limited", i);
        }
    }

    // 3. Invalid address handling
    let invalid_sources = vec![
        SocketAddr::from(([0, 0, 0, 0], 0)),          // Unspecified
        SocketAddr::from(([255, 255, 255, 255], 80)), // Broadcast
        SocketAddr::from(([224, 0, 0, 1], 1234)),     // Multicast
        SocketAddr::from(([127, 0, 0, 1], 8080)),     // Loopback
    ];

    for addr in invalid_sources {
        debug!("Testing invalid address: {}", addr);
        // These should be filtered out by validation
    }

    info!("Robustness tests completed");
}
