//! Minimal raw STUN client for NAT type detection.
//! Sends STUN Binding Requests to two servers and compares mapped ports.

use std::net::SocketAddr;
use tokio::net::UdpSocket;
use tracing::{debug, warn};

use crate::types::{NatMapping, NatType};

const STUN_SERVERS: &[&str] = &[
    "stun.l.google.com:19302",
    "stun.cloudflare.com:3478",
];

const STUN_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(3);

/// STUN Binding Request (RFC 5389): 20 bytes
/// Type: 0x0001 (Binding Request), Length: 0, Magic: 0x2112A442, Transaction ID: 12 random bytes
fn build_binding_request() -> [u8; 20] {
    let mut buf = [0u8; 20];
    // Message type: Binding Request (0x0001)
    buf[0] = 0x00;
    buf[1] = 0x01;
    // Message length: 0
    buf[2] = 0x00;
    buf[3] = 0x00;
    // Magic cookie: 0x2112A442
    buf[4] = 0x21;
    buf[5] = 0x12;
    buf[6] = 0xA4;
    buf[7] = 0x42;
    // Transaction ID: 12 random bytes
    use rand::Rng;
    let mut rng = rand::rng();
    rng.fill(&mut buf[8..20]);
    buf
}

/// Parse XOR-MAPPED-ADDRESS from a STUN Binding Response.
/// Returns the mapped SocketAddr or None if not found/parseable.
fn parse_xor_mapped_address(resp: &[u8], txn_id: &[u8; 12]) -> Option<SocketAddr> {
    if resp.len() < 20 {
        return None;
    }
    // Verify it's a Binding Response (0x0101)
    if resp[0] != 0x01 || resp[1] != 0x01 {
        return None;
    }
    let magic: [u8; 4] = [0x21, 0x12, 0xA4, 0x42];

    // Walk attributes
    let msg_len = u16::from_be_bytes([resp[2], resp[3]]) as usize;
    let end = std::cmp::min(20 + msg_len, resp.len());
    let mut pos = 20;
    while pos + 4 <= end {
        let attr_type = u16::from_be_bytes([resp[pos], resp[pos + 1]]);
        let attr_len = u16::from_be_bytes([resp[pos + 2], resp[pos + 3]]) as usize;
        pos += 4;
        if pos + attr_len > end {
            break;
        }
        // XOR-MAPPED-ADDRESS = 0x0020, MAPPED-ADDRESS = 0x0001
        if attr_type == 0x0020 && attr_len >= 8 {
            // byte 0: reserved, byte 1: family (0x01=IPv4, 0x02=IPv6)
            let family = resp[pos + 1];
            if family == 0x01 {
                // IPv4
                let xport = u16::from_be_bytes([resp[pos + 2], resp[pos + 3]])
                    ^ u16::from_be_bytes([magic[0], magic[1]]);
                let xip = [
                    resp[pos + 4] ^ magic[0],
                    resp[pos + 5] ^ magic[1],
                    resp[pos + 6] ^ magic[2],
                    resp[pos + 7] ^ magic[3],
                ];
                let addr = SocketAddr::new(
                    std::net::IpAddr::V4(std::net::Ipv4Addr::new(xip[0], xip[1], xip[2], xip[3])),
                    xport,
                );
                return Some(addr);
            } else if family == 0x02 && attr_len >= 20 {
                // IPv6: XOR with magic + txn_id
                let xport = u16::from_be_bytes([resp[pos + 2], resp[pos + 3]])
                    ^ u16::from_be_bytes([magic[0], magic[1]]);
                let mut ip6 = [0u8; 16];
                let xor_key: Vec<u8> = magic.iter().chain(txn_id.iter()).copied().collect();
                for i in 0..16 {
                    ip6[i] = resp[pos + 4 + i] ^ xor_key[i];
                }
                let addr = SocketAddr::new(
                    std::net::IpAddr::V6(std::net::Ipv6Addr::from(ip6)),
                    xport,
                );
                return Some(addr);
            }
        }
        // Pad to 4-byte boundary
        pos += (attr_len + 3) & !3;
    }
    None
}

/// Query a single STUN server and return the mapped address.
async fn stun_query(sock: &UdpSocket, server: &str) -> Option<SocketAddr> {
    use std::net::ToSocketAddrs;
    let server_addr = match server.to_socket_addrs() {
        Ok(mut addrs) => addrs.next()?,
        Err(e) => {
            debug!(server, error = %e, "STUN DNS resolution failed");
            return None;
        }
    };

    let request = build_binding_request();
    let txn_id: [u8; 12] = request[8..20].try_into().unwrap();

    if let Err(e) = sock.send_to(&request, server_addr).await {
        debug!(server, error = %e, "STUN send failed");
        return None;
    }

    let mut buf = [0u8; 256];
    match tokio::time::timeout(STUN_TIMEOUT, sock.recv_from(&mut buf)).await {
        Ok(Ok((len, _))) => parse_xor_mapped_address(&buf[..len], &txn_id),
        Ok(Err(e)) => {
            debug!(server, error = %e, "STUN recv failed");
            None
        }
        Err(_) => {
            debug!(server, "STUN query timed out (3s)");
            None
        }
    }
}

/// Detect NAT type by comparing mapped addresses from two STUN servers.
/// Must be called with the local port we're interested in (for Public detection).
/// Also returns the NatMapping classification for the advanced NAT profile.
pub async fn detect_nat_type(local_port: u16) -> (NatType, NatMapping) {
    let sock = match UdpSocket::bind("0.0.0.0:0").await {
        Ok(s) => s,
        Err(e) => {
            warn!(error = %e, "Failed to bind UDP socket for STUN");
            return (NatType::Unknown, NatMapping::Unknown);
        }
    };

    let local_addr = sock.local_addr().ok();

    // Query both servers from the same socket
    let result1 = stun_query(&sock, STUN_SERVERS[0]).await;
    let result2 = stun_query(&sock, STUN_SERVERS[1]).await;

    match (result1, result2) {
        (Some(addr1), Some(addr2)) => {
            debug!(
                server1 = STUN_SERVERS[0], mapped1 = %addr1,
                server2 = STUN_SERVERS[1], mapped2 = %addr2,
                local_port,
                "STUN results"
            );
            // If mapped port matches our local port, we might be public/no-NAT
            if let Some(local) = local_addr {
                if addr1.port() == local.port() && addr2.port() == local.port() {
                    return (NatType::Public, NatMapping::EndpointIndependent);
                }
            }
            // Same mapped port from both = cone NAT (Easy / EIM)
            // Different ports = symmetric NAT (Hard / EDM)
            if addr1.port() == addr2.port() {
                (NatType::Easy, NatMapping::EndpointIndependent)
            } else {
                (NatType::Hard, NatMapping::EndpointDependent)
            }
        }
        (Some(addr), None) | (None, Some(addr)) => {
            debug!(mapped = %addr, "Only one STUN server responded, assuming Easy");
            (NatType::Easy, NatMapping::EndpointIndependent)
        }
        (None, None) => {
            warn!("Both STUN servers unreachable, NAT type unknown");
            (NatType::Unknown, NatMapping::Unknown)
        }
    }
}