#include "mcp.h"

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

#include <fcntl.h>
#include <signal.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <unistd.h>

#include <cjson/cJSON.h>

static int failures = 0;
static int checks = 0;

#define CHECK(cond, msg)                                                     \
    do {                                                                     \
        checks++;                                                            \
        if (!(cond)) {                                                       \
            failures++;                                                      \
            fprintf(stderr, "FAIL %s:%d: %s\n", __FILE__, __LINE__, (msg)); \
        }                                                                    \
    } while (0)

#define CHECK_EQ(a, b, msg)                                                  \
    do {                                                                     \
        checks++;                                                            \
        if (strcmp((a) ? (a) : "(null)", (b)) != 0) {                        \
            failures++;                                                      \
            fprintf(stderr, "FAIL %s:%d: %s (got `%s`, want `%s`)\n",        \
                    __FILE__, __LINE__, (msg), (a) ? (a) : "(null)", (b));   \
        }                                                                    \
    } while (0)

#define CHECK_CONT(hay, needle, msg)                                         \
    do {                                                                     \
        checks++;                                                            \
        if (!hay || !strstr(hay, needle)) {                                  \
            failures++;                                                      \
            fprintf(stderr,                                                  \
                    "FAIL %s:%d: %s (got `%s`, want substring `%s`)\n",      \
                    __FILE__, __LINE__, (msg), hay ? hay : "(null)",         \
                    (needle));                                               \
        }                                                                    \
    } while (0)

static void mem_free(mcp_memory *m)
{
    sl_free(&m->responses);
    sl_free(&m->sent);
}

/* helpers: push a queued JSON-RPC response / read the nth sent message */

static void push_resp(mcp_memory *m, unsigned long id, const char *result)
{
    char *line = xstrfmt("{\"jsonrpc\":\"2.0\",\"id\":%lu,\"result\":%s}", id,
                         result);
    sl_push(&m->responses, line);
}

static cJSON *sent_msg(mcp_memory *m, size_t i)
{
    if (i >= m->sent.len)
        return NULL;
    return cJSON_Parse(m->sent.items[i]);
}

static void test_parses_stdio_config(void)
{
    mcp_server_config cfg;
    CHECK(mcp_server_config_from_json(
              "{\"command\":\"node\",\"args\":[\"server.js\",\"--fast\"],"
              "\"env\":{\"API_KEY\":\"k\"}}",
              &cfg) == 1,
          "stdio config parses");
    CHECK_EQ(cfg.command, "node", "command");
    CHECK(cfg.arg_count == 2, "args count");
    CHECK(cfg.arg_count == 2 && strcmp(cfg.args[0], "server.js") == 0,
          "arg 0");
    CHECK(cfg.arg_count == 2 && strcmp(cfg.args[1], "--fast") == 0, "arg 1");
    CHECK(cfg.env_count == 1 && strcmp(cfg.env[0].key, "API_KEY") == 0,
          "env key");
    CHECK(cfg.env_count == 1 && strcmp(cfg.env[0].val, "k") == 0, "env value");
    CHECK(cfg.url == NULL, "no url");
    mcp_server_config_free(&cfg);
}

static void test_parses_http_config(void)
{
    mcp_server_config cfg;
    CHECK(mcp_server_config_from_json(
              "{\"url\":\"https://example.com/mcp\","
              "\"headers\":{\"Authorization\":\"Bearer t\"}}",
              &cfg) == 1,
          "http config parses");
    CHECK_EQ(cfg.url, "https://example.com/mcp", "url");
    CHECK(cfg.header_count == 1 &&
              strcmp(cfg.headers[0].key, "Authorization") == 0,
          "header key");
    CHECK(cfg.header_count == 1 &&
              strcmp(cfg.headers[0].val, "Bearer t") == 0,
          "header value");
    CHECK(cfg.command == NULL, "no command");
    mcp_server_config_free(&cfg);
}

