#![no_std]
#![no_main]

use lancer_user::net;
use lancer_user::syscall;

const STATE_NAMES: [&[u8]; 11] = [
    b"Closed",
    b"Listen",
    b"SynSent",
    b"SynRecv",
    b"Established",
    b"FinWait1",
    b"FinWait2",
    b"CloseWait",
    b"Closing",
    b"LastAck",
    b"TimeWait",
];

#[unsafe(no_mangle)]
pub extern "C" fn lancer_main() -> ! {
    let (rx, tx) = match net::init() {
        Some(pair) => pair,
        None => {
            lancer_user::show!(net, error, "netsock init failed");
            syscall::exit();
        }
    };

    if !net::has_relay() {
        lancer_user::io::write_bytes(b"networking not available (use SSH)\n");
        syscall::exit();
    }

    let mut args_buf = [0u8; 64];
    let _ = net::recv_args(&rx, &mut args_buf);

    lancer_user::io::write_bytes(
        b"Proto  Local             Remote            State         SndQ  RcvQ\n",
    );

    net::send_request(&tx, &[net::MSG_NETSTAT_REQUEST]);

    let mut resp = [0u8; 64];
    let mut done = false;
    core::iter::from_fn(|| match done {
        true => None,
        false => {
            let n = net::recv_response_blocking(&rx, &mut resp);
            match n < 1 {
                true => None,
                false => match resp[0] {
                    net::MSG_NETSTAT_DONE => {
                        done = true;
                        None
                    }
                    net::MSG_NETSTAT_ENTRY if n >= 18 => {
                        format_tcp_entry(&resp[1..n]);
                        Some(())
                    }
                    net::MSG_NETSTAT_ENTRY_UDP if n >= 4 => {
                        format_udp_entry(&resp[1..n]);
                        Some(())
                    }
                    _ => Some(()),
                },
            }
        }
    })
    .take(32)
    .count();

    syscall::exit()
}

fn format_tcp_entry(data: &[u8]) {
    let state_idx = data[1] as usize;
    let local_ip = &data[2..6];
    let local_port = ((data[6] as u16) << 8) | data[7] as u16;
    let remote_ip = &data[8..12];
    let remote_port = ((data[12] as u16) << 8) | data[13] as u16;
    let send_q = ((data[14] as u16) << 8) | data[15] as u16;
    let recv_q = ((data[16] as u16) << 8) | data[17] as u16;

    let mut out = [b' '; 128];
    let mut pos = 0usize;

    pos = append(&mut out, pos, b"tcp    ");

    let ip_len = net::write_ip_decimal(&mut out[pos..], local_ip);
    pos += ip_len;
    out[pos] = b':';
    pos += 1;
    let port_len = net::write_u16_decimal(&mut out[pos..], local_port);
    pos += port_len;

    pad_to(&mut out, &mut pos, 25);

    let has_remote = remote_ip.iter().any(|&b| b != 0) || remote_port != 0;
    match has_remote {
        true => {
            let rip_len = net::write_ip_decimal(&mut out[pos..], remote_ip);
            pos += rip_len;
            out[pos] = b':';
            pos += 1;
            let rport_len = net::write_u16_decimal(&mut out[pos..], remote_port);
            pos += rport_len;
        }
        false => {
            out[pos] = b'-';
            pos += 1;
        }
    }

    pad_to(&mut out, &mut pos, 43);

    let state_name = match state_idx < STATE_NAMES.len() {
        true => STATE_NAMES[state_idx],
        false => b"?",
    };
    pos = append(&mut out, pos, state_name);

    pad_to(&mut out, &mut pos, 57);

    let sq_len = net::write_u16_decimal(&mut out[pos..], send_q);
    pos += sq_len;

    pad_to(&mut out, &mut pos, 63);

    let rq_len = net::write_u16_decimal(&mut out[pos..], recv_q);
    pos += rq_len;

    out[pos] = b'\n';
    pos += 1;

    lancer_user::io::write_bytes(&out[..pos]);
}

fn format_udp_entry(data: &[u8]) {
    let local_port = ((data[0] as u16) << 8) | data[1] as u16;
    let recv_q = ((data[2] as u16) << 8) | data[3] as u16;

    let mut out = [b' '; 128];
    let mut pos = 0usize;

    pos = append(&mut out, pos, b"udp    0.0.0.0:");
    let port_len = net::write_u16_decimal(&mut out[pos..], local_port);
    pos += port_len;

    pad_to(&mut out, &mut pos, 25);

    out[pos] = b'-';
    pos += 1;

    pad_to(&mut out, &mut pos, 43);

    out[pos] = b'-';
    pos += 1;

    pad_to(&mut out, &mut pos, 63);

    let rq_len = net::write_u16_decimal(&mut out[pos..], recv_q);
    pos += rq_len;

    out[pos] = b'\n';
    pos += 1;

    lancer_user::io::write_bytes(&out[..pos]);
}

fn append(buf: &mut [u8], pos: usize, data: &[u8]) -> usize {
    let n = data.len().min(buf.len() - pos);
    buf[pos..pos + n].copy_from_slice(&data[..n]);
    pos + n
}

fn pad_to(buf: &mut [u8], pos: &mut usize, target: usize) {
    (*pos..target).for_each(|i| {
        buf[i] = b' ';
    });
    *pos = (*pos).max(target);
}