use std::net::SocketAddr;
use tokio::net::UdpSocket;
use byteorder::{BigEndian, WriteBytesExt};
use rand::RngCore;
use std::sync::Arc;
use tokio::time::{sleep, Duration};

async fn get_public_addr(socket: &UdpSocket, stun_server: &str) -> Result<SocketAddr, Box<dyn std::error::Error>> {
    let mut request = Vec::with_capacity(20);
    request.write_u16::<BigEndian>(0x0001)?;
    request.write_u16::<BigEndian>(0)?;
    request.write_u32::<BigEndian>(0x2112A442)?;
    let mut transaction_id = [0u8; 12];
    rand::thread_rng().fill_bytes(&mut transaction_id);
    request.extend_from_slice(&transaction_id);

    socket.send_to(&request, stun_server).await?;
    let mut buf = [0u8; 1024];
    let (len, _) = socket.recv_from(&mut buf).await?;
    
    if len < 20 { return Err("Response too short".into()); }
    if u16::from_be_bytes([buf[0], buf[1]]) != 0x0101 { return Err("Not a binding response".into()); }

    let mut pos = 20;
    while pos + 4 <= len {
        let attr_type = u16::from_be_bytes([buf[pos], buf[pos+1]]);
        let attr_len = u16::from_be_bytes([buf[pos+2], buf[pos+3]]) as usize;
        pos += 4;
        if attr_type == 0x0020 {
            let xored_port = u16::from_be_bytes([buf[pos], buf[pos+1]]);
            let xored_ip = [buf[pos+2], buf[pos+3], buf[pos+4], buf[pos+5]];
            let port = xored_port ^ 0x2112;
            let ip = [xored_ip[0] ^ 0x21, xored_ip[1] ^ 0x12, xored_ip[2] ^ 0xA4, xored_ip[3] ^ 0x42];
            return Ok(SocketAddr::new(std::net::Ipv4Addr::new(ip[0], ip[1], ip[2], ip[3]), port));
        }
        pos += attr_len;
    }
    Err("XOR-MAPPED-ADDRESS not found".into())
}

async fn detect_nat_type(socket: &UdpSocket) -> Result<String, Box<dyn std::error::Error>> {
    // To detect Symmetric NAT, we query two different STUN servers.
    // If they report different public ports for the same internal socket, it's Symmetric.
    let server1 = "stun.l.google.com:19302";
    let server2 = "stun1.l.google.com:19302"; // Using another Google STUN server if available, or similar

    let addr1 = get_public_addr(socket, server1).await?;
    let addr2 = get_public_addr(socket, server2).await?;

    if addr1.port() == addr2.port() {
        Ok("Cone NAT".to_string())
    } else {
        Ok("Symmetric NAT".to_string())
    }
}

async fn punch_hole(socket: Arc<UdpSocket>, peer_addr: SocketAddr) {
    println!("Starting hole punching to {}", peer_addr);
    for i in 0..15 {
        let msg = format!("Punch {} from {}", i, socket.local_addr().unwrap());
        let _ = socket.send_to(msg.as_bytes(), peer_addr).await;
        sleep(Duration::from_millis(300)).await;
    }
}

async fn listen_for_peer(socket: Arc<UdpSocket>) {
    let mut buf = [0u8; 1024];
    loop {
        if let Ok((len, addr)) = socket.recv_from(&mut buf).await {
            let msg = String::from_utf8_lossy(&buf[..len]);
            println!("Received from {}: {}", addr, msg);
            if msg.contains("Punch") {
                println!("Hole punched successfully from {}!", addr);
            }
        }
    }
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let socket = Arc::new(UdpSocket::bind("0.0.0.0:0").await?);
    println!("Local address: {}", socket.local_addr()?);

    let nat_type = detect_nat_type(&socket).await.unwrap_or_else(|_| "Unknown".to_string());
    println!("Detected NAT Type: {}", nat_type);

    let public_addr = get_public_addr(&socket, "stun.l.google.com:19302").await?;
    println!("Public address: {}", public_addr);

    println!("Enter peer's public address (IP:PORT) or 'skip':");
    let mut input = String::new();
    std::io::stdin().read_line(&mut input)?;
    let input = input.trim();

    if input != "skip" {
        let peer_addr: SocketAddr = input.parse().expect("Invalid address");
        let s_clone = Arc::clone(&socket);
        tokio::spawn(async move {
            punch_hole(s_clone, peer_addr).await;
        });
    }

    listen_for_peer(socket).await;
    Ok(())
}