/* Sub-agents end to end against a local fake OpenAI-style SSE server: the
 * parent calls task twice (explore, in parallel), each child reads a file
 * and reports, the parent sees both reports. Also: tool restrictions in the
 * child, progress hooks, transcripts, roles from .agent/agents, depth
 * limit, unknown roles and cancellation. */
#define _POSIX_C_SOURCE 200809L
#ifndef _DEFAULT_SOURCE
#define _DEFAULT_SOURCE
#endif

#include "cJSON.h"
#include "cst_agent.h"
#include "cst_subagent.h"
#include "cst_tools.h"
#include "cst_util.h"
#include "os_compat.h"

#define CST_TEST_SOCKETS
#include "cst_test_util.h"

#include <ctype.h>
#include <curl/curl.h>
#include <pthread.h>
#include <stdatomic.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

static int nfail;

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

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

static atomic_int g_child_active, g_child_max, g_child_reqs;
static atomic_int g_child_saw_task, g_child_saw_write, g_child_role_prompt;
static atomic_int g_parent_saw_task;

static void send_all(test_sock_t fd, const char *s, size_t n) {
  while (n > 0) {
    long w = send(fd, s, (int)n, MSG_NOSIGNAL);
    if (w <= 0)
      return;
    s += w;
    n -= (size_t)w;
  }
}

/* Portable Content-Length lookup (avoid GNU strcasestr; undeclared under
 * strict C11 it is assumed to return int and the pointer truncates). */
static const char *find_content_length(const char *hdrs) {
  static const char key[] = "content-length:";
  for (const char *p = hdrs; *p; p++) {
    size_t i = 0;
    while (key[i] && tolower((unsigned char)p[i]) == (unsigned char)key[i])
      i++;
    if (key[i] == '\0')
      return p + (sizeof(key) - 1);
  }
  return NULL;
}

static void sse(test_sock_t fd, cJSON *delta, const char *finish) {
  cJSON *o = cJSON_CreateObject();
  cJSON *ch = cJSON_AddArrayToObject(o, "choices");
  cJSON *c0 = cJSON_CreateObject();
  cJSON_AddNumberToObject(c0, "index", 0);
  cJSON_AddItemToObject(c0, "delta", delta);
  if (finish)
    cJSON_AddStringToObject(c0, "finish_reason", finish);
  else
    cJSON_AddNullToObject(c0, "finish_reason");
  cJSON_AddItemToArray(ch, c0);
  if (finish) {
    cJSON *u = cJSON_AddObjectToObject(o, "usage");
    cJSON_AddNumberToObject(u, "prompt_tokens", 1200);
    cJSON_AddNumberToObject(u, "completion_tokens", 20);
  }
  char *txt = cJSON_PrintUnformatted(o);
  char *line = malloc(strlen(txt) + 16);
  int n = sprintf(line, "data: %s\n\n", txt);
  send_all(fd, line, (size_t)n);
  free(line);
  free(txt);
  cJSON_Delete(o);
}

static void reply_text(test_sock_t fd, const char *text) {
  cJSON *d = cJSON_CreateObject();
  cJSON_AddStringToObject(d, "role", "assistant");
  cJSON_AddStringToObject(d, "content", text);
  sse(fd, d, NULL);
  sse(fd, cJSON_CreateObject(), "stop");
  send_all(fd, "data: [DONE]\n\n", 14);
}

/* calls: {name, args_json} pairs */
static void reply_calls(test_sock_t fd, const char *const *calls, int n) {
  for (int i = 0; i < n; i++) {
    cJSON *d = cJSON_CreateObject();
    cJSON *tc = cJSON_AddArrayToObject(d, "tool_calls");
    cJSON *c = cJSON_CreateObject();
    cJSON_AddNumberToObject(c, "index", i);
    char id[32];
    snprintf(id, sizeof(id), "call_%d", i);
    cJSON_AddStringToObject(c, "id", id);
    cJSON_AddStringToObject(c, "type", "function");
    cJSON *f = cJSON_AddObjectToObject(c, "function");
    cJSON_AddStringToObject(f, "name", calls[2 * i]);
    cJSON_AddStringToObject(f, "arguments", calls[2 * i + 1]);
    cJSON_AddItemToArray(tc, c);
    sse(fd, d, NULL);
  }
  sse(fd, cJSON_CreateObject(), "tool_calls");
  send_all(fd, "data: [DONE]\n\n", 14);
}

