/* MCP client: the test binary doubles as a stdio MCP server (--server).
 * Covers the handshake, paginated tools/list, registration, tools/call
 * content mapping, isError, server->client requests, stdout noise,
 * cancellation, a crashing server and bad configs. */
#define _POSIX_C_SOURCE 200809L

#include "cJSON.h"
#include "cst_agent.h"
#include "cst_attach.h"
#include "cst_mcp.h"
#include "cst_tools.h"
#include "os_compat.h"

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

#include "cst_test_util.h"

static int nfail;
static char g_tmp[OS_PATH_MAX];
static cst_agent_t *g_agent;

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

/* ---------- fake server ---------- */

static void reply(const cJSON *id, cJSON *result) {
  cJSON *m = cJSON_CreateObject();
  cJSON_AddStringToObject(m, "jsonrpc", "2.0");
  cJSON_AddItemToObject(m, "id", cJSON_Duplicate(id, 1));
  cJSON_AddItemToObject(m, "result", result);
  char *s = cJSON_PrintUnformatted(m);
  printf("%s\n", s);
  fflush(stdout);
  free(s);
  cJSON_Delete(m);
}

static cJSON *tool(const char *name, const char *desc) {
  cJSON *t = cJSON_CreateObject();
  cJSON_AddStringToObject(t, "name", name);
  cJSON_AddStringToObject(t, "description", desc);
  cJSON *s = cJSON_AddObjectToObject(t, "inputSchema");
  cJSON_AddStringToObject(s, "type", "object");
  cJSON *p = cJSON_AddObjectToObject(s, "properties");
  cJSON *n = cJSON_AddObjectToObject(p, "n");
  cJSON_AddStringToObject(n, "type", "integer");
  cJSON *x = cJSON_AddObjectToObject(p, "text");
  cJSON_AddStringToObject(x, "type", "string");
  return t;
}

static cJSON *text_result(const char *text, int is_error) {
  cJSON *r = cJSON_CreateObject();
  cJSON *c = cJSON_AddArrayToObject(r, "content");
  cJSON *it = cJSON_CreateObject();
  cJSON_AddStringToObject(it, "type", "text");
  cJSON_AddStringToObject(it, "text", text);
  cJSON_AddItemToArray(c, it);
  if (is_error)
    cJSON_AddTrueToObject(r, "isError");
  return r;
}

