#define _POSIX_C_SOURCE 200809L

#include "cJSON.h"
#include "cst_agent.h"

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

/* Drives cst_llm_chat_stream through CST_REPLAY_SSE with synthetic and
 * recorded (tests/fixtures/sse) streams, at several network chunk sizes. */

static int nfail;

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

#define MAX_CALLS 16

typedef struct {
  char text[65536];
  size_t text_len;
  int ncalls;
  char ids[MAX_CALLS][64];
  char names[MAX_CALLS][64];
  char *args[MAX_CALLS]; /* NULL = invalid */
  int args_null[MAX_CALLS];
  int nopen;
  int open_before_first_call; /* opens seen before the first call fired */
} capture_t;

static void on_delta(const char *d, size_t n, void *ud) {
  capture_t *c = ud;
  if (c->text_len + n < sizeof(c->text)) {
    memcpy(c->text + c->text_len, d, n);
    c->text_len += n;
    c->text[c->text_len] = 0;
  }
}

static void on_tool(const char *id, const char *name, const char *args,
                    void *ud) {
  capture_t *c = ud;
  if (c->ncalls >= MAX_CALLS)
    return;
  int i = c->ncalls++;
  snprintf(c->ids[i], sizeof(c->ids[i]), "%s", id ? id : "");
  snprintf(c->names[i], sizeof(c->names[i]), "%s", name ? name : "");
  c->args[i] = args ? strdup(args) : NULL;
  c->args_null[i] = args == NULL;
  if (i == 0)
    c->open_before_first_call = c->nopen;
}

static void on_open(int index, const char *id, const char *name, void *ud) {
  (void)index;
  (void)id;
  (void)name;
  capture_t *c = ud;
  c->nopen++;
}

static void capture_free(capture_t *c) {
  for (int i = 0; i < c->ncalls; i++)
    free(c->args[i]);
}

static char g_dir[256];

static void write_fixture(const char *sse) {
  char path[512];
  snprintf(path, sizeof(path), "%s/0000.sse", g_dir);
  FILE *fp = fopen(path, "wb");
  fputs(sse, fp);
  fclose(fp);
}

static int replay(const char *dir, int seq, const char *chunk, capture_t *c,
                  cst_llm_stats_t *st) {
  memset(c, 0, sizeof(*c));
  memset(st, 0, sizeof(*st));
  setenv("CST_REPLAY_SSE", dir, 1);
  setenv("CST_REPLAY_CHUNK", chunk, 1);
  cst_llm_trace_reset();
  for (int i = 0; i < seq; i++) {
    /* advance the sequence counter to the wanted file */
    cst_llm_request_t skip;
    memset(&skip, 0, sizeof(skip));
    skip.api_base = "http://replay";
    skip.model = "m";
    char e[64];
    capture_t tmp;
    memset(&tmp, 0, sizeof(tmp));
    skip.userdata = &tmp;
    cst_llm_chat_stream(&skip, e, sizeof(e));
  }
  cst_llm_request_t r;
  memset(&r, 0, sizeof(r));
  r.api_base = "http://replay";
  r.model = "m";
  r.on_delta = on_delta;
  r.on_native_tool = on_tool;
  r.on_native_open = on_open;
  r.userdata = c;
  r.stats = st;
  char err[256];
  return cst_llm_chat_stream(&r, err, sizeof(err));
}

static const char *CHUNKS[] = {"1", "7", "64", "4096"};
#define NCHUNKS 4