static void test_rejects_no_transport(void)
{
    mcp_server_config cfg;
    CHECK(mcp_server_config_from_json("{\"foo\":\"bar\"}", &cfg) == 0,
          "object without transport");
    mcp_server_config_free(&cfg);
    CHECK(mcp_server_config_from_json("42", &cfg) == 0, "non-object");
    mcp_server_config_free(&cfg);
    CHECK(mcp_server_config_from_json(NULL, &cfg) == 0, "null");
    mcp_server_config_free(&cfg);
}

static void test_handshake_discovery_call(void)
{
    mcp_memory mem;
    mcp_memory_init(&mem);
    push_resp(&mem, 1,
              "{\"protocolVersion\":\"2024-11-05\",\"capabilities\":{},\"serverInfo\":{\"name\":\"fake\",\"version\":\"0\"}}");
    push_resp(&mem, 2,
              "{\"tools\":[{\"name\":\"add\",\"description\":\"Add two "
              "integers\",\"inputSchema\":{\"type\":\"object\",\"properties\":"
              "{\"a\":{\"type\":\"integer\"},\"b\":{\"type\":\"integer\"}},"
              "\"required\":[\"a\",\"b\"]}}]}");
    push_resp(&mem, 3, "{\"content\":[{\"type\":\"text\",\"text\":\"7\"}]}");

    char *err = NULL;
    mcp_client *c = mcp_connect_memory(&mem, &err);
    CHECK(c != NULL, "connect memory");

    CHECK(mcp_initialize(c, &err) == 1, "handshake");
    CHECK(err == NULL, "no handshake error");

    mcp_tool *tools = NULL;
    size_t n = 0;
    err = NULL;
    CHECK(mcp_list_tools(c, &tools, &n, &err) == 1, "tools/list");
    CHECK(n == 1, "one tool");
    if (n == 1) {
        CHECK_EQ(tools[0].name, "add", "tool name");
        CHECK_EQ(tools[0].description, "Add two integers", "tool description");
        CHECK_CONT(tools[0].input_schema, "\"type\":\"object\"", "schema type");
        CHECK_CONT(tools[0].input_schema, "\"required\"", "schema required");
    }

    err = NULL;
    char *out = mcp_call_tool(c, "add", "{\"a\":3,\"b\":4}", &err);
    CHECK(out != NULL, "tools/call");
    CHECK_EQ(out, "7", "call result");
    free(out);

    /* wire order and ids must be exactly the MCP handshake */
    CHECK(mem.sent.len == 4, "four messages sent");
    const char *want[] = {"initialize", "notifications/initialized",
                          "tools/list", "tools/call"};
    for (size_t i = 0; i < 4 && i < mem.sent.len; i++) {
        cJSON *m = sent_msg(&mem, i);
        cJSON *method = m ? cJSON_GetObjectItemCaseSensitive(m, "method") : NULL;
        CHECK(cJSON_IsString(method) && strcmp(method->valuestring, want[i]) == 0,
              "method order");
        cJSON_Delete(m);
    }
    {
        cJSON *m0 = sent_msg(&mem, 0);
        cJSON *m2 = sent_msg(&mem, 2);
        cJSON *m3 = sent_msg(&mem, 3);
        CHECK(m0 && (unsigned long long)cJSON_GetObjectItemCaseSensitive(m0, "id")->valuedouble == 1,
              "initialize id 1");
        CHECK(m2 && (unsigned long long)cJSON_GetObjectItemCaseSensitive(m2, "id")->valuedouble == 2,
              "tools/list id 2");
        CHECK(m3 && (unsigned long long)cJSON_GetObjectItemCaseSensitive(m3, "id")->valuedouble == 3,
              "tools/call id 3");
        if (m3) {
            cJSON *params = cJSON_GetObjectItemCaseSensitive(m3, "params");
            cJSON *name = params ? cJSON_GetObjectItemCaseSensitive(params, "name") : NULL;
            CHECK(cJSON_IsString(name) && strcmp(name->valuestring, "add") == 0,
                  "call params name");
            cJSON *args = params ? cJSON_GetObjectItemCaseSensitive(params, "arguments") : NULL;
            CHECK(args && (long long)cJSON_GetObjectItemCaseSensitive(args, "a")->valuedouble == 3,
                  "call params a=3");
        }
        cJSON_Delete(m0);
        cJSON_Delete(m2);
        cJSON_Delete(m3);
    }
    mcp_tools_free(tools, n);
    mcp_client_free(c);
    mem_free(&mem);
}

