// 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>

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");
    }

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