// testNet.c -- exercises the calog network library over loopback. Each test runs two
// concurrent scripts on separate context threads: a server (listen/bind, then blocking
// accept/recv) and a client (connect/send). A serverReady signal gates the client so there
// is no bind/connect race. Covers TCP (stream round-trip) and UDP (datagram round-trip,
// including the sender address reported by udpRecvFrom).

#define _POSIX_C_SOURCE 200809L

#include "calog.h"
#include "calogNet.h"

#include <stdatomic.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>

#define CHECK(cond, msg) checkImpl((cond), (msg), __FILE__, __LINE__)

#define PUMP_LIMIT 8000
#define TCP_PORT   47811
#define UDP_PORT   47812
#define ENET_PORT  47813

static CalogT         *calog            = NULL;
static _Atomic bool    serverListening   = false;
static _Atomic int32_t finishedCount     = 0;
static _Atomic int32_t errorCount        = 0;
static char            serverGot[64]     = { 0 };
static _Atomic int32_t serverGotLen      = -1;
static char            clientGot[64]     = { 0 };
static _Atomic int32_t clientGotLen      = -1;
static int32_t         testsRun          = 0;
static int32_t         testsFailed       = 0;

static void    checkImpl(bool condition, const char *message, const char *file, int32_t line);
static void    copyReport(CalogValueT *arg, char *dest, _Atomic int32_t *lenOut);
static int32_t nativeFinished(CalogValueT *args, int32_t argCount, CalogValueT *result, void *userData);
static int32_t nativeReportClient(CalogValueT *args, int32_t argCount, CalogValueT *result, void *userData);
static int32_t nativeReportServer(CalogValueT *args, int32_t argCount, CalogValueT *result, void *userData);
static int32_t nativeServerReady(CalogValueT *args, int32_t argCount, CalogValueT *result, void *userData);
static void    onError(uint64_t contextId, const char *message, void *userData);
static void    pumpUntilFinished(void);
static void    pumpUntilListening(void);
static void    resetState(void);
static void    testEnet(void);
static void    testEnetSafety(void);
static void    testTcp(void);
static void    testUdp(void);


static void checkImpl(bool condition, const char *message, const char *file, int32_t line) {
    testsRun++;
    if (!condition) {
        testsFailed++;
        printf("FAIL %s:%d  %s\n", file, line, message);
    }
}


static void copyReport(CalogValueT *arg, char *dest, _Atomic int32_t *lenOut) {
    size_t length;

    length = (size_t)arg->as.s.length;
    if (length >= 64) {
        length = 63;
    }
    memcpy(dest, arg->as.s.bytes, length);
    dest[length] = '\0';
    atomic_store(lenOut, (int32_t)length);
}


static int32_t nativeFinished(CalogValueT *args, int32_t argCount, CalogValueT *result, void *userData) {
    (void)args;
    (void)argCount;
    (void)userData;
    atomic_fetch_add(&finishedCount, 1);
    calogValueNil(result);
    return calogOkE;
}


static int32_t nativeReportClient(CalogValueT *args, int32_t argCount, CalogValueT *result, void *userData) {
    (void)userData;
    calogValueNil(result);
    if (argCount != 1 || args[0].type != calogStringE) {
        return calogFail(result, calogErrArgE, "reportClient expects one string");
    }
    copyReport(&args[0], clientGot, &clientGotLen);
    return calogOkE;
}


static int32_t nativeReportServer(CalogValueT *args, int32_t argCount, CalogValueT *result, void *userData) {
    (void)userData;
    calogValueNil(result);
    if (argCount != 1 || args[0].type != calogStringE) {
        return calogFail(result, calogErrArgE, "reportServer expects one string");
    }
    copyReport(&args[0], serverGot, &serverGotLen);
    return calogOkE;
}


static int32_t nativeServerReady(CalogValueT *args, int32_t argCount, CalogValueT *result, void *userData) {
    (void)args;
    (void)argCount;
    (void)userData;
    atomic_store(&serverListening, true);
    calogValueNil(result);
    return calogOkE;
}


static void onError(uint64_t contextId, const char *message, void *userData) {
    (void)contextId;
    (void)userData;
    fprintf(stderr, "  [script error] %s\n", (message != NULL) ? message : "(null)");
    atomic_fetch_add(&errorCount, 1);
}