static int server(void) {
  fprintf(stderr, "fake server starting\n");
  printf("this line is not JSON-RPC\n");
  fflush(stdout);
  char line[65536];
  while (fgets(line, sizeof(line), stdin)) {
    cJSON *m = cJSON_Parse(line);
    const cJSON *id = cJSON_GetObjectItemCaseSensitive(m, "id");
    const cJSON *method = cJSON_GetObjectItemCaseSensitive(m, "method");
    const char *mt = cJSON_IsString(method) ? method->valuestring : "";
    const cJSON *params = cJSON_GetObjectItemCaseSensitive(m, "params");
    if (strcmp(mt, "initialize") == 0) {
      cJSON *r = cJSON_CreateObject();
      cJSON_AddStringToObject(r, "protocolVersion", "2025-06-18");
      cJSON_AddObjectToObject(cJSON_AddObjectToObject(r, "capabilities"),
                              "tools");
      cJSON *si = cJSON_AddObjectToObject(r, "serverInfo");
      cJSON_AddStringToObject(si, "name", "fake");
      reply(id, r);
    } else if (strcmp(mt, "tools/list") == 0) {
      const cJSON *cur = cJSON_GetObjectItemCaseSensitive(params, "cursor");
      cJSON *r = cJSON_CreateObject();
      cJSON *tools = cJSON_AddArrayToObject(r, "tools");
      if (!cur) {
        cJSON_AddItemToArray(tools, tool("echo", "Echo the arguments."));
        cJSON_AddItemToArray(tools, tool("odd.name/x", "Sanitized name."));
        cJSON_AddStringToObject(r, "nextCursor", "page2");
      } else {
        cJSON_AddItemToArray(tools, tool("fail", "Always fails."));
        cJSON_AddItemToArray(tools, tool("mixed", "Mixed content."));
        cJSON_AddItemToArray(tools, tool("hang", "Never answers."));
        cJSON_AddItemToArray(tools, tool("crash", "Exits."));
        cJSON_AddItemToArray(tools, tool("pic", "Returns images."));
      }
      reply(id, r);
    } else if (strcmp(mt, "tools/call") == 0) {
      const char *name = cJSON_GetObjectItemCaseSensitive(params, "name")
                             ->valuestring;
      const cJSON *args = cJSON_GetObjectItemCaseSensitive(params, "arguments");
      if (strcmp(name, "echo") == 0) {
        /* a server->client request first; the client must answer it */
        printf("{\"jsonrpc\":\"2.0\",\"id\":\"s1\",\"method\":\"ping\"}\n");
        fflush(stdout);
        char pong[4096];
        int ok = fgets(pong, sizeof(pong), stdin) && strstr(pong, "\"s1\"") &&
                 strstr(pong, "\"result\"");
        char *a = cJSON_PrintUnformatted(args);
        char out[8192];
        const char *env = getenv("FAKE_ENV");
        snprintf(out, sizeof(out), "args=%s ping=%s env=%s", a,
                 ok ? "ok" : "missing", env ? env : "");
        free(a);
        reply(id, text_result(out, 0));
      } else if (strcmp(name, "fail") == 0) {
        reply(id, text_result("bad input", 1));
      } else if (strcmp(name, "mixed") == 0) {
        cJSON *r = text_result("first", 0);
        cJSON *c = cJSON_GetObjectItemCaseSensitive(r, "content");
        cJSON *img = cJSON_CreateObject();
        cJSON_AddStringToObject(img, "type", "image");
        cJSON_AddStringToObject(img, "mimeType", "image/png");
        cJSON_AddStringToObject(img, "data", "AAAA");
        cJSON_AddItemToArray(c, img);
        cJSON *res = cJSON_CreateObject();
        cJSON_AddStringToObject(res, "type", "resource");
        cJSON *rr = cJSON_AddObjectToObject(res, "resource");
        cJSON_AddStringToObject(rr, "uri", "file:///x");
        cJSON_AddStringToObject(rr, "text", "resource text");
        cJSON_AddItemToArray(c, res);
        reply(id, r);
      } else if (strcmp(name, "pic") == 0) {
        cJSON *r = text_result("a chart", 0);
        cJSON *c = cJSON_GetObjectItemCaseSensitive(r, "content");
        const char *datas[] = {
            "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAC"
            "hwGA60e6kgAAAABJRU5ErkJggg==",
            "AAAA"};
        for (int i = 0; i < 2; i++) {
          cJSON *img = cJSON_CreateObject();
          cJSON_AddStringToObject(img, "type", "image");
          cJSON_AddStringToObject(img, "mimeType", "image/png");
          cJSON_AddStringToObject(img, "data", datas[i]);
          cJSON_AddItemToArray(c, img);
        }
        reply(id, r);
      } else if (strcmp(name, "crash") == 0) {
        fprintf(stderr, "fatal: fake crash\n");
        exit(3);
      }
      /* hang: no reply */
    }
    cJSON_Delete(m);
  }
  return 0;
}

/* ---------- client ---------- */

static int call(const char *tool, const char *k, const char *v,
                cst_tool_result_t *r) {
  ctl_tool_call_t c;
  memset(&c, 0, sizeof(c));
  snprintf(c.name, sizeof(c.name), "%s", tool);
  if (k) {
    snprintf(c.args[0].key, sizeof(c.args[0].key), "%s", k);
    c.args[0].value = (char *)v;
    c.args[0].value_len = strlen(v);
    c.nargs = 1;
  }
  cst_tool_ctx_t ctx;
  memset(&ctx, 0, sizeof(ctx));
  ctx.workspace_root = g_tmp;
  ctx.agent = g_agent;
  memset(r, 0, sizeof(*r));
  return cst_tool_run(&ctx, &c, r);
}