static void test_text_reasoning_usage(void) {
  write_fixture(
      "data: {\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\","
      "\"reasoning_content\":\"plan \"}}]}\n\n"
      "data: {\"choices\":[{\"index\":0,\"delta\":{\"reasoning_content\":"
      "\"more\"}}]}\r\n\r\n"
      ": keep-alive comment\n\n"
      "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"Hel\"}}]}\n\n"
      "data:{\"choices\":[{\"index\":0,\"delta\":{\"content\":\"lo\\t\\u4e2d\"},"
      "\"finish_reason\":\"stop\"}]}\n\n"
      "data: {\"choices\":[],\"usage\":{\"prompt_tokens\":11,"
      "\"completion_tokens\":5,\"prompt_cache_hit_tokens\":8}}\n\n"
      "data: [DONE]\n\n");
  for (int k = 0; k < NCHUNKS; k++) {
    capture_t c;
    cst_llm_stats_t st;
    int rc = replay(g_dir, 0, CHUNKS[k], &c, &st);
    EXPECT(rc == 0, "text stream ok");
    EXPECT(strcmp(c.text, "<think>plan more</think>Hello\t\xe4\xb8\xad") == 0,
           "reasoning wrapped + content decoded");
    EXPECT(strcmp(st.finish_reason, "stop") == 0, "finish_reason");
    EXPECT(st.prompt_tokens == 11 && st.completion_tokens == 5 &&
               st.cached_tokens == 8,
           "usage incl. deepseek cache hits");
    EXPECT(c.ncalls == 0, "no calls");
    capture_free(&c);
  }
}

/* Two parallel calls, fragmented arguments, no [DONE] and no final newline. */
static void test_parallel_calls(void) {
  write_fixture(
      "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":"
      "\"c1\",\"type\":\"function\",\"function\":{\"name\":\"read_file\","
      "\"arguments\":\"\"}}]}}]}\n\n"
      "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,"
      "\"function\":{\"arguments\":\"{\\\"pa\"}}]}}]}\n\n"
      "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,"
      "\"function\":{\"arguments\":\"th\\\":\\\"a.c\\\"}\"}}]}}]}\n\n"
      "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":1,\"id\":"
      "\"c2\",\"function\":{\"name\":\"run_shell\",\"arguments\":"
      "\"{\\\"cmd\\\":\"}}]}}]}\n\n"
      "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":1,"
      "\"function\":{\"arguments\":\"\\\"make\\\\ttest\\\"}\"}}]}}]}\n\n"
      "data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"tool_calls\"}]}");
  for (int k = 0; k < NCHUNKS; k++) {
    capture_t c;
    cst_llm_stats_t st;
    int rc = replay(g_dir, 0, CHUNKS[k], &c, &st);
    EXPECT(rc == 0, "parallel stream ok");
    EXPECT(c.ncalls == 2, "two calls");
    if (c.ncalls == 2) {
      EXPECT(strcmp(c.ids[0], "c1") == 0 && strcmp(c.names[0], "read_file") == 0,
             "call 1 id/name");
      EXPECT(c.args[0] && strcmp(c.args[0], "{\"path\":\"a.c\"}") == 0,
             "call 1 args reassembled");
      EXPECT(strcmp(c.ids[1], "c2") == 0 && strcmp(c.names[1], "run_shell") == 0,
             "call 2 id/name");
      EXPECT(c.args[1] && strcmp(c.args[1], "{\"cmd\":\"make\\ttest\"}") == 0,
             "call 2 args reassembled");
      /* native-inc: call 1 is complete as soon as index 1 starts, before its
       * open event and long before finish_reason */
      EXPECT(c.open_before_first_call == 1, "call 1 fired when index 1 began");
    }
    EXPECT(c.nopen == 2, "two opens");
    EXPECT(strcmp(st.finish_reason, "tool_calls") == 0, "finish tool_calls");
    capture_free(&c);
  }
}

/* Gateways that reuse index 0 for every call but change the id. */
static void test_same_index_new_id(void) {
  write_fixture(
      "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":"
      "\"a\",\"function\":{\"name\":\"read_file\",\"arguments\":"
      "\"{\\\"path\\\":\\\"x\\\"}\"}}]}}]}\n\n"
      "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":"
      "\"b\",\"function\":{\"name\":\"read_file\",\"arguments\":"
      "\"{\\\"path\\\":\\\"y\\\"}\"}}]}}]}\n\n"
      "data: [DONE]\n\n");
  capture_t c;
  cst_llm_stats_t st;
  replay(g_dir, 0, "4096", &c, &st);
  EXPECT(c.ncalls == 2, "same index, new id -> two calls");
  if (c.ncalls == 2)
    EXPECT(c.args[1] && strstr(c.args[1], "\"y\""), "second call args");
  capture_free(&c);
}