static void test_server_error_is_surfaced(void)
{
    mcp_memory mem;
    mcp_memory_init(&mem);
    sl_push(&mem.responses,
            xstrdup("{\"jsonrpc\":\"2.0\",\"id\":1,\"error\":{\"code\":-32601,"
                    "\"message\":\"Method not found\"}}"));
    char *err = NULL;
    mcp_client *c = mcp_connect_memory(&mem, &err);
    err = NULL;
    char *res = mcp_request(c, "initialize", "{}", &err);
    CHECK(res == NULL, "request failed");
    CHECK(err && strstr(err, "server error -32601") != NULL, "error code");
    CHECK(err && strstr(err, "Method not found") != NULL, "error message");
    free(err);
    mcp_client_free(c);
    mem_free(&mem);
}

static void test_notification_lines_are_skipped(void)
{
    mcp_memory mem;
    mcp_memory_init(&mem);
    sl_push(&mem.responses, xstrdup(
                                "{\"jsonrpc\":\"2.0\",\"method\":"
                                "\"notifications/message\",\"params\":{}}"));
    push_resp(&mem, 1, "{\"ok\":true}");
    char *err = NULL;
    mcp_client *c = mcp_connect_memory(&mem, &err);
    err = NULL;
    char *res = mcp_request(c, "ping", "{}", &err);
    CHECK(res != NULL, "response past the notification");
    CHECK_CONT(res, "\"ok\":true", "result ok");
    free(res);
    mcp_client_free(c);
    mem_free(&mem);
}

static void test_tool_invoke_roundtrip(void)
{
    mcp_memory mem;
    mcp_memory_init(&mem);
    push_resp(&mem, 1,
              "{\"tools\":[{\"name\":\"echo\",\"description\":\"echo\","
              "\"inputSchema\":{\"type\":\"object\"}}]}");
    push_resp(&mem, 2, "{\"content\":[{\"type\":\"text\",\"text\":\"pong\"}]}");
    char *err = NULL;
    mcp_client *c = mcp_connect_memory(&mem, &err);

    mcp_tool *tools = NULL;
    size_t n = 0;
    err = NULL;
    CHECK(mcp_list_tools(c, &tools, &n, &err) == 1, "tools");
    CHECK(n == 1, "one tool");
    if (n == 1) {
        CHECK_EQ(tools[0].name, "echo", "tool name");
        err = NULL;
        char *out = mcp_tool_invoke(&tools[0], "{\"msg\":\"hi\"}", &err);
        CHECK(out != NULL, "execute");
        CHECK_CONT(out, "pong", "server output");
        free(out);
    }
    mcp_tools_free(tools, n);
    mcp_client_free(c);
    mem_free(&mem);
}

static void test_tool_invoke_wraps_errors(void)
{
    mcp_memory mem;
    mcp_memory_init(&mem);
    push_resp(&mem, 1,
              "{\"tools\":[{\"name\":\"boom\",\"description\":\"\",\"inputSchema\":{\"type\":\"object\"}}]}");
    push_resp(&mem, 2,
              "{\"isError\":true,\"content\":[{\"type\":\"text\",\"text\":\"kaboom\"}]}");
    char *err = NULL;
    mcp_client *c = mcp_connect_memory(&mem, &err);
    mcp_tool *tools = NULL;
    size_t n = 0;
    err = NULL;
    CHECK(mcp_list_tools(c, &tools, &n, &err) == 1, "tools");
    if (n == 1) {
        err = NULL;
        char *out = mcp_tool_invoke(&tools[0], "{}", &err);
        CHECK(out == NULL, "isError surfaces as failure");
        CHECK(err && strstr(err, "MCP tool `boom` failed") != NULL,
              "wrapped error");
        CHECK(err && strstr(err, "kaboom") != NULL, "server text included");
        free(err);
    }
    mcp_tools_free(tools, n);
    mcp_client_free(c);
    mem_free(&mem);
}

