use std::net::Ipv4Addr;
use tokio::net::UdpSocket;
use tokio::time::{timeout, Duration};

const STUN_SERVER: &str = "stun.cloudflare.com:3478";
const STUN_TIMEOUT: Duration = Duration::from_secs(3);

// STUN constants
const STUN_BINDING_REQUEST: u16 = 0x0001;
const STUN_MAGIC_COOKIE: u32 = 0x2112A442;
const ATTR_XOR_MAPPED_ADDRESS: u16 = 0x0020;
const ATTR_MAPPED_ADDRESS: u16 = 0x0001;

/// Discover our public IP via STUN Binding Request (RFC 5389).
/// Returns `None` on timeout or parse failure.
pub async fn discover_public_ip() -> Option<String> {
    discover_public_ip_from(STUN_SERVER).await
}

pub async fn discover_public_ip_from(server: &str) -> Option<String> {
    let result = timeout(STUN_TIMEOUT, async {
        let socket = UdpSocket::bind("0.0.0.0:0").await.ok()?;
        socket.connect(server).await.ok()?;

        // Build STUN Binding Request (20 bytes)
        let mut request = [0u8; 20];
        // Message Type: Binding Request (0x0001)
        request[0..2].copy_from_slice(&STUN_BINDING_REQUEST.to_be_bytes());
        // Message Length: 0 (no attributes)
        request[2..4].copy_from_slice(&0u16.to_be_bytes());
        // Magic Cookie
        request[4..8].copy_from_slice(&STUN_MAGIC_COOKIE.to_be_bytes());
        // Transaction ID: 12 random bytes
        let txn_id: [u8; 12] = rand_bytes();
        request[8..20].copy_from_slice(&txn_id);

        socket.send(&request).await.ok()?;

        let mut buf = [0u8; 512];
        let n = socket.recv(&mut buf).await.ok()?;
        if n < 20 {
            return None;
        }

        parse_stun_response(&buf[..n], &txn_id)
    })
    .await;

    match result {
        Ok(ip) => ip,
        Err(_) => None, // timeout
    }
}

fn parse_stun_response(data: &[u8], _txn_id: &[u8; 12]) -> Option<String> {
    if data.len() < 20 {
        return None;
    }

    // Verify it's a Binding Response (0x0101)
    let msg_type = u16::from_be_bytes([data[0], data[1]]);
    if msg_type != 0x0101 {
        return None;
    }

    let msg_len = u16::from_be_bytes([data[2], data[3]]) as usize;
    let magic = u32::from_be_bytes([data[4], data[5], data[6], data[7]]);

    // Parse attributes
    let attrs = &data[20..std::cmp::min(20 + msg_len, data.len())];
    let mut offset = 0;

    while offset + 4 <= attrs.len() {
        let attr_type = u16::from_be_bytes([attrs[offset], attrs[offset + 1]]);
        let attr_len = u16::from_be_bytes([attrs[offset + 2], attrs[offset + 3]]) as usize;
        offset += 4;

        if offset + attr_len > attrs.len() {
            break;
        }

        let attr_data = &attrs[offset..offset + attr_len];

        match attr_type {
            ATTR_XOR_MAPPED_ADDRESS if attr_data.len() >= 8 => {
                let family = attr_data[1];
                if family == 0x01 {
                    // IPv4
                    let port_xored = u16::from_be_bytes([attr_data[2], attr_data[3]]);
                    let _port = port_xored ^ (STUN_MAGIC_COOKIE >> 16) as u16;
                    let ip_xored = u32::from_be_bytes([
                        attr_data[4],
                        attr_data[5],
                        attr_data[6],
                        attr_data[7],
                    ]);
                    let ip = ip_xored ^ magic;
                    return Some(Ipv4Addr::from(ip).to_string());
                }
            }
            ATTR_MAPPED_ADDRESS if attr_data.len() >= 8 => {
                let family = attr_data[1];
                if family == 0x01 {
                    // IPv4 (non-XOR fallback)
                    let ip = u32::from_be_bytes([
                        attr_data[4],
                        attr_data[5],
                        attr_data[6],
                        attr_data[7],
                    ]);
                    return Some(Ipv4Addr::from(ip).to_string());
                }
            }
            _ => {}
        }

        // Pad to 4-byte boundary
        offset += (attr_len + 3) & !3;
    }

    None
}

/// Generate 12 random bytes for transaction ID.
fn rand_bytes() -> [u8; 12] {
    let mut bytes = [0u8; 12];
    // Use a simple approach: mix timestamp + counter
    let now = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap_or_default();
    let nanos = now.as_nanos();
    bytes[0..8].copy_from_slice(&(nanos as u64).to_le_bytes());
    // Fill remaining with process-id based data
    let pid = std::process::id();
    bytes[8..12].copy_from_slice(&pid.to_le_bytes());
    bytes
}