//! NAT traversal via UDP hole-punching (R-03).
//!
//! Implements cooperative UDP hole-punching: both peers simultaneously
//! send probes to each other's public address to open NAT mappings.
//! This allows direct peer-to-peer connectivity before falling back
//! to the relay for fully symmetric NAT scenarios.
//!
//! Flow:
//!  1. Both peers connect to the relay and report their observed public addr
//!  2. Relay exchanges public addresses via PLP CONNECT payloads
//!  3. Both peers start hole-punch probes simultaneously (via countdown)
//!  4. On first successful probe reply, direct path is established
//!  5. If hole-punch fails within timeout, relay forwarding continues

use std::net::SocketAddr;
use std::sync::Arc;
use tokio::net::UdpSocket;
use tokio::time::{Duration, sleep, timeout};
use tracing::{debug, info, warn};

/// Result of a hole-punch attempt.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum HolePunchResult {
    /// Direct connection established to peer_addr
    Success { peer_addr: SocketAddr, rtt_ms: u64 },
    /// Hole punching failed — will use relay
    Failed { reason: String },
    /// Symmetric NAT detected — relay required
    SymmetricNat,
}

/// Configuration for hole-punching.
#[derive(Debug, Clone)]
pub struct HolePunchConfig {
    /// Max time to attempt hole-punching before giving up
    pub timeout: Duration,
    /// Interval between probe packets
    pub probe_interval: Duration,
    /// Maximum number of probe attempts
    pub max_probes: u32,
}

impl Default for HolePunchConfig {
    fn default() -> Self {
        Self {
            timeout: Duration::from_secs(5),
            probe_interval: Duration::from_millis(100),
            max_probes: 40,
        }
    }
}

/// Probe payload — 8 bytes: magic (4) + sequence (4)
const PROBE_MAGIC: &[u8; 4] = b"PLSP"; // Pulse Signalling Probe 
const PROBE_ACK_MAGIC: &[u8; 4] = b"PSPA"; // Pulse Signalling Probe Ack

/// Attempt UDP hole-punching to a peer's public address.
/// `local_socket` is already bound; `peer_addr` is the peer's observed public address
/// as reported by the relay.
pub async fn attempt_hole_punch(
    local_socket: Arc<UdpSocket>,
    peer_addr: SocketAddr,
    config: HolePunchConfig,
) -> HolePunchResult {
    info!("[NAT] Starting hole-punch to {}", peer_addr);

    let socket_recv = local_socket.clone();
    let start_ms = std::time::Instant::now();

    // Spawn the receiver task
    let recv_handle = tokio::spawn(async move {
        let mut buf = [0u8; 32];
        loop {
            match timeout(Duration::from_millis(100), socket_recv.recv_from(&mut buf)).await {
                Ok(Ok((len, addr))) if len >= 4 => {
                    if &buf[..4] == PROBE_MAGIC || &buf[..4] == PROBE_ACK_MAGIC {
                        return Some((addr, start_ms.elapsed().as_millis() as u64));
                    }
                }
                _ => continue,
            }
        }
    });

    // Send probes simultaneously
    let mut probes_sent = 0u32;
    let mut probe_seq = 0u32;
    let send_socket = local_socket.clone();

    loop {
        if probes_sent >= config.max_probes {
            recv_handle.abort();
            return HolePunchResult::Failed {
                reason: format!("Max probes ({}) exhausted", config.max_probes),
            };
        }

        // Build probe: magic + seq
        let mut probe = [0u8; 8];
        probe[..4].copy_from_slice(PROBE_MAGIC);
        probe[4..8].copy_from_slice(&probe_seq.to_be_bytes());

        if let Err(e) = send_socket.send_to(&probe, peer_addr).await {
            warn!("[NAT] Probe send failed: {}", e);
        } else {
            debug!("[NAT] Probe {} sent to {}", probe_seq, peer_addr);
        }

        probes_sent += 1;
        probe_seq += 1;

        // Check if receiver found a reply
        if recv_handle.is_finished() {
            break;
        }

        // Check overall timeout
        if start_ms.elapsed() >= config.timeout {
            recv_handle.abort();
            return HolePunchResult::Failed {
                reason: "Timeout".to_string(),
            };
        }

        sleep(config.probe_interval).await;
    }

    match timeout(Duration::from_millis(200), recv_handle).await {
        Ok(Ok(Some((addr, rtt_ms)))) => {
            info!("[NAT] Hole-punch SUCCESS to {} (RTT {}ms)", addr, rtt_ms);
            HolePunchResult::Success {
                peer_addr: addr,
                rtt_ms,
            }
        }
        _ => HolePunchResult::Failed {
            reason: "No response received".to_string(),
        },
    }
}

