// User application entry points for udp.

#include <stdint.h>
#include <stddef.h>
#include <stdbool.h>
#include <string.h>
#include <ntlite/syscall.h>

// Prints the requested operation.
static void print(const char* s) {
    if (s) { SYS_WRITE_CONSOLE_REQUEST req = { SysSvFromCstr(s) }; SysWriteConsole(&req); }
}

// Prints char.
static void print_char(char c) {
    SYS_WRITE_CONSOLE_REQUEST req = { SysSvFromParts(&c, 1) };
    SysWriteConsole(&req);
}

// Prints u64.
static void print_u64(uint64_t v) {
    char buf[32];
    int n = 0;
    if (v == 0) { print("0"); return; }
    while (v && n < (int)sizeof(buf)) {
        buf[n++] = (char)('0' + (v % 10));
        v /= 10;
    }
    while (n > 0) print_char(buf[--n]);
}

// Prints ipv4.
static void print_ipv4(const uint8_t ip[4]) {
    print_u64(ip[0]); print("."); print_u64(ip[1]); print(".");
    print_u64(ip[2]); print("."); print_u64(ip[3]);
}

// Implements the next token operation.
static const char* next_token(const char* s, char* out, size_t outSize) {
    if (!s || !out || outSize == 0) return NULL;
    while (*s == ' ' || *s == '\t') s++;
    if (*s == 0) { out[0] = 0; return NULL; }

    size_t n = 0;
    while (*s && *s != ' ' && *s != '\t') {
        if (n + 1 < outSize) out[n++] = *s;
        s++;
    }
    out[n] = 0;
    return s;
}

// Parses u8.
static bool parse_u8(const char** p, uint8_t* out) {
    const char* s = *p;
    uint32_t v = 0;
    int digits = 0;
    while (*s >= '0' && *s <= '9') {
        v = v * 10u + (uint32_t)(*s - '0');
        if (v > 255u) return false;
        s++;
        digits++;
    }
    if (digits == 0) return false;
    *out = (uint8_t)v;
    *p = s;
    return true;
}

// Parses ipv4.
static bool parse_ipv4(const char* s, uint8_t ip[4]) {
    if (!s || !ip) return false;
    for (int i = 0; i < 4; i++) {
        if (!parse_u8(&s, &ip[i])) return false;
        if (i != 3) {
            if (*s != '.') return false;
            s++;
        }
    }
    return *s == 0;
}

// Parses port.
static bool parse_port(const char* s, uint16_t* out) {
    if (!s || !*s || !out) return false;
    uint32_t v = 0;
    while (*s >= '0' && *s <= '9') {
        v = v * 10u + (uint32_t)(*s - '0');
        if (v > 65535u) return false;
        s++;
    }
    if (*s != 0 || v == 0) return false;
    *out = (uint16_t)v;
    return true;
}

// Implements the usage operation.
static void usage(void) {
    print("usage:\n");
    print("  udp.exe send <ipv4> <port> [message]\n");
    print("  udp.exe listen <port>\n");
    print("examples:\n");
    print("  udp.exe send 10.0.2.2 9999 hello\n");
    print("  udp.exe listen 9999\n");
}