static void test_invalid_args_and_error(void) {
  write_fixture(
      "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":"
      "\"c1\",\"function\":{\"name\":\"write_file\",\"arguments\":"
      "\"{\\\"path\\\":\\\"a\"}}]}}]}\n\n"
      "data: [DONE]\n\n");
  capture_t c;
  cst_llm_stats_t st;
  replay(g_dir, 0, "3", &c, &st);
  EXPECT(c.ncalls == 1 && c.args_null[0], "truncated args -> NULL");
  capture_free(&c);

  write_fixture("data: {\"error\":{\"message\":\"quota exceeded\"}}\n\n");
  int rc = replay(g_dir, 0, "4096", &c, &st);
  EXPECT(rc != 0, "in-stream error fails the request");
  EXPECT(strstr(st.error, "quota exceeded") != NULL, "error message kept");
  capture_free(&c);
}

static int file_exists(const char *p) { return access(p, R_OK) == 0; }

/* Recorded DeepSeek traces: every call must decode to a JSON object and the
 * decoded result must not depend on the network chunking. */
static void test_recorded(void) {
  const char *dir = "tests/fixtures/sse/deepseek-native-eos";
  char p[512];
  snprintf(p, sizeof(p), "%s/0000.sse", dir);
  if (!file_exists(p)) {
    printf("  (skip recorded traces: %s missing)\n", p);
    return;
  }
  int total_calls = 0;
  for (int seq = 0; seq < 16; seq++) {
    snprintf(p, sizeof(p), "%s/%04d.sse", dir, seq);
    if (!file_exists(p))
      break;
    capture_t ref;
    cst_llm_stats_t st;
    EXPECT(replay(dir, seq, "4096", &ref, &st) == 0, "recorded replay ok");
    for (int i = 0; i < ref.ncalls; i++) {
      cJSON *j = ref.args[i] ? cJSON_Parse(ref.args[i]) : NULL;
      EXPECT(cJSON_IsObject(j), "recorded call args are a JSON object");
      EXPECT(cst_tool_known(ref.names[i]), "recorded call names a known tool");
      cJSON_Delete(j);
    }
    total_calls += ref.ncalls;
    EXPECT(st.prompt_tokens > 0, "recorded usage present");
    capture_t one;
    replay(dir, seq, "1", &one, &st);
    EXPECT(one.ncalls == ref.ncalls, "chunking does not change calls");
    EXPECT(one.text_len == ref.text_len &&
               memcmp(one.text, ref.text, ref.text_len) == 0,
           "chunking does not change text");
    for (int i = 0; i < one.ncalls && i < ref.ncalls; i++)
      EXPECT(one.args[i] && ref.args[i] && strcmp(one.args[i], ref.args[i]) == 0,
             "chunking does not change args");
    capture_free(&one);
    capture_free(&ref);
  }
  EXPECT(total_calls > 0, "recorded native trace contains tool calls");

  const char *hd = "tests/fixtures/sse/deepseek-ctl-heredoc";
  snprintf(p, sizeof(p), "%s/0000.sse", hd);
  if (file_exists(p)) {
    capture_t c;
    cst_llm_stats_t st;
    replay(hd, 0, "5", &c, &st);
    EXPECT(strstr(c.text, "call ") != NULL, "ctl trace carries CTL text");
    EXPECT(c.ncalls == 0, "text protocol: no native calls");
    capture_free(&c);
  }
}

int main(void) {
  snprintf(g_dir, sizeof(g_dir), "/tmp/cst-sse-test-%d", (int)getpid());
  char cmd[300];
  snprintf(cmd, sizeof(cmd), "mkdir -p %s", g_dir);
  if (system(cmd) != 0)
    return 1;
  test_text_reasoning_usage();
  test_parallel_calls();
  test_same_index_new_id();
  test_invalid_args_and_error();
  test_recorded();
  snprintf(cmd, sizeof(cmd), "rm -rf %s", g_dir);
  if (system(cmd) != 0)
    fprintf(stderr, "warning: cleanup failed\n");
  if (nfail) {
    fprintf(stderr, "test_sse_decode: %d failed\n", nfail);
    return 1;
  }
  printf("test_sse_decode: OK\n");
  return 0;
}