static int has_tool(const cJSON *body, const char *name) {
  const cJSON *t = NULL;
  cJSON_ArrayForEach(t, cJSON_GetObjectItemCaseSensitive(body, "tools")) {
    const cJSON *f = cJSON_GetObjectItemCaseSensitive(t, "function");
    const cJSON *n = cJSON_GetObjectItemCaseSensitive(f, "name");
    if (cJSON_IsString(n) && strcmp(n->valuestring, name) == 0)
      return 1;
  }
  return 0;
}

static void handle(test_sock_t fd, const char *body_txt) {
  cJSON *body = cJSON_Parse(body_txt);
  if (!body) {
    const char *err = "HTTP/1.1 400 Bad Request\r\nConnection: close\r\n\r\n";
    send_all(fd, err, strlen(err));
    return;
  }
  const cJSON *msgs = cJSON_GetObjectItemCaseSensitive(body, "messages");
  int nm = cJSON_IsArray(msgs) ? cJSON_GetArraySize(msgs) : 0;
  const cJSON *sys = NULL;
  const char *last_role = "";
  const char *user = "";
  if (nm > 0) {
    sys = cJSON_GetObjectItemCaseSensitive(cJSON_GetArrayItem(msgs, 0),
                                           "content");
    const cJSON *last = cJSON_GetArrayItem(msgs, nm - 1);
    const cJSON *lr = cJSON_GetObjectItemCaseSensitive(last, "role");
    if (cJSON_IsString(lr))
      last_role = lr->valuestring;
  }
  for (int i = 0; i < nm; i++) {
    const cJSON *m = cJSON_GetArrayItem(msgs, i);
    const cJSON *r = cJSON_GetObjectItemCaseSensitive(m, "role");
    const cJSON *c = cJSON_GetObjectItemCaseSensitive(m, "content");
    if (cJSON_IsString(r) && strcmp(r->valuestring, "user") == 0 &&
        cJSON_IsString(c))
      user = c->valuestring;
  }
  const char *hdr = "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\n"
                    "Connection: close\r\n\r\n";
  send_all(fd, hdr, strlen(hdr));
  int child = cJSON_IsString(sys) && strstr(sys->valuestring, "SUB-AGENT");
  if (child) {
    atomic_fetch_add(&g_child_reqs, 1);
    int now = atomic_fetch_add(&g_child_active, 1) + 1;
    int mx = atomic_load(&g_child_max);
    while (now > mx && !atomic_compare_exchange_weak(&g_child_max, &mx, now))
      ;
    if (has_tool(body, "task"))
      atomic_store(&g_child_saw_task, 1);
    if (has_tool(body, "write_file") || has_tool(body, "run_shell"))
      atomic_store(&g_child_saw_write, 1);
    if (strstr(sys->valuestring, "Role instructions:\nYou only read."))
      atomic_store(&g_child_role_prompt, 1);
    if (strstr(user, "Slow")) {
      test_sleep_ms(4000); /* cancelled before this ends */
    } else {
      test_sleep_ms(300);
    }
    if (strcmp(last_role, "user") == 0) {
      const char *calls[] = {"read_file", "{\"path\":\"a.txt\"}"};
      reply_calls(fd, calls, 1);
    } else {
      /* last char of the user prompt (e.g. "Find A" -> "A"); empty-safe */
      char tag = user[0] ? user[strlen(user) - 1] : '?';
      char out[64];
      snprintf(out, sizeof(out), "child report %c", tag);
      reply_text(fd, out);
    }
    atomic_fetch_sub(&g_child_active, 1);
  } else {
    if (has_tool(body, "task"))
      atomic_store(&g_parent_saw_task, 1);
    if (strcmp(last_role, "user") == 0 && strstr(user, "cancel test")) {
      const char *calls[] = {
          "task", "{\"description\":\"slow one\",\"prompt\":\"Slow\","
                  "\"agent\":\"explore\"}"};
      reply_calls(fd, calls, 1);
    } else if (strcmp(last_role, "user") == 0) {
      const char *calls[] = {
          "task",
          "{\"description\":\"find A\",\"prompt\":\"Find A\","
          "\"agent\":\"explore\"}",
          "task",
          "{\"description\":\"find B\",\"prompt\":\"Find B\","
          "\"agent\":\"explore\"}"};
      reply_calls(fd, calls, 2);
    } else {
      cst_buf_t b = {0};
      cst_buf_puts(&b, "parent done:");
      for (int i = 0; i < nm; i++) {
        const cJSON *m = cJSON_GetArrayItem(msgs, i);
        const cJSON *r = cJSON_GetObjectItemCaseSensitive(m, "role");
        const cJSON *c = cJSON_GetObjectItemCaseSensitive(m, "content");
        const char *rep = cJSON_IsString(c) ? strstr(c->valuestring,
                                                     "child report ")
                                            : NULL;
        if (cJSON_IsString(r) && strcmp(r->valuestring, "tool") == 0 && rep)
          cst_buf_printf(&b, " [%.14s]", rep);
      }
      reply_text(fd, b.buf);
      cst_buf_free(&b);
    }
  }
  cJSON_Delete(body);
}