static void pumpUntilFinished(void) {
    struct timespec ts = { 0, 500000 };
    int             i;

    for (i = 0; i < PUMP_LIMIT; i++) {
        calogPump(calog);
        if (atomic_load(&finishedCount) >= 2) {
            return;
        }
        nanosleep(&ts, NULL);
    }
}


static void pumpUntilListening(void) {
    struct timespec ts = { 0, 500000 };
    int             i;

    for (i = 0; i < PUMP_LIMIT; i++) {
        calogPump(calog);
        if (atomic_load(&serverListening) || atomic_load(&errorCount) != 0) {
            return;
        }
        nanosleep(&ts, NULL);
    }
}


static void resetState(void) {
    atomic_store(&serverListening, false);
    atomic_store(&finishedCount, 0);
    atomic_store(&serverGotLen, -1);
    atomic_store(&clientGotLen, -1);
    memset(serverGot, 0, sizeof(serverGot));
    memset(clientGot, 0, sizeof(clientGot));
}


static void testEnet(void) {
    CalogContextT *server;
    CalogContextT *client;
    char           serverScript[768];
    char           clientScript[768];

    resetState();
    // Server: echo each reliable packet, stop when the client disconnects.
    snprintf(serverScript, sizeof(serverScript),
        "local host = enetHost(%d, 4)\n"
        "serverReady()\n"
        "local running = true\n"
        "local count = 0\n"
        "while running and count < 2000 do\n"
        "  local ev = enetService(host, 50)\n"
        "  if ev.type == 'receive' then\n"
        "    reportServer(ev.data)\n"
        "    enetSend(ev.peer, 0, 'pong:' .. ev.data, true)\n"
        "  elseif ev.type == 'disconnect' then\n"
        "    running = false\n"
        "  end\n"
        "  count = count + 1\n"
        "end\n"
        "enetClose(host)\n"
        "finished()", ENET_PORT);
    // Client: connect, send once, read the reply, then disconnect gracefully.
    snprintf(clientScript, sizeof(clientScript),
        "local host = enetHost(0, 1)\n"
        "local peer = enetConnect(host, '127.0.0.1', %d, 1)\n"
        "local finishedClient = false\n"
        "local count = 0\n"
        "while not finishedClient and count < 2000 do\n"
        "  local ev = enetService(host, 50)\n"
        "  if ev.type == 'connect' then\n"
        "    enetSend(peer, 0, 'hi', true)\n"
        "  elseif ev.type == 'receive' then\n"
        "    reportClient(ev.data)\n"
        "    enetDisconnect(peer)\n"
        "  elseif ev.type == 'disconnect' then\n"
        "    finishedClient = true\n"
        "  end\n"
        "  count = count + 1\n"
        "end\n"
        "enetClose(host)\n"
        "finished()", ENET_PORT);

    server = calogContextOpen(calog, &calogLuaEngine);
    calogContextEval(server, serverScript);
    pumpUntilListening();
    client = calogContextOpen(calog, &calogLuaEngine);
    calogContextEval(client, clientScript);
    pumpUntilFinished();

    CHECK(strcmp(serverGot, "hi") == 0, "ENet server received the reliable packet");
    CHECK(strcmp(clientGot, "pong:hi") == 0, "ENet client received the reliable reply over its peer");
    CHECK(atomic_load(&errorCount) == 0, "no errors during the ENet round-trip");

    calogContextClose(server);
    calogContextClose(client);
}


static void testEnetSafety(void) {
    CalogContextT  *ctx;
    struct timespec ts = { 0, 500000 };
    int32_t         before;
    int             i;

    resetState();
    before = atomic_load(&errorCount);
    ctx    = calogContextOpen(calog, &calogLuaEngine);
    // Using a peer handle after its host is closed must raise a clean script error, NOT a
    // use-after-free (enetClose drops the host's peer handles). ASan would flag any UAF.
    calogContextEval(ctx,
        "local h = enetHost(0, 1)\n"
        "local p = enetConnect(h, '127.0.0.1', 59999, 1)\n"
        "enetClose(h)\n"
        "enetSend(p, 0, 'x', true)\n"
        "done()");
    for (i = 0; i < PUMP_LIMIT; i++) {
        calogPump(calog);
        if (atomic_load(&errorCount) > before) {
            break;
        }
        nanosleep(&ts, NULL);
    }

    CHECK(atomic_load(&errorCount) > before, "using an ENet peer handle after enetClose is a clean error, not a use-after-free");

    calogContextClose(ctx);
}