static void test_extract_json_plain_and_sse(void)
{
    char *err = NULL;
    char *v = mcp_extract_json(
        "application/json",
        "{\"jsonrpc\":\"2.0\",\"id\":1,\"result\":{\"a\":1}}", &err);
    CHECK(v != NULL, "plain parses");
    CHECK_CONT(v, "\"a\":1", "plain result");
    free(v);

    const char *sse = "event: message\ndata: {\"jsonrpc\":\"2.0\",\"id\":2,"
                      "\"result\":{\"b\":2}}\n\n";
    err = NULL;
    v = mcp_extract_json("text/event-stream", sse, &err);
    CHECK(v != NULL, "sse parses");
    CHECK_CONT(v, "\"b\":2", "sse result");
    free(v);

    err = NULL;
    v = mcp_extract_json("text/event-stream", "event: message\n\n", &err);
    CHECK(v == NULL, "sse without data line fails");
    CHECK(err && strstr(err, "data:") != NULL, "sse error mentions data:");
    free(err);
}

static void test_connect_all_collects_warnings(void)
{
    const char *names[] = {"good", "bad"};
    const char *configs[] = {
        "{\"command\":\"/definitely/not/a/real/binary-llsh\"}",
        "{\"foo\":\"bar\"}",
    };
    strlist warnings;
    sl_init(&warnings);
    mcp_registry reg;
    mcp_connect_all(names, configs, 2, &reg, &warnings);
    CHECK(reg.tool_count == 0, "no tools from unreachable servers");
    CHECK(warnings.len == 2, "one warning per failed server");
    if (warnings.len == 2) {
        CHECK_CONT(warnings.items[0], "`good`", "good warning");
        CHECK_CONT(warnings.items[1], "`bad`", "bad warning");
        CHECK_CONT(warnings.items[1], "no `command` or `url`", "bad reason");
    }
    for (size_t i = 0; i < warnings.len; i++)
        free(warnings.items[i]);
    free(warnings.items);
    mcp_registry_free(&reg);
}

/* ------------------------------------------------------------------ */
/* read_line_fd (via mcp_read_line) -- buffered stdio reads            */
/* ------------------------------------------------------------------ */

/* Spawn `sh -c cmd` with stdout piped; returns the read end and the pid. */
static int spawn_sh(const char *cmd, pid_t *out_pid)
{
    int pfd[2];
    if (pipe(pfd) != 0)
        return -1;
    pid_t pid = fork();
    if (pid < 0) {
        close(pfd[0]);
        close(pfd[1]);
        return -1;
    }
    if (pid == 0) {
        dup2(pfd[1], STDOUT_FILENO);
        close(pfd[0]);
        close(pfd[1]);
        int dn = open("/dev/null", O_WRONLY);
        if (dn >= 0) {
            dup2(dn, STDERR_FILENO);
            close(dn);
        }
        setpgid(0, 0);
        execl("/bin/sh", "sh", "-c", cmd, (char *)NULL);
        _exit(127);
    }
    close(pfd[1]);
    if (out_pid)
        *out_pid = pid;
    return pfd[0];
}

static void kill_sh(pid_t pid)
{
    if (pid <= 0)
        return;
    kill(-pid, SIGKILL);
    int st;
    waitpid(pid, &st, 0);
}

/* A burst of lines arriving in one read(): the first line must come back
 * alone, with the rest carried for the following calls. */
