use bytes::BytesMut;
use niom_turn::constants::*;
// use niom_turn::stun; // not needed; use specific functions via path when required
use std::net::SocketAddr;
use tokio::net::UdpSocket;
use std::time::Duration;

// Use shared decoder from library: niom_turn::stun::decode_xor_relayed_address

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    tracing_subscriber::fmt::init();
    let server: SocketAddr = "127.0.0.1:3478".parse()?;
    let local = UdpSocket::bind("0.0.0.0:0").await?;

    let username = "testuser";
    let password = "secretpassword";

    // Build Allocate request (method METHOD_ALLOCATE)
    let mut buf = BytesMut::new();
    buf.extend_from_slice(&METHOD_ALLOCATE.to_be_bytes()); // Allocate Request
    buf.extend_from_slice(&0u16.to_be_bytes()); // length placeholder
    buf.extend_from_slice(&MAGIC_COOKIE_BYTES);
    let trans = [13u8; 12];
    buf.extend_from_slice(&trans);

    // USERNAME
    let uname = username.as_bytes();
    buf.extend_from_slice(&ATTR_USERNAME.to_be_bytes());
    buf.extend_from_slice(&(uname.len() as u16).to_be_bytes());
    buf.extend_from_slice(uname);
    while (buf.len() % 4) != 0 { buf.extend_from_slice(&[0u8]); }

    // MESSAGE-INTEGRITY placeholder
    let mi_attr_offset = buf.len();
    buf.extend_from_slice(&ATTR_MESSAGE_INTEGRITY.to_be_bytes());
    buf.extend_from_slice(&((HMAC_SHA1_LEN as u16).to_be_bytes()));
    let mi_val_pos = buf.len();
    buf.extend_from_slice(&[0u8;20]);
    while (buf.len() % 4) != 0 { buf.extend_from_slice(&[0u8]); }

    // fix length
    let total_len = (buf.len() - 20) as u16;
    let len_bytes = total_len.to_be_bytes();
    buf[2] = len_bytes[0];
    buf[3] = len_bytes[1];

    // compute HMAC over bytes up to MI attribute header
    {
    use hmac::{Hmac, Mac};
    use sha1::Sha1;
    type HmacSha1 = Hmac<Sha1>;
    let mut mac = HmacSha1::new_from_slice(password.as_bytes()).expect("HMAC key");
        mac.update(&buf[..mi_attr_offset]);
        let res = mac.finalize().into_bytes();
        for i in 0..20 { buf[mi_val_pos + i] = res[i]; }
    }

    // send Allocate
    local.send_to(&buf, server).await?;

    // receive response
    let mut r = vec![0u8; 1500];
    let (len, _addr) = local.recv_from(&mut r).await?;
    println!("got {} bytes", len);
    let resp = &r[..len];
    // expect success (RESP_BINDING_SUCCESS) with XOR-RELAYED-ADDRESS attr
    if resp.len() < 20 {
        anyhow::bail!("response too short");
    }
    let typ = u16::from_be_bytes([resp[0], resp[1]]);
    println!("resp type 0x{:04x}", typ);
    if typ != RESP_BINDING_SUCCESS {
        anyhow::bail!("expected success response, got 0x{:04x}", typ);
    }
    // parse attributes
    let length = u16::from_be_bytes([resp[2], resp[3]]) as usize;
    let total = 20 + length;
    let mut offset = 20;
    let mut relay_addr_opt: Option<SocketAddr> = None;
    while offset + 4 <= total {
        let atype = u16::from_be_bytes([resp[offset], resp[offset+1]]);
        let alen = u16::from_be_bytes([resp[offset+2], resp[offset+3]]) as usize;
        offset += 4;
        if offset + alen > total { break; }
        println!("attr type=0x{:04x} len={}", atype, alen);
        println!("raw: {}", hex::encode(&resp[offset..offset+alen]));
            if atype == ATTR_XOR_RELAYED_ADDRESS {
                // XOR-RELAYED-ADDRESS: decode via shared library function
                if let Some(sa) = niom_turn::stun::decode_xor_relayed_address(&resp[offset..offset+alen], &trans) {
                    relay_addr_opt = Some(sa);
                }
            }
        offset += alen;
        let pad = (4 - (alen % 4)) % 4;
        offset += pad;
    }

    let relay_addr = match relay_addr_opt {
        Some(a) => a,
        None => anyhow::bail!("no relay address in response"),
    };

    println!("got relayed addr: {}", relay_addr);

    // send test payload to relay addr
    let payload = b"hello-relay";
    local.send_to(payload, relay_addr).await?;

    // wait for forwarded packet (should arrive via server socket) using tokio timeout
    let mut buf2 = vec![0u8; 1500];
    match tokio::time::timeout(Duration::from_secs(2), local.recv_from(&mut buf2)).await {
        Ok(Ok((l, src))) => {
            println!("received {} bytes from {}", l, src);
            let got = &buf2[..l];
            println!("payload: {:?}", got);
            if got == payload { println!("relay test success"); Ok(()) } else { anyhow::bail!("payload mismatch") }
        }
        Ok(Err(e)) => anyhow::bail!("recv error: {:?}", e),
        Err(_) => anyhow::bail!("no forwarded packet received: timeout"),
    }
}