static void *conn_main(void *p) {
  test_sock_t fd = (test_sock_t)(intptr_t)p;
  cst_buf_t req = {0};
  char tmp[8192];
  long need = -1;
  size_t hdr_end = 0;
  for (;;) {
    long n = recv(fd, tmp, sizeof(tmp), 0);
    if (n <= 0)
      break;
    cst_buf_append(&req, tmp, (size_t)n);
    if (!hdr_end) {
      char *e = strstr(req.buf, "\r\n\r\n");
      if (e) {
        hdr_end = (size_t)(e - req.buf) + 4;
        const char *cl = find_content_length(req.buf);
        need = cl ? atol(cl) : 0;
      }
    }
    if (hdr_end && (long)(req.len - hdr_end) >= need)
      break;
  }
  if (hdr_end)
    handle(fd, req.buf + hdr_end);
  cst_buf_free(&req);
  test_sock_close(fd);
  return NULL;
}

static test_sock_t g_lfd;

static void *server_main(void *p) {
  (void)p;
  for (;;) {
    test_sock_t fd = accept(g_lfd, NULL, NULL);
    if (fd == TEST_BAD_SOCK)
      return NULL;
    pthread_t th;
    if (pthread_create(&th, NULL, conn_main, (void *)(intptr_t)fd) == 0)
      pthread_detach(th);
    else
      test_sock_close(fd);
  }
}

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

static atomic_int g_progress;
static char g_progress_last[256];

static void on_tool_event(const char *phase, const ctl_tool_call_t *call,
                          const char *detail, void *ud) {
  (void)call;
  (void)ud;
  if (strcmp(phase, "progress") == 0 && detail) {
    atomic_fetch_add(&g_progress, 1);
    if (strstr(detail, "read_file"))
      snprintf(g_progress_last, sizeof(g_progress_last), "%s", detail);
  }
}

static const char *last_assistant(cst_agent_t *a) {
  cst_session_t *s = cst_agent_session(a);
  for (int i = cst_session_count(s) - 1; i >= 0; i--) {
    const cst_message_t *m = cst_session_get(s, i);
    if (m->role == CST_ROLE_ASSISTANT && m->content && m->content[0])
      return m->content;
  }
  return "";
}

static void put(const char *root, const char *rel, const char *text) {
  char p[1024];
  snprintf(p, sizeof(p), "%s/%s", root, rel);
  for (char *q = p + strlen(root) + 1; *q; q++)
    if (*q == '/') {
      *q = 0;
      test_mkdir(p, 0755);
      *q = '/';
    }
  FILE *fp = fopen(p, "wb");
  if (fp) {
    fputs(text, fp);
    fclose(fp);
  }
}

static void *cancel_later(void *p) {
  test_sleep_ms(800);
  cst_agent_request_cancel(p);
  return NULL;
}