static void test_read_line_burst(void)
{
    pid_t pid = 0;
    int fd = spawn_sh("printf 'a\\nb\\nc\\n'", &pid);
    CHECK(fd >= 0, "server spawned");
    char *err = NULL;
    mcp_client *c = mcp_connect_fds(-1, fd, &err);
    CHECK(c != NULL, "fd client created");
    char *l = mcp_read_line(c, "t", &err);
    CHECK(l && strcmp(l, "a") == 0, "first line of the burst");
    free(l);
    err = NULL;
    l = mcp_read_line(c, "t", &err);
    CHECK(l && strcmp(l, "b") == 0, "second line of the burst");
    free(l);
    err = NULL;
    l = mcp_read_line(c, "t", &err);
    CHECK(l && strcmp(l, "c") == 0, "third line of the burst");
    free(l);
    err = NULL;
    l = mcp_read_line(c, "t", &err);
    CHECK(l == NULL, "EOF after the burst");
    CHECK(err && strstr(err, "server closed") != NULL, "EOF error message");
    free(err);
    mcp_client_free(c);
    kill_sh(pid);
}

/* A line split across two read() calls must be reassembled from the carry
 * buffer. */
static void test_read_line_partial_across_reads(void)
{
    pid_t pid = 0;
    int fd = spawn_sh("printf 'hel'; sleep 0.5; printf 'lo\\nworld\\n'", &pid);
    CHECK(fd >= 0, "server spawned");
    char *err = NULL;
    mcp_client *c = mcp_connect_fds(-1, fd, &err);
    char *l = mcp_read_line(c, "t", &err);
    CHECK(l && strcmp(l, "hello") == 0, "line assembled across two reads");
    free(l);
    err = NULL;
    l = mcp_read_line(c, "t", &err);
    CHECK(l && strcmp(l, "world") == 0, "following line intact");
    free(l);
    mcp_client_free(c);
    kill_sh(pid);
}

/* Timeout with bytes in flight: the partial line is what we got. */
static void test_read_line_timeout_partial(void)
{
    setenv("LLSH_MCP_TIMEOUT", "1", 1);
    pid_t pid = 0;
    int fd = spawn_sh("printf 'partial'; sleep 5", &pid);
    CHECK(fd >= 0, "server spawned");
    char *err = NULL;
    mcp_client *c = mcp_connect_fds(-1, fd, &err);
    char *l = mcp_read_line(c, "t", &err);
    CHECK(l && strcmp(l, "partial") == 0, "partial line returned on timeout");
    free(l);
    mcp_client_free(c);
    kill_sh(pid);
    unsetenv("LLSH_MCP_TIMEOUT");
}

/* Timeout with a silent server: an error, not a hang or a fake line. */
static void test_read_line_timeout_empty(void)
{
    setenv("LLSH_MCP_TIMEOUT", "1", 1);
    pid_t pid = 0;
    int fd = spawn_sh("sleep 5", &pid);
    CHECK(fd >= 0, "server spawned");
    char *err = NULL;
    mcp_client *c = mcp_connect_fds(-1, fd, &err);
    char *l = mcp_read_line(c, "t", &err);
    CHECK(l == NULL, "no line when the server is silent");
    CHECK(err && strstr(err, "timed out") != NULL, "timeout error message");
    free(err);
    mcp_client_free(c);
    kill_sh(pid);
    unsetenv("LLSH_MCP_TIMEOUT");
}

static void test_stdio_line_size_limit(void)
{
    setenv("LLSH_MCP_TIMEOUT", "3", 1);
    pid_t pid = 0;
    int fd = spawn_sh("head -c 4194305 /dev/zero", &pid);
    CHECK(fd >= 0, "large-response server spawned");
    char *err = NULL;
    mcp_client *c = mcp_connect_fds(-1, fd, &err);
    char *line = mcp_read_line(c, "large", &err);
    CHECK(line == NULL, "oversized MCP response line rejected");
    CHECK(err && strstr(err, "4 MiB limit"), "size-limit error returned");
    free(line);
    free(err);
    mcp_client_free(c);
    kill_sh(pid);
    unsetenv("LLSH_MCP_TIMEOUT");
}

