#![no_std]
#![no_main]

use lancer_user::io::write_bytes;
use lancer_user::net;
use lancer_user::syscall;

const DEFAULT_TIMEOUT_S: i64 = 10;
const MAX_TIMEOUT_S: i64 = 60;

#[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() {
        write_bytes(b"networking not available (use SSH)\n");
        syscall::exit();
    }

    let mut args_buf = [0u8; 128];
    let args_len = net::recv_args(&rx, &mut args_buf);
    let args = &args_buf[..args_len];

    let (_cmd_tok, rest) = net::next_token(args);
    let (sub, remainder) = net::next_token(rest);

    match sub {
        _ if net::bytes_eq(sub, b"send") => cmd_send(&tx, remainder),
        _ if net::bytes_eq(sub, b"listen") => cmd_listen(&rx, &tx, remainder),
        _ if net::bytes_eq(sub, b"close") => cmd_close(&tx),
        _ => {
            write_bytes(b"usage: udp send <ip> <port> <data>\n       udp listen <port> [timeout]\n       udp close\n");
        }
    }

    syscall::exit()
}

fn cmd_send(tx: &lancer_core::packet_ring::PacketRingWriter, args: &[u8]) {
    let (ip_tok, r1) = net::next_token(args);
    let ip = match net::parse_ip(ip_tok) {
        Some(ip) => ip,
        None => {
            write_bytes(b"invalid IP\n");
            return;
        }
    };

    let (port_tok, r2) = net::next_token(r1);
    let port = match net::parse_decimal_u16(port_tok) {
        Some(p) => p,
        None => {
            write_bytes(b"invalid port\n");
            return;
        }
    };

    let data = net::skip_spaces(r2);
    match data.is_empty() {
        true => {
            write_bytes(b"no data\n");
        }
        false => {
            let mut msg = [0u8; 62];
            let payload_len = data.len().min(54);
            msg[0] = net::MSG_UDP_SEND;
            msg[1] = ip[0];
            msg[2] = ip[1];
            msg[3] = ip[2];
            msg[4] = ip[3];
            msg[5] = (port >> 8) as u8;
            msg[6] = port as u8;
            msg[7..7 + payload_len].copy_from_slice(&data[..payload_len]);
            net::send_request(tx, &msg[..7 + payload_len]);
            write_bytes(b"sent\n");
        }
    }
}

fn cmd_listen(
    rx: &lancer_core::packet_ring::PacketRingReader,
    tx: &lancer_core::packet_ring::PacketRingWriter,
    args: &[u8],
) {
    let (port_tok, remainder) = net::next_token(args);
    let port = match net::parse_decimal_u16(port_tok) {
        Some(p) => p,
        None => {
            write_bytes(b"invalid port\n");
            return;
        }
    };

    let (timeout_tok, _) = net::next_token(remainder);
    let timeout_s = match timeout_tok.is_empty() {
        true => DEFAULT_TIMEOUT_S,
        false => match net::parse_decimal_u16(timeout_tok) {
            Some(t) => (t as i64).clamp(1, MAX_TIMEOUT_S),
            None => DEFAULT_TIMEOUT_S,
        },
    };

    let bind_msg = [net::MSG_UDP_BIND, (port >> 8) as u8, port as u8];
    net::send_request(tx, &bind_msg);

    let mut resp_buf = [0u8; 64];
    let n = net::recv_response_blocking(rx, &mut resp_buf);
    match n >= 2 && resp_buf[0] == net::MSG_UDP_BOUND {
        true => match resp_buf[1] {
            0 => write_bytes(b"bound, listening...\n"),
            _ => {
                write_bytes(b"bind failed\n");
                return;
            }
        },
        false => {
            write_bytes(b"bind failed\n");
            return;
        }
    }

    let start_ms = syscall::clock_monotonic_ms();
    let deadline_ms = start_ms + timeout_s * 1000;

    let mut running = true;
    core::iter::from_fn(|| match running {
        false => None,
        true => {
            let now = syscall::clock_monotonic_ms();
            match now >= deadline_ms {
                true => {
                    running = false;
                    None
                }
                false => match net::recv_response(rx, &mut resp_buf) {
                    Some(rn) if rn >= 2 && resp_buf[0] == net::MSG_UDP_RECV && rn >= 7 => {
                        let mut out = [0u8; 128];
                        let mut pos = 0usize;
                        let prefix = b"from ";
                        out[..prefix.len()].copy_from_slice(prefix);
                        pos += prefix.len();
                        pos += net::write_ip_decimal(&mut out[pos..], &resp_buf[1..5]);
                        out[pos] = b':';
                        pos += 1;
                        let src_port = ((resp_buf[5] as u16) << 8) | resp_buf[6] as u16;
                        pos += net::write_u16_decimal(&mut out[pos..], src_port);
                        let mid = b": ";
                        out[pos..pos + mid.len()].copy_from_slice(mid);
                        pos += mid.len();
                        write_bytes(&out[..pos]);
                        if rn > 7 {
                            write_bytes(&resp_buf[7..rn]);
                        }
                        write_bytes(b"\n");
                        Some(())
                    }
                    _ => {
                        syscall::notify_wait(net::NETSOCK_NOTIF_SLOT);
                        Some(())
                    }
                },
            }
        }
    })
    .take(10_000)
    .count();

    let close_msg = [net::MSG_UDP_CLOSE];
    net::send_request(tx, &close_msg);
    write_bytes(b"closed\n");
}

fn cmd_close(tx: &lancer_core::packet_ring::PacketRingWriter) {
    let close_msg = [net::MSG_UDP_CLOSE];
    net::send_request(tx, &close_msg);
    write_bytes(b"closed\n");
}