#define _POSIX_C_SOURCE 200809L

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

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

/*
 * Truncation experiment on recorded model streams (tests/fixtures/sse).
 *
 * For every recorded assistant message we rebuild the logical stream the
 * runtime sees and cut it at every byte:
 *   - CTL traces: the content text; executable calls = CALL_READY events.
 *   - native traces: tool calls as a JSON array of {"name","arguments"};
 *       json-whole : execute only when the whole array parses (EOS baseline)
 *       json-object: execute each object once it parses (incremental JSON)
 *
 * Checked properties (the CST soundness claims):
 *   S1 no call is executable on a prefix unless its bytes are complete;
 *   S2 every call executed on a prefix is identical to the same call in the
 *      full stream (prefix consistency, i.e. execution is never revoked);
 *   S3 the number of executable calls is monotone in the prefix length.
 * The printed table reports how early calls become executable (experiment C).
 */

static int nfail;

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

typedef struct {
  char *buf;
  size_t len, cap;
} sbuf_t;

static void sb_add(sbuf_t *b, const char *s, size_t n) {
  if (b->len + n + 1 > b->cap) {
    size_t nc = b->cap ? b->cap * 2 : 1024;
    while (nc < b->len + n + 1)
      nc *= 2;
    b->buf = realloc(b->buf, nc);
    b->cap = nc;
  }
  memcpy(b->buf + b->len, s, n);
  b->len += n;
  b->buf[b->len] = 0;
}

#define MAXC 32

typedef struct {
  sbuf_t content;
  int ncalls;
  char names[MAXC][64];
  sbuf_t args[MAXC];
} trace_t;

static void trace_free(trace_t *t) {
  free(t->content.buf);
  for (int i = 0; i < MAXC; i++)
    free(t->args[i].buf);
  memset(t, 0, sizeof(*t));
}

static int load_trace(const char *path, trace_t *t) {
  memset(t, 0, sizeof(*t));
  FILE *fp = fopen(path, "rb");
  if (!fp)
    return -1;
  char *line = NULL;
  size_t cap = 0;
  ssize_t n;
  while ((n = test_getline(&line, &cap, fp)) > 0) {
    if (strncmp(line, "data:", 5) != 0)
      continue;
    cJSON *j = cJSON_Parse(line + 5);
    if (!j)
      continue;
    const cJSON *ch = cJSON_GetArrayItem(cJSON_GetObjectItem(j, "choices"), 0);
    const cJSON *d = cJSON_GetObjectItem(ch, "delta");
    const cJSON *c = cJSON_GetObjectItem(d, "content");
    if (cJSON_IsString(c))
      sb_add(&t->content, c->valuestring, strlen(c->valuestring));
    const cJSON *tc = NULL;
    cJSON_ArrayForEach(tc, cJSON_GetObjectItem(d, "tool_calls")) {
      const cJSON *ji = cJSON_GetObjectItem(tc, "index");
      int idx = cJSON_IsNumber(ji) ? ji->valueint : 0;
      if (idx < 0 || idx >= MAXC)
        continue;
      if (idx >= t->ncalls)
        t->ncalls = idx + 1;
      const cJSON *fn = cJSON_GetObjectItem(tc, "function");
      const cJSON *nm = cJSON_GetObjectItem(fn, "name");
      const cJSON *ar = cJSON_GetObjectItem(fn, "arguments");
      if (cJSON_IsString(nm) && nm->valuestring[0])
        snprintf(t->names[idx], sizeof(t->names[idx]), "%s", nm->valuestring);
      if (cJSON_IsString(ar))
        sb_add(&t->args[idx], ar->valuestring, strlen(ar->valuestring));
    }
    cJSON_Delete(j);
  }
  free(line);
  fclose(fp);
  return 0;
}

/* --- CTL side ------------------------------------------------------------- */

typedef struct {
  int n;
  char sig[MAXC][512]; /* id|name|k=v... (truncated) as identity */
} ready_t;

static void on_ctl(const ctl_event_t *ev, void *ud) {
  ready_t *r = ud;
  if (ev->type != CTL_EV_CALL_READY || r->n >= MAXC)
    return;
  char *s = r->sig[r->n++];
  int o = snprintf(s, 512, "%s|%s", ev->call->id, ev->call->name);
  for (int i = 0; i < ev->call->nargs && o < 500; i++)
    o += snprintf(s + o, (size_t)(512 - o), "|%s=%.*s", ev->call->args[i].key,
                  (int)(ev->call->args[i].value_len > 200
                            ? 200
                            : ev->call->args[i].value_len),
                  ev->call->args[i].value);
}

static void ctl_ready(const char *s, size_t n, int finish, ready_t *r) {
  memset(r, 0, sizeof(*r));
  ctl_parser_t *p = ctl_parser_create(on_ctl, r);
  ctl_parser_feed(p, s, n);
  if (finish)
    ctl_parser_finish(p);
  ctl_parser_destroy(p);
}

/* --- JSON side ------------------------------------------------------------ */

static void build_json_stream(const trace_t *t, sbuf_t *out) {
  sb_add(out, "[", 1);
  for (int i = 0; i < t->ncalls; i++) {
    char head[160];
    int hn = snprintf(head, sizeof(head), "%s{\"name\":\"%s\",\"arguments\":",
                      i ? "," : "", t->names[i]);
    sb_add(out, head, (size_t)hn);
    if (t->args[i].len)
      sb_add(out, t->args[i].buf, t->args[i].len);
    else
      sb_add(out, "{}", 2);
    sb_add(out, "}", 1);
  }
  sb_add(out, "]", 1);
}

