/* test_sock.c — out-of-process control: wire codec + authentication/admission.
 * Bring up two listeners (one admitting, one rejecting) on abstract sockets, run
 * the core in a separate thread, connect clients and check: policy invoked,
 * valid poller admitted, unauthorized rejected (EOF), reap without a crash. */
#define _GNU_SOURCE
#include "vmsig.h"
#include "vmsig_socket.h"
#include "memctx.h"        /* VMSIG_MEMWRITE_MAX: the adapter's extent bound (private) */
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <stdint.h>
#include <pthread.h>
#include <stdatomic.h>
#include <sys/socket.h>
#include <sys/un.h>
#include <unistd.h>
#include <stddef.h>
#include <time.h>

static int g_fail = 0;
#define CHECK(cond, msg) do { \
    if (!(cond)) { printf("  FAIL: %s\n", (msg)); g_fail = 1; } \
} while (0)

static atomic_int g_auth = 0;
static atomic_int g_deny = 0;
static atomic_int g_admit = 0;
static atomic_int g_reject = 0;

static void audit_cb(void* ud, const vmsig_audit* a) {
    (void)ud;
    if (a->kind == VMSIG_AUDIT_ADMIT)       atomic_fetch_add(&g_admit, 1);
    else if (a->kind == VMSIG_AUDIT_REJECT) atomic_fetch_add(&g_reject, 1);
}

static vmsig_grant pol_ok(uint32_t uid, uint32_t pid, void* ud) {
    (void)pid; (void)ud;
    atomic_fetch_add(&g_auth, 1);
    vmsig_grant g; memset(&g, 0, sizeof g);
    g.principal = uid; g.endpoint_mask = 1u << 0;
    g.source_mask = 0xFFFFFFFFu; g.cap_mask = VMSIG_CAP_OBSERVE;
    return g;
}
static vmsig_grant pol_deny(uint32_t uid, uint32_t pid, void* ud) {
    (void)uid; (void)pid; (void)ud;
    atomic_fetch_add(&g_deny, 1);
    vmsig_grant g; memset(&g, 0, sizeof g);   /* empty => reject */
    return g;
}

static uint64_t now_ns(void) {
    struct timespec ts;
    clock_gettime(CLOCK_MONOTONIC, &ts);
    return (uint64_t)ts.tv_sec * 1000000000ull + (uint64_t)ts.tv_nsec;
}

static int connect_abstract(const char* name) {
    int fd = socket(AF_UNIX, SOCK_STREAM, 0);
    if (fd < 0) return -1;
    struct sockaddr_un a; memset(&a, 0, sizeof a); a.sun_family = AF_UNIX;
    size_t n = strlen(name);
    a.sun_path[0] = 0;
    memcpy(a.sun_path + 1, name + 1, n - 1);
    socklen_t alen = (socklen_t)(offsetof(struct sockaddr_un, sun_path) + n);
    if (connect(fd, (struct sockaddr*)&a, alen) < 0) { close(fd); return -1; }
    return fd;
}

static void* loop_main(void* p) { vmsig_core_run((vmsig_core*)p); return NULL; }

static void wait_atomic(atomic_int* a, int want, int ms) {
    for (int i = 0; i < ms; i++) {
        if (atomic_load(a) >= want) return;
        struct timespec t = { .tv_sec = 0, .tv_nsec = 1000000 };
        nanosleep(&t, NULL);
    }
}

static void test_wire(void) {
    printf("test_wire\n");
    vmsig_event ev; memset(&ev, 0, sizeof ev);
    ev.kind = VMSIG_EV_CMD_VM; ev.source = VMSIG_SRC_VMHOST; ev.dir = VMSIG_DIR_DOWN;
    ev.prio = VMSIG_PRIO_HIGH; ev.endpoint = 0; ev.corr = 0xABCD;
    for (int i = 0; i < 48; i++) ev.inln[i] = (uint8_t)i;

    vmsig_wire w; vmsig_wire_encode(&w, &ev);
    vmsig_event d;
    CHECK(vmsig_wire_decode(&w, &d) == 0, "decode ok");
    CHECK(d.kind == ev.kind && d.source == ev.source &&
          d.endpoint == ev.endpoint && d.corr == ev.corr, "frame fields match");
    CHECK(memcmp(d.inln, ev.inln, 48) == 0, "inln matches");

    vmsig_wire bad = w; bad.magic = 0; vmsig_event x;
    CHECK(vmsig_wire_decode(&bad, &x) == -1, "bad magic rejected");
}

/* ===== variant B: socket CMD_MEMWRITE with a length-prefixed SRC tail (>INLINE) =====
 * Stub memctx adapter (no VM): proves the socket transport carries a frame + SRC tail
 * (with partial recv), routes through cap -> grant -> lease-gate -> adapter, and gets
 * ACT_ACK ok=1. Also: an over-cap len is a framing violation that closes the connection. */