int main(void) {
  test_sock_init();
  curl_global_init(CURL_GLOBAL_DEFAULT);
  g_lfd = socket(AF_INET, SOCK_STREAM, 0);
  struct sockaddr_in addr;
  memset(&addr, 0, sizeof(addr));
  addr.sin_family = AF_INET;
  addr.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
  socklen_t al = sizeof(addr);
  if (bind(g_lfd, (struct sockaddr *)&addr, sizeof(addr)) != 0 ||
      listen(g_lfd, 16) != 0 ||
      getsockname(g_lfd, (struct sockaddr *)&addr, &al) != 0) {
    perror("listen");
    return 1;
  }
  pthread_t srv;
  pthread_create(&srv, NULL, server_main, NULL);

  char root[600], ws[700];
  if (test_tmpdir(root, sizeof(root), "cst-sub") != 0)
    return 1;
  snprintf(ws, sizeof(ws), "%s/ws", root);
  setenv("XDG_CONFIG_HOME", root, 1);
  put(root, "ws/a.txt", "alpha\n");
  put(root, "ws/.agent/agents/reviewer.md",
      "---\nname: reviewer\ndescription: Reviews diffs.\ntools: read_file, "
      "grep\nreadonly: true\nmax_steps: 4\n---\nReview carefully.\n");
  put(root, "ws/.agent/agents/explore.md",
      "---\ndescription: Custom explorer.\nreadonly: yes\n---\nYou only "
      "read. Custom.\n");
  put(root, "ws/.agent/sessions/.keep", "");

  /* roles */
  cst_agent_role_t roles[CST_ROLES_MAX];
  int nr = cst_roles_list(ws, roles, CST_ROLES_MAX);
  CHECK(nr == 4, "three built-ins + reviewer (explore overridden)");
  CHECK(nr == 4 && strcmp(roles[0].name, "general") == 0 &&
            strcmp(roles[3].name, "reviewer") == 0 &&
            strcmp(roles[3].tools, ",read_file,grep,") == 0 &&
            roles[3].read_only && roles[3].max_steps == 4 &&
            strcmp(roles[3].prompt, "Review carefully.\n") == 0,
        "custom role parsed");
  CHECK(nr == 4 && strcmp(roles[1].name, "explore") == 0 &&
            strcmp(roles[1].source, "workspace") == 0 &&
            strcmp(roles[1].desc, "Custom explorer.") == 0,
        "workspace overrides a built-in");
  cst_roles_free(roles, nr);
  char *ix = cst_subagent_index(ws);
  CHECK(ix && strstr(ix, "- reviewer (read-only): Reviews diffs.") &&
            strstr(ix, "- general: "),
        "roles index");
  free(ix);

  CHECK(cst_subagent_register_tool(ws) == 0 && cst_tool_find("task"),
        "task registered");
  const cst_tool_def_t *td = cst_tool_find("task");
  ctl_tool_call_t pc;
  memset(&pc, 0, sizeof(pc));
  snprintf(pc.name, sizeof(pc.name), "task");
  pc.nargs = 1;
  snprintf(pc.args[0].key, sizeof(pc.args[0].key), "agent");
  pc.args[0].value = "explore";
  CHECK(td->parallel && td->parallel(&pc, td->ud), "explore runs in parallel");
  pc.args[0].value = "general";
  CHECK(!td->parallel(&pc, td->ud), "general does not");
  pc.args[0].value = "reviewer";
  CHECK(td->parallel(&pc, td->ud), "read-only custom role does");

  /* parent agent */
  cst_agent_config_t cfg;
  memset(&cfg, 0, sizeof(cfg));
  snprintf(cfg.workspace_root, sizeof(cfg.workspace_root), "%s", ws);
  snprintf(cfg.session_path, sizeof(cfg.session_path),
           "%s/.agent/sessions/parent1.jsonl", ws);
  snprintf(cfg.api_base, sizeof(cfg.api_base), "http://127.0.0.1:%d",
           ntohs(addr.sin_port));
  snprintf(cfg.model, sizeof(cfg.model), "m");
  cfg.protocol = CST_PROTO_NATIVE_EOS;
  cfg.max_steps = 6;
  cfg.auto_approve = 1;
  cst_agent_t *a = cst_agent_create(&cfg);
  CHECK(a != NULL, "parent created");
  cst_agent_hooks_t h;
  memset(&h, 0, sizeof(h));
  h.on_tool_event = on_tool_event;
  char err[512];
  double t0 = os_now_ms();
  int rc = cst_agent_run_turn_ex(a, "go", &h, err, sizeof(err));
  double ms = os_now_ms() - t0;
  CHECK(rc == 0, err);
  const char *fin = last_assistant(a);
  CHECK(strstr(fin, "[child report A") && strstr(fin, "[child report B"),
        "parent got both reports");
  CHECK(atomic_load(&g_parent_saw_task), "parent offered task");
  CHECK(!atomic_load(&g_child_saw_task), "child not offered task");
  CHECK(!atomic_load(&g_child_saw_write), "explore child has no write/shell");
  CHECK(atomic_load(&g_child_role_prompt), "role instructions in the prompt");
  CHECK(atomic_load(&g_child_max) == 2, "two children ran at once");
  CHECK(atomic_load(&g_child_reqs) == 4, "two model calls per child");
  CHECK(ms < 2400, "parallel children overlap in time");
  CHECK(atomic_load(&g_progress) >= 4 &&
            strstr(g_progress_last, "explore · read_file a.txt"),
        "progress events");

  cst_subagent_rec_t recs[8];
  int nrec = cst_subagent_list(ws, "parent1", recs, 8);
  CHECK(nrec == 2 && strcmp(recs[0].file, "sub-2") == 0 &&
            strcmp(recs[0].status, "done") == 0 && recs[0].steps == 2 &&
            recs[0].tools == 1 && strcmp(recs[0].role, "explore") == 0,
        "transcript records");
  char tp[1024];
  CHECK(cst_subagent_path(ws, "parent1", "sub-1", tp, sizeof(tp)) == 0 &&
            access(tp, F_OK) == 0,
        "transcript file");
  char ids[8][128];
  int nids = cst_agent_store_list(ws, ids, 8);
  CHECK(nids == 1 && strcmp(ids[0], "parent1") == 0,
        "sub-agent transcripts are not sessions");

  /* direct calls: unknown role, depth limit */
  cst_tool_ctx_t ctx;
  memset(&ctx, 0, sizeof(ctx));
  ctx.workspace_root = ws;
  ctx.agent = a;
  ctl_tool_call_t c;
  memset(&c, 0, sizeof(c));
  snprintf(c.name, sizeof(c.name), "task");
  c.nargs = 3;
  snprintf(c.args[0].key, sizeof(c.args[0].key), "description");
  c.args[0].value = "x";
  snprintf(c.args[1].key, sizeof(c.args[1].key), "prompt");
  c.args[1].value = "y";
  snprintf(c.args[2].key, sizeof(c.args[2].key), "agent");
  c.args[2].value = "nobody";
  for (int i = 0; i < 3; i++)
    c.args[i].value_len = strlen(c.args[i].value);
  cst_tool_result_t r;
  memset(&r, 0, sizeof(r));
  cst_tool_run(&ctx, &c, &r);
  CHECK(r.exit_code == 1 && r.out &&
            strstr(r.out, "available: general, explore, plan, reviewer"),
        "unknown role lists the roles");
  cst_tool_result_free(&r);
  cst_agent_config_mut(a)->depth = 1;
  c.args[2].value = "explore";
  c.args[2].value_len = 7;
  memset(&r, 0, sizeof(r));
  cst_tool_run(&ctx, &c, &r);
  CHECK(r.exit_code == 1 && r.out && strstr(r.out, "cannot start sub-agents"),
        "depth limit");
  cst_tool_result_free(&r);
  CHECK(!cst_tool_active(cst_agent_config(a), td), "task hidden at depth 1");
  cst_agent_config_mut(a)->depth = 0;

  /* cancelling the parent stops the child */
  pthread_t cth;
  pthread_create(&cth, NULL, cancel_later, a);
  t0 = os_now_ms();
  rc = cst_agent_run_turn_ex(a, "cancel test", &h, err, sizeof(err));
  ms = os_now_ms() - t0;
  pthread_join(cth, NULL);
  CHECK(rc != 0 && ms < 3500, "cancel reaches the child");
  nrec = cst_subagent_list(ws, "parent1", recs, 8);
  CHECK(nrec == 3 && strcmp(recs[0].status, "cancelled") == 0,
        "cancelled run recorded");

  cst_agent_destroy(a);
  cst_tool_unregister_owner("builtin:task");
  CHECK(cst_agent_store_delete(ws, "parent1") == 0, "delete session");
  CHECK(cst_subagent_list(ws, "parent1", recs, 8) == 0,
        "delete removes sub-agent transcripts");

  if (test_rmrf(root) != 0)
    fprintf(stderr, "note: cleanup failed\n");
  test_sock_close(g_lfd);
  if (nfail) {
    fprintf(stderr, "test_subagent: %d failure(s)\n", nfail);
    return 1;
  }
  printf("test_subagent: OK\n");
  return 0;
}