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
}

#[cfg(test)]
mod tests {
    use super::*;

    /// Build a synthetic STUN Binding Response with given attributes.
    fn build_stun_response(attrs: &[(u16, &[u8])]) -> Vec<u8> {
        let mut attrs_bytes = Vec::new();
        for &(attr_type, attr_data) in attrs {
            attrs_bytes.extend_from_slice(&attr_type.to_be_bytes());
            attrs_bytes.extend_from_slice(&(attr_data.len() as u16).to_be_bytes());
            attrs_bytes.extend_from_slice(attr_data);
            // Pad to 4-byte boundary
            let pad = (4 - (attr_data.len() % 4)) % 4;
            attrs_bytes.extend(std::iter::repeat(0u8).take(pad));
        }

        let mut response = Vec::new();
        // msg_type = 0x0101 (Binding Response)
        response.extend_from_slice(&0x0101u16.to_be_bytes());
        // message length
        response.extend_from_slice(&(attrs_bytes.len() as u16).to_be_bytes());
        // magic cookie
        response.extend_from_slice(&STUN_MAGIC_COOKIE.to_be_bytes());
        // transaction ID (12 bytes)
        response.extend_from_slice(&[0u8; 12]);
        // attributes
        response.extend_from_slice(&attrs_bytes);
        response
    }

    #[test]
    fn test_xor_mapped_address_ipv4() {
        // IP 203.0.113.1 = 0xCB007101, XOR'd with magic 0x2112A442 = 0xEA12D543
        let attr_data: [u8; 8] = [
            0x00, 0x01, // reserved + family (IPv4)
            0x11, 0x2B, // port XOR'd with 0x2112 (port 0x3039 = 12345)
            0xEA, 0x12, 0xD5, 0x43, // IP XOR'd
        ];
        let data = build_stun_response(&[(ATTR_XOR_MAPPED_ADDRESS, &attr_data)]);
        let txn_id = [0u8; 12];
        let result = parse_stun_response(&data, &txn_id);
        assert_eq!(result, Some("203.0.113.1".to_string()));
    }

    #[test]
    fn test_mapped_address_fallback_ipv4() {
        // IP 192.168.1.1 = 0xC0A80101 (no XOR)
        let attr_data: [u8; 8] = [
            0x00, 0x01, // reserved + family (IPv4)
            0x00, 0x50, // port 80
            0xC0, 0xA8, 0x01, 0x01, // IP
        ];
        let data = build_stun_response(&[(ATTR_MAPPED_ADDRESS, &attr_data)]);
        let txn_id = [0u8; 12];
        let result = parse_stun_response(&data, &txn_id);
        assert_eq!(result, Some("192.168.1.1".to_string()));
    }

    #[test]
    fn test_response_too_short() {
        let data = vec![0u8; 19]; // < 20 bytes
        let txn_id = [0u8; 12];
        assert_eq!(parse_stun_response(&data, &txn_id), None);
    }

    #[test]
    fn test_wrong_msg_type() {
        // Build with correct helper then overwrite msg_type to 0x0001 (Binding Request)
        let mut data = build_stun_response(&[]);
        data[0] = 0x00;
        data[1] = 0x01;
        let txn_id = [0u8; 12];
        assert_eq!(parse_stun_response(&data, &txn_id), None);
    }

    #[test]
    fn test_no_mapped_address_attributes() {
        // Valid response with no attributes
        let data = build_stun_response(&[]);
        let txn_id = [0u8; 12];
        assert_eq!(parse_stun_response(&data, &txn_id), None);
    }

    #[test]
    fn test_xor_preferred_over_mapped() {
        // XOR gives 203.0.113.1, MAPPED gives 192.168.1.1
        let xor_data: [u8; 8] = [
            0x00, 0x01,
            0x11, 0x2B,
            0xEA, 0x12, 0xD5, 0x43,
        ];
        let mapped_data: [u8; 8] = [
            0x00, 0x01,
            0x00, 0x50,
            0xC0, 0xA8, 0x01, 0x01,
        ];
        // XOR listed first — should be preferred
        let data = build_stun_response(&[
            (ATTR_XOR_MAPPED_ADDRESS, &xor_data),
            (ATTR_MAPPED_ADDRESS, &mapped_data),
        ]);
        let txn_id = [0u8; 12];
        let result = parse_stun_response(&data, &txn_id);
        assert_eq!(result, Some("203.0.113.1".to_string()));
    }

    #[test]
    fn test_truncated_attribute_data() {
        // Attribute claims 8 bytes but only 4 are present
        let mut data = build_stun_response(&[]);
        // Manually append a truncated XOR_MAPPED_ADDRESS attribute
        let attr_type = ATTR_XOR_MAPPED_ADDRESS.to_be_bytes();
        let attr_len = 8u16.to_be_bytes(); // claims 8 bytes
        let truncated = [0x00, 0x01, 0x11, 0x2B]; // only 4 bytes
        // Update message length
        let new_msg_len = (attr_type.len() + attr_len.len() + truncated.len()) as u16;
        data[2..4].copy_from_slice(&new_msg_len.to_be_bytes());
        data.extend_from_slice(&attr_type);
        data.extend_from_slice(&attr_len);
        data.extend_from_slice(&truncated);

        let txn_id = [0u8; 12];
        // Should return None, not panic
        assert_eq!(parse_stun_response(&data, &txn_id), 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
}