// Runs send.
static int run_send(const char* cmd) {
    char tok[256];
    uint8_t ip[4];
    uint16_t port = 0;

    cmd = next_token(cmd, tok, sizeof(tok));
    if (!tok[0] || !parse_ipv4(tok, ip)) { usage(); return 1; }
    cmd = next_token(cmd, tok, sizeof(tok));
    if (!tok[0] || !parse_port(tok, &port)) { usage(); return 1; }

    // Skip leading whitespace, then take the entire remaining command line.
    if (cmd) {
        while (*cmd == ' ' || *cmd == '\t') {
            cmd++;
        }
    }

    const char* msg = "hello from nt-lite udp\n";
    if (cmd && *cmd != '\0') {
        msg = cmd;
    }

    uint32_t sock = 0;
    uint32_t st = SysCreateSocket(SYS_AF_INET, SYS_SOCK_DGRAM, SYS_IPPROTO_UDP, &sock);
    if (st != SYS_STATUS_SUCCESS) {
        print("udp: SysCreateSocket failed "); print_u64(st); print("\n");
        return 1;
    }

    SYS_SOCKET_ADDRESS_IPV4 remote;
    memset(&remote, 0, sizeof(remote));
    remote.Family = SYS_AF_INET;
    remote.Port = port;
    remote.Address[0] = ip[0];
    remote.Address[1] = ip[1];
    remote.Address[2] = ip[2];
    remote.Address[3] = ip[3];

    size_t sent = 0;
    SYS_SOCKET_SEND_TO_REQUEST req;
    memset(&req, 0, sizeof(req));
    req.Buffer = msg;
    req.Length = strlen(msg);
    req.RemoteAddress = &remote;
    req.RemoteAddressSize = sizeof(remote);
    req.SentOut = &sent;

    st = SysSendToSocket(sock, &req);
    SysCloseHandle(sock);
    if (st != SYS_STATUS_SUCCESS) {
        print("udp: SysSendToSocket failed "); print_u64(st); print("\n");
        return 1;
    }

    print("udp: sent "); print_u64(sent); print(" bytes to "); print_ipv4(ip); print(":"); print_u64(port); print("\n");
    return 0;
}

// Runs listen.
static int run_listen(const char* cmd) {
    char tok[64];
    uint16_t port = 0;

    (void)next_token(cmd, tok, sizeof(tok));
    if (!tok[0] || !parse_port(tok, &port)) { usage(); return 1; }

    uint32_t sock = 0;
    uint32_t st = SysCreateSocket(SYS_AF_INET, SYS_SOCK_DGRAM, SYS_IPPROTO_UDP, &sock);
    if (st != SYS_STATUS_SUCCESS) {
        print("udp: SysCreateSocket failed "); print_u64(st); print("\n");
        return 1;
    }

    SYS_SOCKET_ADDRESS_IPV4 bindAddr;
    memset(&bindAddr, 0, sizeof(bindAddr));
    bindAddr.Family = SYS_AF_INET;
    bindAddr.Port = port;

    st = SysBindSocket(sock, &bindAddr, sizeof(bindAddr));
    if (st != SYS_STATUS_SUCCESS) {
        print("udp: SysBindSocket failed "); print_u64(st); print("\n");
        SysCloseHandle(sock);
        return 1;
    }

    print("udp: listening on 0.0.0.0:"); print_u64(port); print(" for one datagram\n");

    char buf[1025];
    SYS_SOCKET_ADDRESS_IPV4 remote;
    size_t remoteSize = sizeof(remote);
    size_t got = 0;
    SYS_SOCKET_RECV_FROM_REQUEST req;
    memset(&req, 0, sizeof(req));
    req.Buffer = buf;
    req.Length = sizeof(buf) - 1;
    req.RemoteAddressOut = &remote;
    req.RemoteAddressSizeInOut = &remoteSize;
    req.ReceivedOut = &got;

    st = SysRecvFromSocket(sock, &req);
    SysCloseHandle(sock);
    if (st != SYS_STATUS_SUCCESS) {
        print("udp: SysRecvFromSocket failed "); print_u64(st); print("\n");
        return 1;
    }

    buf[got] = 0;
    print("udp: received "); print_u64(got); print(" bytes from ");
    if (remote.Family == SYS_AF_INET) {
        print_ipv4(remote.Address); print(":"); print_u64(remote.Port);
    } else {
        print("unknown");
    }
    print("\n");
    print(buf);
    if (got && buf[got - 1] != '\n') print("\n");
    return 0;
}

// Runs the requested operation.
int main(void) {
    const char* cmd = SysGetPeb()->CommandLine.Data;
    char tok[64];
    cmd = next_token(cmd, tok, sizeof(tok)); // program path
    cmd = next_token(cmd, tok, sizeof(tok)); // verb
    if (!tok[0]) { usage(); SysExitProcess(1); }

    int rc;
    if (strcmp(tok, "send") == 0) {
        rc = run_send(cmd);
    } else if (strcmp(tok, "listen") == 0) {
        rc = run_listen(cmd);
    } else {
        usage();
        rc = 1;
    }

    SysExitProcess(rc);
    return rc;
}