static int json_whole(const char *s, size_t n) {
  cJSON *j = cJSON_ParseWithLength(s, n);
  int k = cJSON_IsArray(j) ? cJSON_GetArraySize(j) : 0;
  cJSON_Delete(j);
  return k;
}

static int json_objects(const char *s, size_t n) {
  int k = 0;
  size_t off = 0;
  while (off < n) {
    cst_json_tool_t t;
    const char *p = memchr(s + off, '{', n - off);
    if (!p)
      break;
    if (!cst_json_try_parse_tool(p, n - (size_t)(p - s), &t))
      break;
    k++;
    /* skip past this object */
    const char *end = NULL;
    cJSON *j = cJSON_ParseWithLengthOpts(p, n - (size_t)(p - s), &end, 0);
    cJSON_Delete(j);
    if (!end || end <= p)
      break;
    off = (size_t)(end - s);
  }
  return k;
}

/* --- experiment ----------------------------------------------------------- */

static void report_row(const char *label, size_t n, const int *at_pct,
                       int total) {
  printf("  %-26s bytes=%-6zu calls=%d  executable@25/50/75/90/100%%: "
         "%d/%d/%d/%d/%d\n",
         label, n, total, at_pct[0], at_pct[1], at_pct[2], at_pct[3],
         at_pct[4]);
}

static const int PCTS[5] = {25, 50, 75, 90, 100};

static void run_ctl(const char *label, const sbuf_t *c) {
  ready_t full;
  ctl_ready(c->buf, c->len, 1, &full);
  int prev = 0, at[5] = {0};
  size_t first_ready_at = 0;
  for (size_t cut = 0; cut <= c->len; cut++) {
    ready_t r;
    ctl_ready(c->buf, cut, cut == c->len, &r);
    EXPECT(r.n >= prev, "S3 monotone");
    EXPECT(r.n <= full.n, "S1 never more than the full stream");
    for (int i = 0; i < r.n && i < full.n; i++)
      EXPECT(strcmp(r.sig[i], full.sig[i]) == 0, "S2 prefix-consistent call");
    if (r.n > 0 && prev == 0)
      first_ready_at = cut;
    prev = r.n;
    for (int k = 0; k < 5; k++)
      if (cut == c->len * (size_t)PCTS[k] / 100)
        at[k] = r.n;
  }
  report_row(label, c->len, at, full.n);
  if (full.n > 0)
    printf("  %-26s first call executable at byte %zu (%.0f%% of message)\n",
           "", first_ready_at, 100.0 * (double)first_ready_at / (double)c->len);
}

static void run_native(const char *label, const trace_t *t) {
  sbuf_t js = {0};
  build_json_stream(t, &js);
  int whole_full = json_whole(js.buf, js.len);
  int obj_full = json_objects(js.buf, js.len);
  EXPECT(whole_full == t->ncalls, "json-whole parses at full length");
  EXPECT(obj_full == t->ncalls, "json-object parses all at full length");
  int at_w[5] = {0}, at_o[5] = {0}, prev_o = 0;
  for (size_t cut = 0; cut <= js.len; cut++) {
    int w = json_whole(js.buf, cut);
    int o = json_objects(js.buf, cut);
    if (cut < js.len)
      EXPECT(w == 0, "json-whole executes nothing before the end");
    EXPECT(o >= prev_o, "S3 monotone (json-object)");
    prev_o = o;
    for (int k = 0; k < 5; k++) {
      if (cut == js.len * (size_t)PCTS[k] / 100) {
        at_w[k] = w;
        at_o[k] = o;
      }
    }
  }
  char l[80];
  snprintf(l, sizeof(l), "%s json-whole", label);
  report_row(l, js.len, at_w, whole_full);
  snprintf(l, sizeof(l), "%s json-object", label);
  report_row(l, js.len, at_o, obj_full);
  free(js.buf);
}

int main(void) {
  static const char *dirs[] = {"tests/fixtures/sse/deepseek-ctl-heredoc",
                               "tests/fixtures/sse/deepseek-ctl",
                               "tests/fixtures/sse/deepseek-native-eos"};
  int traces = 0;
  for (size_t d = 0; d < sizeof(dirs) / sizeof(dirs[0]); d++) {
    printf("%s\n", dirs[d]);
    for (int seq = 0; seq < 64; seq++) {
      char path[512];
      snprintf(path, sizeof(path), "%s/%04d.sse", dirs[d], seq);
      if (access(path, R_OK) != 0)
        break;
      trace_t t;
      if (load_trace(path, &t) != 0)
        continue;
      traces++;
      char label[64];
      snprintf(label, sizeof(label), "#%04d", seq);
      if (t.ncalls > 0)
        run_native(label, &t);
      else if (t.content.len > 0 && strstr(t.content.buf, "call "))
        run_ctl(label, &t.content);
      else
        printf("  %-26s (final answer, no calls)\n", label);
      trace_free(&t);
    }
  }
  if (traces == 0)
    printf("(no recorded traces found; run from the repo root)\n");
  if (nfail) {
    fprintf(stderr, "test_replay_real: %d failed\n", nfail);
    return 1;
  }
  printf("test_replay_real: OK\n");
  return 0;
}