// Transport-independent test of the agent's connection handler + exec streamer. Drives
// handle_connection() over a socketpair (so it works where vsock loopback isn't available), exercising
// ping, exec stdout/stderr/exit, and the stdin round-trip — the exact NDJSON framing the host speaks.
#define _GNU_SOURCE
#define AGENT_NO_MAIN
#include "../src/agent.c"

#include <pthread.h>
#include <sys/stat.h>

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

static void *server_thread(void *arg) {
    int fd = (int)(intptr_t)arg;
    handle_connection(fd);
    close(fd);
    return NULL;
}

static char *read_line_fd(int fd) {
    static char buf[1 << 20];
    size_t n = 0;
    while (n < sizeof buf - 1) {
        char c;
        ssize_t r = read(fd, &c, 1);
        if (r <= 0) return n ? (buf[n] = 0, buf) : NULL;
        if (c == '\n') break;
        buf[n++] = c;
    }
    buf[n] = 0;
    return buf;
}

// Spawn a server end for one connection; returns the client fd.
static int open_conn(pthread_t *t) {
    int sv[2];
    socketpair(AF_UNIX, SOCK_STREAM, 0, sv);
    pthread_create(t, NULL, server_thread, (void *)(intptr_t)sv[0]);
    return sv[1];
}

int main(void) {
    // ---- ping ----
    {
        pthread_t t;
        int c = open_conn(&t);
        const char *ping = "{\"op\":\"ping\"}\n";
        write(c, ping, strlen(ping));
        char *reply = read_line_fd(c);
        CHECK(reply && strstr(reply, "\"ok\":true"), "ping ok:true");
        CHECK(reply && strstr(reply, "\"version\":2"), "ping version 2");
        CHECK(reply && strstr(reply, "\"os\":\"linux\""), "ping os:linux");
        close(c);
        pthread_join(t, NULL);
    }

    // ---- exec: stdout + stderr + nonzero exit ----
    {
        pthread_t t;
        int c = open_conn(&t);
        const char *ex = "{\"op\":\"exec\",\"command\":\"printf hello; printf oops 1>&2; exit 7\"}\n";
        write(c, ex, strlen(ex));
        char sout[4096] = "", serr[4096] = "";
        int code = -999;
        for (;;) {
            char *l = read_line_fd(c);
            if (!l) break;
            char *tag = NULL, *d = NULL;
            if (json_get_string(l, "t", &tag) != 0) continue;
            if (strcmp(tag, "o") == 0 && json_get_string(l, "d", &d) == 0) { size_t n; unsigned char *b = base64_decode(d, strlen(d), &n); memcpy(sout + strlen(sout), b, n); free(b); }
            else if (strcmp(tag, "e") == 0 && json_get_string(l, "d", &d) == 0) { size_t n; unsigned char *b = base64_decode(d, strlen(d), &n); memcpy(serr + strlen(serr), b, n); free(b); }
            else if (strcmp(tag, "x") == 0) { long v; json_get_number(l, "code", &v); code = (int)v; free(tag); free(d); break; }
            free(tag); free(d);
        }
        close(c);
        pthread_join(t, NULL);
        CHECK(strcmp(sout, "hello") == 0, "exec stdout == 'hello'");
        CHECK(strcmp(serr, "oops") == 0, "exec stderr == 'oops'");
        CHECK(code == 7, "exec exit code == 7");
    }

    // ---- exec: stdin round-trip via `cat` ----
    {
        pthread_t t;
        int c = open_conn(&t);
        const char *cat = "{\"op\":\"exec\",\"command\":\"cat\"}\n";
        write(c, cat, strlen(cat));
        const char *inframe = "{\"t\":\"i\",\"d\":\"cGluZwo=\"}\n";  // base64("ping\n")
        write(c, inframe, strlen(inframe));
        const char *eof = "{\"t\":\"eof\"}\n";
        write(c, eof, strlen(eof));
        char cout[4096] = "";
        for (;;) {
            char *l = read_line_fd(c);
            if (!l) break;
            char *tag = NULL, *d = NULL;
            if (json_get_string(l, "t", &tag) != 0) continue;
            if (strcmp(tag, "o") == 0 && json_get_string(l, "d", &d) == 0) { size_t n; unsigned char *b = base64_decode(d, strlen(d), &n); memcpy(cout + strlen(cout), b, n); free(b); }
            else if (strcmp(tag, "x") == 0) { free(tag); free(d); break; }
            free(tag); free(d);
        }
        close(c);
        pthread_join(t, NULL);
        CHECK(strcmp(cout, "ping\n") == 0, "exec stdin->stdout round-trip");
    }

    // ---- exec: command with quotes/escapes in the JSON string ----
    {
        pthread_t t;
        int c = open_conn(&t);
        // command: printf '%s' "a\"b"  → expects: a"b
        const char *ex = "{\"op\":\"exec\",\"command\":\"printf '%s' 'a\\\"b'\"}\n";
        write(c, ex, strlen(ex));
        char sout[256] = "";
        for (;;) {
            char *l = read_line_fd(c);
            if (!l) break;
            char *tag = NULL, *d = NULL;
            if (json_get_string(l, "t", &tag) != 0) continue;
            if (strcmp(tag, "o") == 0 && json_get_string(l, "d", &d) == 0) { size_t n; unsigned char *b = base64_decode(d, strlen(d), &n); memcpy(sout + strlen(sout), b, n); free(b); }
            else if (strcmp(tag, "x") == 0) { free(tag); free(d); break; }
            free(tag); free(d);
        }
        close(c);
        pthread_join(t, NULL);
        CHECK(strcmp(sout, "a\"b") == 0, "exec command with escaped quote");
    }

    // ---- shell-events: drain a spool dir, then truncate ----
    {
        // Build a throwaway spool dir with two *.jsonl batches (plus a non-matching file to ignore).
        char dir[] = "/tmp/nash-drain-XXXXXX";
        CHECK(mkdtemp(dir) != NULL, "shell-events: made temp spool dir");
        char p1[256], p2[256], p3[256];
        snprintf(p1, sizeof p1, "%s/aaa.jsonl", dir);
        snprintf(p2, sizeof p2, "%s/bbb.jsonl", dir);
        snprintf(p3, sizeof p3, "%s/ignore.txt", dir);
        const char *b1 = "{\"type\":\"shell-batch\",\"shellId\":\"aaa\",\"events\":[]}\n";
        const char *b2 = "{\"type\":\"shell-batch\",\"shellId\":\"bbb\",\"events\":[]}\n";
        FILE *f;
        f = fopen(p1, "w"); fwrite(b1, 1, strlen(b1), f); fclose(f);
        f = fopen(p2, "w"); fwrite(b2, 1, strlen(b2), f); fclose(f);
        f = fopen(p3, "w"); fwrite("nope\n", 1, 5, f); fclose(f);

        pthread_t t;
        int c = open_conn(&t);
        char reqbuf[512];
        snprintf(reqbuf, sizeof reqbuf, "{\"op\":\"shell-events\",\"dir\":\"%s\",\"truncate\":1}\n", dir);
        write(c, reqbuf, strlen(reqbuf));
        char *reply = read_line_fd(c);
        char events[4096] = "";
        int ndrained = 0;
        if (reply) {
            CHECK(strstr(reply, "\"ok\":true") != NULL, "shell-events ok:true");
            CHECK(strstr(reply, "\"files\":2") != NULL, "shell-events drained 2 files");
            CHECK(strstr(reply, "\"truncated\":false") != NULL, "shell-events not over cap");
            char *ev = NULL;
            if (json_get_string(reply, "events", &ev) == 0) {
                size_t n; unsigned char *b = base64_decode(ev, strlen(ev), &n);
                memcpy(events, b, n); events[n] = 0; free(b); free(ev);
            }
        }
        close(c);
        pthread_join(t, NULL);
        // Both batches present (order is readdir-dependent, so just check each substring).
        CHECK(strstr(events, "\"shellId\":\"aaa\"") && strstr(events, "\"shellId\":\"bbb\""),
              "shell-events returned both batches");
        // truncate:1 → the .jsonl files are gone, the non-matching file remains.
        if (access(p1, F_OK) != 0) ndrained++;
        if (access(p2, F_OK) != 0) ndrained++;
        CHECK(ndrained == 2, "shell-events truncated drained files");
        CHECK(access(p3, F_OK) == 0, "shell-events left non-.jsonl file alone");
        unlink(p3);
        rmdir(dir);
    }

    // ---- shell-events: missing spool dir is not an error ----
    {
        pthread_t t;
        int c = open_conn(&t);
        const char *req = "{\"op\":\"shell-events\",\"dir\":\"/tmp/nash-does-not-exist-xyz\"}\n";
        write(c, req, strlen(req));
        char *reply = read_line_fd(c);
        CHECK(reply && strstr(reply, "\"ok\":true") && strstr(reply, "\"files\":0"),
              "shell-events: missing dir → ok, 0 files");
        close(c);
        pthread_join(t, NULL);
    }

    printf(fails ? "\n%d FAILURES\n" : "\nALL PASS\n", fails);
    return fails ? 1 : 0;
}