/* ------------------------------------------------------------------ */
/* mcp_connect_stdio -- real fork/pipe wiring                          */
/* ------------------------------------------------------------------ */

/* The child must be able to write to stdout and the parent to read it
 * back: the pipe ends were historically swapped, which made every stdio
 * server fail. A child that answers on its own (and stays alive so the
 * parent's write to its stdin does not hit EPIPE) exercises both ends. */
static void test_connect_stdio_roundtrip(void)
{
    mcp_server_config cfg;
    memset(&cfg, 0, sizeof cfg);
    cfg.command = xstrdup("/bin/sh");
    cfg.args = (char **)xmalloc(2 * sizeof *cfg.args);
    cfg.args[0] = xstrdup("-c");
    cfg.args[1] =
        xstrdup("printf '{\"jsonrpc\":\"2.0\",\"id\":1,\"result\":\"pong\"}"
                "\\n'; sleep 1");
    cfg.arg_count = 2;

    char *err = NULL;
    mcp_client *c = mcp_connect_stdio(&cfg, &err);
    CHECK(c != NULL, "stdio client spawned");
    if (!c) {
        if (err)
            free(err);
        mcp_server_config_free(&cfg);
        return;
    }
    char *res = mcp_request(c, "ping", "{}", &err);
    CHECK(res != NULL, "request round-trips to the child's stdout");
    if (res)
        CHECK(strstr(res, "pong") != NULL, "child response read back");
    free(res);
    if (err)
        free(err);
    mcp_client_free(c);
    mcp_server_config_free(&cfg);
}

static void test_connect_stdio_sends_newline_terminated_request(void)
{
    setenv("LLSH_MCP_TIMEOUT", "2", 1);
    mcp_server_config cfg;
    memset(&cfg, 0, sizeof cfg);
    cfg.command = xstrdup("/bin/sh");
    cfg.args = (char **)xmalloc(2 * sizeof *cfg.args);
    cfg.args[0] = xstrdup("-c");
    cfg.args[1] = xstrdup(
        "IFS= read -r req || exit 2; printf '{\"jsonrpc\":\"2.0\","
        "\"id\":1,\"result\":\"newline-ok\"}\\n'; sleep 1");
    cfg.arg_count = 2;
    char *err = NULL;
    mcp_client *c = mcp_connect_stdio(&cfg, &err);
    CHECK(c != NULL, "stdio client spawned");
    char *res = c ? mcp_request(c, "ping", "{}", &err) : NULL;
    CHECK(res && strstr(res, "newline-ok"),
          "server receives complete newline-framed request");
    free(res);
    free(err);
    mcp_client_free(c);
    mcp_server_config_free(&cfg);
    unsetenv("LLSH_MCP_TIMEOUT");
}

int main(void)
{
    /* Writes to a pipe whose server already died (e.g. a failed exec) must
     * surface as EPIPE, not kill the process with SIGPIPE. */
    signal(SIGPIPE, SIG_IGN);
    test_parses_stdio_config();
    test_parses_http_config();
    test_rejects_no_transport();
    test_handshake_discovery_call();
    test_server_error_is_surfaced();
    test_notification_lines_are_skipped();
    test_tool_invoke_roundtrip();
    test_tool_invoke_wraps_errors();
    test_extract_json_plain_and_sse();
    test_connect_all_collects_warnings();
    test_read_line_burst();
    test_read_line_partial_across_reads();
    test_read_line_timeout_partial();
    test_read_line_timeout_empty();
    test_stdio_line_size_limit();
    test_connect_stdio_roundtrip();
    test_connect_stdio_sends_newline_terminated_request();
    printf("%d checks, %d failures\n", checks, failures);
    return failures ? 1 : 0;
}