static void testTcp(void) {
    CalogContextT *server;
    CalogContextT *client;
    char           serverScript[512];
    char           clientScript[512];

    resetState();
    snprintf(serverScript, sizeof(serverScript),
        "local l = tcpListen(%d)\n"
        "serverReady()\n"
        "local c = tcpAccept(l)\n"
        "local d = tcpRecv(c, 100)\n"
        "reportServer(d)\n"
        "tcpSend(c, 'pong:' .. d)\n"
        "tcpClose(c)\n"
        "tcpClose(l)\n"
        "finished()", TCP_PORT);
    snprintf(clientScript, sizeof(clientScript),
        "local c = tcpConnect('127.0.0.1', %d)\n"
        "tcpSend(c, 'ping')\n"
        "local r = tcpRecv(c, 100)\n"
        "reportClient(r)\n"
        "tcpClose(c)\n"
        "finished()", TCP_PORT);

    server = calogContextOpen(calog, &calogLuaEngine);
    calogContextEval(server, serverScript);
    pumpUntilListening();
    client = calogContextOpen(calog, &calogLuaEngine);
    calogContextEval(client, clientScript);
    pumpUntilFinished();

    CHECK(strcmp(serverGot, "ping") == 0, "TCP server received the client's stream payload");
    CHECK(strcmp(clientGot, "pong:ping") == 0, "TCP client received the server's reply");
    CHECK(atomic_load(&errorCount) == 0, "no errors during the TCP round-trip");

    calogContextClose(server);
    calogContextClose(client);
}


static void testUdp(void) {
    CalogContextT *server;
    CalogContextT *client;
    char           serverScript[512];
    char           clientScript[512];

    resetState();
    // The server replies to the address udpRecvFrom reports, so the client need not bind a
    // known port -- this also checks the {data, host, port} map.
    snprintf(serverScript, sizeof(serverScript),
        "local s = udpOpen(%d)\n"
        "serverReady()\n"
        "local m = udpRecvFrom(s, 100)\n"
        "reportServer(m.data)\n"
        "udpSendTo(s, m.host, m.port, 'pong:' .. m.data)\n"
        "udpClose(s)\n"
        "finished()", UDP_PORT);
    snprintf(clientScript, sizeof(clientScript),
        "local c = udpOpen(0)\n"
        "udpSendTo(c, '127.0.0.1', %d, 'hi')\n"
        "local m = udpRecvFrom(c, 100)\n"
        "reportClient(m.data)\n"
        "udpClose(c)\n"
        "finished()", UDP_PORT);

    server = calogContextOpen(calog, &calogLuaEngine);
    calogContextEval(server, serverScript);
    pumpUntilListening();
    client = calogContextOpen(calog, &calogLuaEngine);
    calogContextEval(client, clientScript);
    pumpUntilFinished();

    CHECK(strcmp(serverGot, "hi") == 0, "UDP server received the datagram");
    CHECK(strcmp(clientGot, "pong:hi") == 0, "UDP client received the reply routed to its reported address");
    CHECK(atomic_load(&errorCount) == 0, "no errors during the UDP round-trip");

    calogContextClose(server);
    calogContextClose(client);
}


int main(void) {
    calog = calogCreate();
    if (calog == NULL) {
        printf("calog create failed\n");
        return 1;
    }
    calogSetErrorHandler(calog, onError, NULL);
    calogRegister(calog, "serverReady", nativeServerReady, NULL);
    calogRegister(calog, "reportServer", nativeReportServer, NULL);
    calogRegister(calog, "reportClient", nativeReportClient, NULL);
    calogRegister(calog, "finished", nativeFinished, NULL);
    if (calogNetRegister(calog) != calogOkE) {
        printf("calogNetRegister failed\n");
        return 1;
    }

    testTcp();
    testUdp();
    testEnet();
    testEnetSafety();

    calogDestroy(calog);
    calogNetShutdown();

    printf("\n%d checks, %d failed\n", testsRun, testsFailed);
    fflush(stdout);
    if (testsFailed != 0) {
        return 1;
    }
    return 0;
}