//! STUN Server Implementation
//! RFC 5389 compliant STUN server for NAT traversal

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

/// STUN Message Types
const BINDING_REQUEST: u16 = 0x0001;
const BINDING_RESPONSE: u16 = 0x0101;

/// STUN Attribute Types
const ATTR_MAPPED_ADDRESS: u16 = 0x0001;
const ATTR_XOR_MAPPED_ADDRESS: u16 = 0x0020;
const ATTR_SOFTWARE: u16 = 0x8022;
const ATTR_FINGERPRINT: u16 = 0x8028;

/// STUN Magic Cookie (RFC 5389)
const MAGIC_COOKIE: u32 = 0x2112A442;

/// STUN Server
pub struct StunServer {
    port: u16,
}

impl StunServer {
    pub fn new(port: u16) -> Self {
        Self { port }
    }

    /// Start the STUN server
    pub async fn start(&self) -> anyhow::Result<()> {
        let addr = format!("0.0.0.0:{}", self.port);
        let socket = UdpSocket::bind(&addr).await?;
        let socket = Arc::new(socket);
        
        info!("STUN Server listening on UDP port {}", self.port);

        loop {
            let mut buf = [0u8; 1024];
            match socket.recv_from(&mut buf).await {
                Ok((len, src_addr)) => {
                    let socket_clone = socket.clone();
                    let data = buf[..len].to_vec();
                    
                    tokio::spawn(async move {
                        if let Err(e) = handle_stun_request(&socket_clone, &data, src_addr).await {
                            warn!("Failed to handle STUN request: {}", e);
                        }
                    });
                }
                Err(e) => {
                    warn!("STUN recv error: {}", e);
                }
            }
        }
    }
}

/// Handle a STUN request
async fn handle_stun_request(
    socket: &UdpSocket,
    data: &[u8],
    src_addr: SocketAddr,
) -> anyhow::Result<()> {
    // Validate minimum STUN header size (20 bytes)
    if data.len() < 20 {
        return Ok(());
    }

    // Parse STUN header
    let msg_type = u16::from_be_bytes([data[0], data[1]]);
    let msg_len = u16::from_be_bytes([data[2], data[3]]);
    let magic = u32::from_be_bytes([data[4], data[5], data[6], data[7]]);
    let transaction_id = &data[8..20];

    // Verify magic cookie
    if magic != MAGIC_COOKIE {
        debug!("Invalid STUN magic cookie from {}", src_addr);
        return Ok(());
    }

    // Only handle Binding Request
    if msg_type != BINDING_REQUEST {
        debug!("Ignoring non-binding STUN message type: 0x{:04x}", msg_type);
        return Ok(());
    }

    debug!("STUN Binding Request from {}", src_addr);

    // Build Binding Response
    let response = build_binding_response(transaction_id, src_addr)?;

    // Send response
    socket.send_to(&response, src_addr).await?;
    debug!("STUN Binding Response sent to {}", src_addr);

    Ok(())
}

/// Build a STUN Binding Response
fn build_binding_response(transaction_id: &[u8], mapped_addr: SocketAddr) -> anyhow::Result<Vec<u8>> {
    let mut response = Vec::with_capacity(64);

    // Build attributes first to calculate length
    let xor_mapped_attr = build_xor_mapped_address(mapped_addr, transaction_id);
    let software_attr = build_software_attribute();
    
    let attrs_len = xor_mapped_attr.len() + software_attr.len();

    // STUN Header (20 bytes)
    // Message Type: Binding Response
    response.extend_from_slice(&BINDING_RESPONSE.to_be_bytes());
    // Message Length (excluding header)
    response.extend_from_slice(&(attrs_len as u16).to_be_bytes());
    // Magic Cookie
    response.extend_from_slice(&MAGIC_COOKIE.to_be_bytes());
    // Transaction ID (12 bytes)
    response.extend_from_slice(transaction_id);

    // Attributes
    response.extend_from_slice(&xor_mapped_attr);
    response.extend_from_slice(&software_attr);

    Ok(response)
}

/// Build XOR-MAPPED-ADDRESS attribute (RFC 5389)
fn build_xor_mapped_address(addr: SocketAddr, transaction_id: &[u8]) -> Vec<u8> {
    let mut attr = Vec::with_capacity(12);

    match addr {
        SocketAddr::V4(v4) => {
            // Attribute Type
            attr.extend_from_slice(&ATTR_XOR_MAPPED_ADDRESS.to_be_bytes());
            // Attribute Length (8 bytes for IPv4)
            attr.extend_from_slice(&8u16.to_be_bytes());
            // Reserved (1 byte) + Family (1 byte)
            attr.push(0x00);
            attr.push(0x01); // IPv4
            
            // XOR'd Port
            let port = v4.port() ^ ((MAGIC_COOKIE >> 16) as u16);
            attr.extend_from_slice(&port.to_be_bytes());
            
            // XOR'd Address
            let ip_bytes = v4.ip().octets();
            let magic_bytes = MAGIC_COOKIE.to_be_bytes();
            for i in 0..4 {
                attr.push(ip_bytes[i] ^ magic_bytes[i]);
            }
        }
        SocketAddr::V6(v6) => {
            // Attribute Type
            attr.extend_from_slice(&ATTR_XOR_MAPPED_ADDRESS.to_be_bytes());
            // Attribute Length (20 bytes for IPv6)
            attr.extend_from_slice(&20u16.to_be_bytes());
            // Reserved (1 byte) + Family (1 byte)
            attr.push(0x00);
            attr.push(0x02); // IPv6
            
            // XOR'd Port
            let port = v6.port() ^ ((MAGIC_COOKIE >> 16) as u16);
            attr.extend_from_slice(&port.to_be_bytes());
            
            // XOR'd Address (XOR with magic cookie + transaction ID)
            let ip_bytes = v6.ip().octets();
            let magic_bytes = MAGIC_COOKIE.to_be_bytes();
            for i in 0..4 {
                attr.push(ip_bytes[i] ^ magic_bytes[i]);
            }
            for i in 4..16 {
                attr.push(ip_bytes[i] ^ transaction_id[i - 4]);
            }
        }
    }

    attr
}

/// Build SOFTWARE attribute
fn build_software_attribute() -> Vec<u8> {
    let software = b"EasyRemote STUN Server";
    let mut attr = Vec::with_capacity(4 + software.len() + 4);

    // Attribute Type
    attr.extend_from_slice(&ATTR_SOFTWARE.to_be_bytes());
    
    // Attribute Length
    let len = software.len() as u16;
    attr.extend_from_slice(&len.to_be_bytes());
    
    // Value
    attr.extend_from_slice(software);
    
    // Padding to 4-byte boundary
    let padding = (4 - (software.len() % 4)) % 4;
    for _ in 0..padding {
        attr.push(0x00);
    }

    attr
}

/// Get the default local STUN server URL
pub fn get_local_stun_url(port: u16, public_ip: Option<&str>) -> String {
    let host = public_ip.unwrap_or("localhost");
    format!("stun:{}:{}", host, port)
}