/// Detect whether a NAT mapping is symmetric by comparing the observed
/// public address across two different remote endpoints.
pub fn detect_symmetric_nat(
    observed_from_relay_a: SocketAddr,
    observed_from_relay_b: SocketAddr,
) -> bool {
    // If port differs between two external vantage points, it's symmetric NAT
    observed_from_relay_a.port() != observed_from_relay_b.port()
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::Arc;
    use tokio::net::UdpSocket;
    use tokio::time::Duration;

    #[test]
    fn symmetric_nat_detects_port_change() {
        let a: SocketAddr = "203.0.113.10:11111".parse().unwrap();
        let b: SocketAddr = "203.0.113.10:22222".parse().unwrap();
        assert!(detect_symmetric_nat(a, b));
    }

    #[test]
    fn non_symmetric_nat_same_port() {
        let a: SocketAddr = "203.0.113.10:11111".parse().unwrap();
        let b: SocketAddr = "198.51.100.20:11111".parse().unwrap();
        assert!(!detect_symmetric_nat(a, b));
    }

    #[tokio::test]
    async fn hole_punch_succeeds_with_local_responder() {
        let peer_socket = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap());
        let peer_addr = peer_socket.local_addr().unwrap();

        let responder = peer_socket.clone();
        tokio::spawn(async move {
            let mut buf = [0u8; 64];
            loop {
                let (len, src) = match responder.recv_from(&mut buf).await {
                    Ok(v) => v,
                    Err(_) => return,
                };
                if len >= 8 && &buf[..4] == PROBE_MAGIC {
                    let mut ack = [0u8; 8];
                    ack[..4].copy_from_slice(PROBE_ACK_MAGIC);
                    ack[4..8].copy_from_slice(&buf[4..8]);
                    let _ = responder.send_to(&ack, src).await;
                    return;
                }
            }
        });

        let local_socket = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap());
        let config = HolePunchConfig {
            timeout: Duration::from_secs(2),
            probe_interval: Duration::from_millis(10),
            max_probes: 100,
        };

        let result = attempt_hole_punch(local_socket, peer_addr, config).await;
        match result {
            HolePunchResult::Success {
                peer_addr: got,
                rtt_ms,
            } => {
                assert_eq!(got, peer_addr);
                assert!(rtt_ms <= 2_000);
            }
            other => panic!("expected Success, got: {:?}", other),
        }
    }

    #[tokio::test]
    async fn hole_punch_fails_when_max_probes_exhausted() {
        let local_socket = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap());
        let unreachable_peer: SocketAddr = "127.0.0.1:9".parse().unwrap();

        let config = HolePunchConfig {
            timeout: Duration::from_secs(5),
            probe_interval: Duration::from_millis(1),
            max_probes: 1,
        };

        let result = attempt_hole_punch(local_socket, unreachable_peer, config).await;
        match result {
            HolePunchResult::Failed { reason } => {
                assert!(reason.contains("Max probes"));
            }
            other => panic!("expected Failed, got: {:?}", other),
        }
    }

    #[tokio::test]
    async fn hole_punch_times_out_without_responder() {
        let local_socket = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap());
        let blackhole_peer: SocketAddr = "127.0.0.1:9".parse().unwrap();

        let config = HolePunchConfig {
            timeout: Duration::from_millis(80),
            probe_interval: Duration::from_millis(10),
            max_probes: 100,
        };

        let result = attempt_hole_punch(local_socket, blackhole_peer, config).await;
        match result {
            HolePunchResult::Failed { reason } => {
                assert!(reason == "Timeout" || reason.contains("Timeout"));
            }
            other => panic!("expected Failed, got: {:?}", other),
        }
    }
}