/* Policy granting a MEMWRITE-capable poller (cap MEMWRITE|MEMCTX|OBSERVE). */
static vmsig_grant pol_mw(uint32_t uid, uint32_t pid, void* ud) {
    (void)pid; (void)ud;
    vmsig_grant g; memset(&g, 0, sizeof g);
    g.principal = uid; g.endpoint_mask = 1ull << 0;
    g.source_mask = 0xFFFFFFFFu;
    g.cap_mask = VMSIG_CAP_MEMWRITE | VMSIG_CAP_MEMCTX | VMSIG_CAP_OBSERVE;
    g.arb_prio = 10;
    return g;
}

/* Encode + write a single fixed frame. */
static int send_frame(int fd, const vmsig_event* ev) {
    vmsig_wire w; vmsig_wire_encode(&w, ev);
    return (write(fd, &w, sizeof w) == (ssize_t)sizeof w) ? 0 : -1;
}

/* Read fixed frames until an ACT_ACK with the wanted corr; return its ok flag (-1 on
 * timeout/EOF). The ACK inln layout from mc_memwrite_ack: {int ok; uint32 corr; uint32 origin}. */
static int wait_ack(int fd, uint32_t want_corr, int ms) {
    struct timeval tv = { .tv_sec = 0, .tv_usec = 200 * 1000 };
    setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof tv);
    uint64_t deadline = now_ns() + (uint64_t)ms * 1000000ull;
    vmsig_wire w; uint8_t* p = (uint8_t*)&w; size_t off = 0;
    while (now_ns() < deadline) {
        ssize_t n = read(fd, p + off, sizeof w - off);
        if (n <= 0) continue;                       /* timeout/EOF retry within deadline */
        off += (size_t)n;
        if (off < sizeof w) continue;
        off = 0;
        vmsig_event ev;
        if (vmsig_wire_decode(&w, &ev) != 0) continue;
        if (ev.kind == VMSIG_EV_ACT_ACK && ev.corr == want_corr) {
            int ok; memcpy(&ok, ev.inln, sizeof ok);
            return ok;
        }
    }
    return -1;
}

static void test_memwrite_tail(void) {
    printf("test_memwrite_tail\n");
    vmsig_ctx*  ctx  = vmsig_ctx_new();
    vmsig_core* core = vmsig_core_new(ctx);
    CHECK(vmsig_core_add_adapter(core, vmsig_memctx_ops(), NULL, 0) >= 0, "add memctx stub");
    const char* MW = "@vmsig-sock-mw-test";
    CHECK(vmsig_socket_attach(core, MW, pol_mw, NULL) == 0, "attach mw listener");

    pthread_t th;
    pthread_create(&th, NULL, loop_main, core);

    int fd = connect_abstract(MW);
    CHECK(fd >= 0, "client connected (mw)");
    if (fd >= 0) {
        /* acquire the MEMWRITE lease */
        vmsig_event d; memset(&d, 0, sizeof d);
        d.kind = VMSIG_EV_CMD_ACQUIRE; d.source = VMSIG_SRC_MEMCTX; d.dir = VMSIG_DIR_DOWN;
        d.endpoint = 0; d.prio = VMSIG_PRIO_HIGH;
        vmsig_lease_req lr = { VMSIG_LEASE_MEMWRITE, 0 };
        memcpy(d.inln, &lr, sizeof lr);
        CHECK(send_frame(fd, &d) == 0, "send ACQUIRE");

        /* happy path: CMD_MEMWRITE(PAYLOAD, len=64) + 64-byte tail, written in two halves
         * to exercise the TAIL-phase partial accumulation. */
        const uint32_t len = 64u;
        uint8_t src[64];
        for (uint32_t i = 0; i < len; i++) src[i] = (uint8_t)(i + 1);
        vmsig_event mwe; memset(&mwe, 0, sizeof mwe);
        mwe.kind = VMSIG_EV_CMD_MEMWRITE; mwe.source = VMSIG_SRC_MEMCTX; mwe.dir = VMSIG_DIR_DOWN;
        mwe.endpoint = 0; mwe.prio = VMSIG_PRIO_HIGH; mwe.corr = 0x101;
        vmsig_memwrite mw = { 0, 0x1000, len, VMSIG_MW_SRC_PAYLOAD };
        memcpy(mwe.inln, &mw, sizeof mw);
        CHECK(send_frame(fd, &mwe) == 0, "send CMD_MEMWRITE frame (PAYLOAD)");
        CHECK(write(fd, src, 32) == 32, "send SRC tail part 1");
        struct timespec ts = { .tv_sec = 0, .tv_nsec = 5 * 1000000 };
        nanosleep(&ts, NULL);                         /* let the loop accumulate a partial tail */
        CHECK(write(fd, src + 32, 32) == 32, "send SRC tail part 2");
        CHECK(wait_ack(fd, 0x101, 1000) == 1, "B: payload-tail write ACKs ok=1 (stub)");

        /* negative: an over-cap PAYLOAD len is a framing-contract violation. The server closes
         * the connection — it cannot safely skip the promised tail, and draining an arbitrary
         * length would be a DoS. Verify no ACK arrives and the socket reaches EOF (conn shut). */
        memset(&mwe.inln, 0, sizeof mwe.inln);
        mwe.corr = 0x102;
        vmsig_memwrite mw2 = { 0, 0x2000, VMSIG_MEMWRITE_MAX + 1u, VMSIG_MW_SRC_PAYLOAD };
        memcpy(mwe.inln, &mw2, sizeof mw2);
        CHECK(send_frame(fd, &mwe) == 0, "send CMD_MEMWRITE frame (over-cap)");
        /* No ACK arrives; the server shuts the conn, so the socket drains to EOF. A 1s recv
         * timeout bounds the wait if the server wrongly kept the connection open. */
        struct timeval rtv = { .tv_sec = 1, .tv_usec = 0 };
        setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, &rtv, sizeof rtv);
        uint8_t junk[80]; ssize_t rr;
        while ((rr = read(fd, junk, sizeof junk)) > 0) { /* drain any in-flight, then EOF */ }
        CHECK(rr == 0, "B: over-cap closed the connection (EOF)");

        close(fd);
    }

    struct timespec t = { .tv_sec = 0, .tv_nsec = 50 * 1000000 };
    nanosleep(&t, NULL);
    vmsig_core_stop(core);
    pthread_join(th, NULL);
    vmsig_core_free(core);
    vmsig_ctx_free(ctx);
}