static const cst_mcp_info_t *info_of(const char *name, cst_mcp_info_t *buf,
                                     int n) {
  for (int i = 0; i < n; i++)
    if (strcmp(buf[i].name, name) == 0)
      return &buf[i];
  return NULL;
}

int main(int argc, char **argv) {
  if (argc > 1 && strcmp(argv[1], "--server") == 0)
    return server();

  char self[OS_PATH_MAX];
  if (os_realpath(argv[0], self, sizeof(self)) != 0) {
    fprintf(stderr, "cannot resolve %s\n", argv[0]);
    return 1;
  }
  if (os_temp_dir(g_tmp, sizeof(g_tmp)) != 0)
    snprintf(g_tmp, sizeof(g_tmp), ".");
  setenv("CST_TEST_SECRET", "s3cret", 1);
  char cfg[OS_PATH_MAX + 128];
  snprintf(cfg, sizeof(cfg),
           "{\"command\": \"%s\", \"args\": [\"--server\"], "
           "\"env\": {\"FAKE_ENV\": \"v=${CST_TEST_SECRET}\"}, "
           "\"autoApprove\": [\"echo\"]}",
           self);
  CHECK(cst_mcp_add_config_json("fake", cfg, NULL, "user") == 0, "add fake");
  CHECK(cst_mcp_add_config_json("fake", cfg, NULL, "user") != 0,
        "duplicate name rejected");
  CHECK(cst_mcp_add_config_json("remote", "{\"url\": \"https://x\"}", NULL,
                                "user") == 0,
        "url server recorded");
  CHECK(cst_mcp_add_config_json("ghost",
                                "{\"command\": \"/nonexistent/cst-mcp\"}",
                                NULL, "user") == 0,
        "ghost server recorded");
  CHECK(cst_mcp_add_config_json("off", "{\"command\": \"true\", "
                                       "\"disabled\": true}",
                                NULL, "user") == 0,
        "disabled server recorded");
  int ready = -1, total = -1;
  cst_mcp_counts(&ready, &total);
  CHECK(ready == 0 && total == 3, "counts before start (disabled excluded)");

  cst_mcp_start_all(g_tmp);
  cst_mcp_wait_ready(10000);

  cst_mcp_info_t mi[CST_MCP_MAX];
  int n = cst_mcp_info(mi, CST_MCP_MAX);
  const cst_mcp_info_t *f = info_of("fake", mi, n);
  CHECK(f && f->state == CST_MCP_READY && f->ntools == 7,
        "fake ready with 7 tools over two pages");
  const cst_mcp_info_t *rm = info_of("remote", mi, n);
  CHECK(rm && rm->state == CST_MCP_FAILED && strstr(rm->error, "stdio"),
        "url servers unsupported");
  const cst_mcp_info_t *gh = info_of("ghost", mi, n);
  CHECK(gh && gh->state != CST_MCP_READY && gh->error[0],
        "missing command fails");
  const cst_mcp_info_t *of = info_of("off", mi, n);
  CHECK(of && of->state == CST_MCP_OFF, "disabled stays off");
  cst_mcp_counts(&ready, &total);
  CHECK(ready == 1 && total == 3, "counts after start");

  const cst_tool_def_t *d = cst_tool_find("mcp__fake__echo");
  CHECK(d && d->perm == CST_TP_READ && d->schema_json &&
            strstr(d->schema_json, "\"integer\""),
        "echo registered, autoApprove -> read, schema kept");
  d = cst_tool_find("mcp__fake__fail");
  CHECK(d && d->perm == CST_TP_NET, "others need approval (net)");
  CHECK(cst_tool_find("mcp__fake__odd_name_x"), "tool name sanitized");

  cst_tool_result_t r;
  call("mcp__fake__echo", "n", "42", &r);
  CHECK(r.exit_code == 0 && r.out && strstr(r.out, "args={\"n\":42}"),
        "typed arguments");
  CHECK(r.out && strstr(r.out, "ping=ok"), "server ping answered");
  CHECK(r.out && strstr(r.out, "env=v=s3cret"), "env with ${VAR} expansion");
  cst_tool_result_free(&r);

  call("mcp__fake__fail", NULL, NULL, &r);
  CHECK(r.exit_code == 1 && r.out && strstr(r.out, "bad input"),
        "isError -> exit 1");
  cst_tool_result_free(&r);

  call("mcp__fake__mixed", NULL, NULL, &r);
  CHECK(r.exit_code == 0 && r.out && strstr(r.out, "first\n[image image/png") &&
            strstr(r.out, "resource text"),
        "content items mapped to text");
  CHECK(r.nattach == 0, "no session: images stay text");
  cst_tool_result_free(&r);

  cst_agent_config_t acfg;
  memset(&acfg, 0, sizeof(acfg));
  snprintf(acfg.workspace_root, sizeof(acfg.workspace_root), "%s", g_tmp);
  char sdir[1024];
  snprintf(sdir, sizeof(sdir), "%.900s/cst-mcp-test-%d", g_tmp, (int)getpid());
  os_mkdir_p(sdir, 0755);
  snprintf(acfg.session_path, sizeof(acfg.session_path), "%s/s.jsonl", sdir);
  g_agent = cst_agent_create(&acfg);
  CHECK(g_agent != NULL, "agent for the session store");
  call("mcp__fake__pic", NULL, NULL, &r);
  CHECK(r.exit_code == 0 && r.nattach == 1 && r.attach &&
            strcmp(r.attach[0].mime, "image/png") == 0 &&
            strcmp(r.attach[0].origin, "mcp") == 0 &&
            r.attach[0].width == 1 && r.attach[0].height == 1,
        "image content stored and attached");
  CHECK(r.out && strstr(r.out, "a chart\n[image ") &&
            strstr(r.out, "shown to you as an image") &&
            strstr(r.out, "not shown:"),
        "one line per image, bad data falls back to text");
  if (r.nattach == 1) {
    char full[OS_PATH_MAX];
    CHECK(cst_attach_full_path(acfg.session_path, &r.attach[0], full,
                               sizeof(full)) == 0 &&
              access(full, R_OK) == 0,
          "stored file exists");
  }
  cst_tool_result_free(&r);
  cst_agent_destroy(g_agent);
  g_agent = NULL;
  test_rmrf(sdir);

  volatile int cancel = 1;
  os_set_cancel_flag(&cancel);
  double t0 = os_now_ms();
  call("mcp__fake__hang", NULL, NULL, &r);
  CHECK(r.exit_code == 1 && r.out && strstr(r.out, "cancelled") &&
            os_now_ms() - t0 < 5000,
        "cancel stops waiting");
  cst_tool_result_free(&r);
  os_set_cancel_flag(NULL);

  call("mcp__fake__crash", NULL, NULL, &r);
  CHECK(r.exit_code == 1 && r.out && strstr(r.out, "fake crash"),
        "crash reported with the stderr tail");
  cst_tool_result_free(&r);
  n = cst_mcp_info(mi, CST_MCP_MAX);
  f = info_of("fake", mi, n);
  CHECK(f && f->state == CST_MCP_EXITED, "crashed server marked exited");
  call("mcp__fake__echo", "n", "1", &r);
  CHECK(r.exit_code == 1 && r.out && strstr(r.out, "not running"),
        "calls after a crash fail fast");
  cst_tool_result_free(&r);

  cst_mcp_stop_all();
  CHECK(!cst_tool_find("mcp__fake__echo"), "stop unregisters tools");
  cst_mcp_counts(&ready, &total);
  CHECK(total == 0, "stop forgets configs");

  /* a live server stops cleanly too */
  CHECK(cst_mcp_add_config_json("fake2", cfg, NULL, "user") == 0, "re-add");
  cst_mcp_start_all(g_tmp);
  cst_mcp_wait_ready(10000);
  CHECK(cst_tool_find("mcp__fake2__echo"), "second server ready");
  t0 = os_now_ms();
  cst_mcp_stop_all();
  CHECK(os_now_ms() - t0 < 3000, "stop is prompt");

  if (nfail) {
    fprintf(stderr, "test_mcp: %d failure(s)\n", nfail);
    return 1;
  }
  printf("test_mcp: OK\n");
  return 0;
}