// User application entry points for tcpserve.

#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 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: tcpserve.exe [port]\n");
    print("default: tcpserve.exe 8080\n");
}

// Runs the requested operation.
int main(void) {
    uint16_t port = 8080;

    const char* cmd = SysGetPeb()->CommandLine.Data;
    char tok[64];
    cmd = next_token(cmd, tok, sizeof(tok)); // program path
    if (cmd) {
        (void)next_token(cmd, tok, sizeof(tok));
        if (tok[0] && !parse_port(tok, &port)) {
            usage();
            SysExitProcess(1);
        }
    }

    uint32_t listener = 0;
    uint32_t st = SysCreateSocket(SYS_AF_INET, SYS_SOCK_STREAM, SYS_IPPROTO_TCP, &listener);
    if (st != SYS_STATUS_SUCCESS) {
        print("tcpserve: SysCreateSocket failed "); print_u64(st); print("\n");
        SysExitProcess(1);
    }

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

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

    st = SysListenSocket(listener, 4);
    if (st != SYS_STATUS_SUCCESS) {
        print("tcpserve: SysListenSocket failed "); print_u64(st); print("\n");
        SysCloseHandle(listener);
        SysExitProcess(1);
    }

    print("tcpserve: listening on 0.0.0.0:"); print_u64(port); print("\n");
    print("tcpserve: accepting one connection...\n");

    SYS_SOCKET_ADDRESS_IPV4 remote;
    size_t remoteSize = sizeof(remote);
    uint32_t client = 0;
    st = SysAcceptSocket(listener, &remote, &remoteSize, &client);
    if (st != SYS_STATUS_SUCCESS) {
        print("tcpserve: SysAcceptSocket failed "); print_u64(st); print("\n");
        SysCloseHandle(listener);
        SysExitProcess(1);
    }

    print("tcpserve: accepted from ");
    if (remote.Family == SYS_AF_INET) {
        print_ipv4(remote.Address);
        print(":"); print_u64(remote.Port);
    } else {
        print("unknown");
    }
    print("\n");

    char request[1025];
    size_t got = 0;
    st = SysRecvSocket(client, request, sizeof(request) - 1, 0, &got);
    if (st != SYS_STATUS_SUCCESS) {
        print("tcpserve: SysRecvSocket failed "); print_u64(st); print("\n");
        SysCloseHandle(client);
        SysCloseHandle(listener);
        SysExitProcess(1);
    }
    request[got] = 0;
    print("tcpserve: received "); print_u64(got); print(" bytes\n");
    if (got) {
        print("--- request begins ---\n");
        print(request);
        print("\n--- request ends ---\n");
    }

    const char body[] = "hello from nt-lite tcpserve\n";
    char response[256];
    size_t off = 0;
#define APPEND_LITERAL(x) do { const char* _s = (x); while (*_s && off + 1 < sizeof(response)) response[off++] = *_s++; response[off] = 0; } while (0)
    APPEND_LITERAL("HTTP/1.0 200 OK\r\nContent-Type: text/plain\r\nConnection: close\r\nContent-Length: ");
    char lenbuf[24];
    uint64_t len = sizeof(body) - 1;
    int n = 0;
    if (len == 0) lenbuf[n++] = '0';
    else {
        char rev[24];
        while (len && n < (int)sizeof(rev)) { rev[n++] = (char)('0' + (len % 10)); len /= 10; }
        for (int i = n - 1; i >= 0 && off + 1 < sizeof(response); --i) response[off++] = rev[i];
        response[off] = 0;
        n = 0;
    }
    if (n) {
        for (int i = 0; i < n && off + 1 < sizeof(response); ++i) response[off++] = lenbuf[i];
        response[off] = 0;
    }
    APPEND_LITERAL("\r\n\r\n");
#undef APPEND_LITERAL

    size_t sent = 0;
    st = SysSendSocket(client, response, off, 0, &sent);
    if (st == SYS_STATUS_SUCCESS)
        st = SysSendSocket(client, body, sizeof(body) - 1, 0, &sent);
    if (st != SYS_STATUS_SUCCESS)
        print("tcpserve: SysSendSocket failed "), print_u64(st), print("\n");

    SysShutdownSocket(client, SYS_SHUTDOWN_BOTH);
    SysCloseHandle(client);
    SysCloseHandle(listener);
    print("tcpserve: done\n");
    SysExitProcess(0);
    return 0;
}