int main(void) {
    test_wire();
    test_memwrite_tail();

    printf("test_socket\n");
    vmsig_ctx*  ctx  = vmsig_ctx_new();
    vmsig_core* core = vmsig_core_new(ctx);
    vmsig_core_set_audit(core, audit_cb, NULL);
    const char* OK   = "@vmsig-sock-ok-test";
    const char* DENY = "@vmsig-sock-deny-test";
    CHECK(vmsig_socket_attach(core, OK,   pol_ok,   NULL) == 0, "attach ok listener");
    CHECK(vmsig_socket_attach(core, DENY, pol_deny, NULL) == 0, "attach deny listener");

    pthread_t th;
    pthread_create(&th, NULL, loop_main, core);

    /* valid poller: connect -> policy -> admission */
    int c1 = connect_abstract(OK);
    CHECK(c1 >= 0, "client connected (ok)");
    wait_atomic(&g_auth, 1, 1000);
    CHECK(atomic_load(&g_auth) >= 1, "policy invoked — poller authenticated/admitted");
    if (c1 >= 0) close(c1);   /* disconnect -> deferred reap (no crash) */

    /* unauthorized: connect -> server closes -> EOF on the client */
    int c2 = connect_abstract(DENY);
    CHECK(c2 >= 0, "client connected (deny)");
    wait_atomic(&g_deny, 1, 1000);
    CHECK(atomic_load(&g_deny) >= 1, "deny policy invoked");
    if (c2 >= 0) {
        struct timeval tv = { .tv_sec = 1, .tv_usec = 0 };
        setsockopt(c2, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof tv);
        char b; ssize_t r = read(c2, &b, 1);
        CHECK(r == 0, "connection rejected by server (EOF)");
        close(c2);
    }

    /* slot reuse: churn > MAX_CONTROLS(64). Without returning slots the listener
     * would die after 64 cycles. Each cycle: connect(ok) -> wait auth++ -> close. */
    int base = atomic_load(&g_auth);
    const int churn = 70;
    for (int k = 0; k < churn; k++) {
        int fc = connect_abstract(OK);
        if (fc < 0) { CHECK(0, "churn connect"); break; }
        wait_atomic(&g_auth, base + k + 1, 1000);
        close(fc);
        struct timespec ts = { .tv_sec = 0, .tv_nsec = 2 * 1000000 };
        nanosleep(&ts, NULL);     /* let the loop reap before the next connection */
    }
    CHECK(atomic_load(&g_auth) >= base + churn,
          "slots reused: churn > MAX_CONTROLS admitted");

    /* audit recorded admissions and rejections */
    CHECK(atomic_load(&g_admit) >= 1, "audit: poller admission");
    CHECK(atomic_load(&g_reject) >= 1, "audit: rejection (deny listener)");

    struct timespec t = { .tv_sec = 0, .tv_nsec = 50 * 1000000 };
    nanosleep(&t, NULL);          /* let the loop process the reaps */
    vmsig_core_stop(core);
    pthread_join(th, NULL);
    vmsig_core_free(core);
    vmsig_ctx_free(ctx);

    printf("socket tests: %s\n", g_fail ? "FAIL" : "PASS");
    return g_fail ? 1 : 0;
}