#define _POSIX_C_SOURCE 200809L

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

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

/* Runs a few replayed turns (tool call, Plan mode on, approval change, a
 * reply cut off by max tokens) and checks the recorded request bodies: every
 * request must extend the previous one byte for byte, so the provider's
 * prefix cache keeps hitting. */

static int nfail;

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

#define USAGE                                                                  \
  "data: {\"choices\":[],\"usage\":{\"prompt_tokens\":1000,"                   \
  "\"completion_tokens\":20,\"prompt_cache_hit_tokens\":900,"                  \
  "\"prompt_cache_miss_tokens\":100}}\n\n"                                     \
  "data: [DONE]\n\n"

#define TOOL(id, name, args)                                                   \
  "data: {\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":0,"  \
  "\"id\":\"" id "\",\"type\":\"function\",\"function\":{\"name\":\"" name     \
  "\",\"arguments\":\"" args "\"}}]},\"finish_reason\":null}]}\n\n"            \
  "data: {\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":"           \
  "\"tool_calls\"}]}\n\n" USAGE

#undef TEXT /* <windows.h> has one */
#define TEXT(s, finish)                                                        \
  "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"" s "\"},"        \
  "\"finish_reason\":null}]}\n\n"                                              \
  "data: {\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"" finish  \
  "\"}]}\n\n" USAGE

#define THINK(s)                                                               \
  "data: {\"choices\":[{\"index\":0,\"delta\":{\"reasoning_content\":\"" s     \
  "\"},\"finish_reason\":null}]}\n\n"

static const char *FIXTURES[] = {
    /* turn 1 */
    THINK("look at the file") TOOL("c1", "read_file", "{\\\"path\\\":\\\"a.txt\\\"}"),
    TEXT("done one", "stop"),
    /* turn 2, Plan mode */
    TOOL("c2", "write_file",
         "{\\\"path\\\":\\\"b.txt\\\",\\\"content\\\":\\\"x\\\"}"),
    TEXT("plan reply", "stop"),
    /* turn 3, Build mode + yolo; cut off, then continued */
    TEXT("part one", "length"),
    TEXT(" part two", "stop"),
    /* turn 4 */
    TEXT("bye", "stop"),
};
#define NFIX ((int)(sizeof(FIXTURES) / sizeof(FIXTURES[0])))

static void write_file(const char *dir, const char *name, const char *body) {
  char p[2048];
  snprintf(p, sizeof(p), "%s/%s", dir, name);
  FILE *fp = fopen(p, "wb");
  EXPECT(fp != NULL, "write fixture");
  if (fp) {
    fputs(body, fp);
    fclose(fp);
  }
}

static cJSON *load_req(const char *dir, int seq) {
  char p[2048];
  snprintf(p, sizeof(p), "%s/%04d.req.json", dir, seq);
  FILE *fp = fopen(p, "rb");
  if (!fp)
    return NULL;
  cst_buf_t b = {0};
  char buf[4096];
  size_t n;
  while ((n = fread(buf, 1, sizeof(buf), fp)) > 0)
    cst_buf_append(&b, buf, n);
  fclose(fp);
  cJSON *j = b.buf ? cJSON_Parse(b.buf) : NULL;
  cst_buf_free(&b);
  return j;
}

static const cJSON *msgs_of(const cJSON *req) {
  return cJSON_GetObjectItemCaseSensitive(req, "messages");
}

static int is_prefix_msg(const cJSON *m) {
  return cJSON_IsTrue(cJSON_GetObjectItemCaseSensitive(m, "prefix"));
}

static const char *str_of(const cJSON *o, const char *key) {
  const cJSON *v = cJSON_GetObjectItemCaseSensitive(o, key);
  return cJSON_IsString(v) ? v->valuestring : NULL;
}

static char *print(const cJSON *j) { return j ? cJSON_PrintUnformatted(j) : NULL; }

/* messages of `a` (minus a request-only prefix message) start `b` */
static int extends(const cJSON *a, const cJSON *b) {
  const cJSON *ma = msgs_of(a), *mb = msgs_of(b);
  int na = cJSON_GetArraySize(ma), nb = cJSON_GetArraySize(mb);
  if (na > 0 && is_prefix_msg(cJSON_GetArrayItem(ma, na - 1)))
    na--;
  if (na > nb)
    return 0;
  for (int i = 0; i < na; i++) {
    char *x = print(cJSON_GetArrayItem(ma, i));
    char *y = print(cJSON_GetArrayItem(mb, i));
    int same = x && y && strcmp(x, y) == 0;
    if (!same)
      fprintf(stderr, "  message %d differs:\n    %.300s\n    %.300s\n", i,
              x ? x : "(null)", y ? y : "(null)");
    free(x);
    free(y);
    if (!same)
      return 0;
  }
  return 1;
}

static const cJSON *last_msg(const cJSON *req) {
  const cJSON *m = msgs_of(req);
  return cJSON_GetArrayItem(m, cJSON_GetArraySize(m) - 1);
}

static int any_msg_has(const cJSON *req, const char *key, const char *needle) {
  const cJSON *m;
  cJSON_ArrayForEach(m, msgs_of(req)) {
    const char *s = str_of(m, key);
    if (s && strstr(s, needle))
      return 1;
  }
  return 0;
}

static int run(cst_agent_t *a, const char *text) {
  char err[512];
  int rc = cst_agent_run_turn(a, text, err, sizeof(err));
  if (rc != 0)
    fprintf(stderr, "  turn \"%s\": %s\n", text, err);
  return rc;
}

/* 8-byte PNG signature + IHDR header: enough for the attachment sniffer. */
static size_t tiny_png(unsigned char *b) {
  static const unsigned char sig[8] = {0x89, 'P', 'N', 'G', '\r', '\n', 0x1A,
                                       '\n'};
  memset(b, 0, 32);
  memcpy(b, sig, 8);
  b[11] = 13;
  memcpy(b + 12, "IHDR", 4);
  b[19] = 2; /* 2 x 2 */
  b[23] = 2;
  return 32;
}

static int tools_have(const cJSON *req, const char *name) {
  char *t = print(cJSON_GetObjectItemCaseSensitive(req, "tools"));
  int has = t && strstr(t, name) != NULL;
  free(t);
  return has;
}

static int has_image_url(const cJSON *req) {
  char *t = print(msgs_of(req));
  int has = t && strstr(t, "image_url") != NULL;
  free(t);
  return has;
}

/* A text-only model never sees images or view_image; a model of unknown
 * vision that answers 400 to an image is retried with notes, and stays
 * that way so later requests extend the retried one. */
static void test_vision(const char *root) {
  char dir[1100], sse[1200], rec[1200];
  snprintf(dir, sizeof(dir), "%s/vision", root);
  snprintf(sse, sizeof(sse), "%s/sse", dir);
  snprintf(rec, sizeof(rec), "%s/rec", dir);
  os_mkdir_p(sse, 0755);
  write_file(sse, "0000.sse", TEXT("blind reply", "stop"));
  write_file(sse, "0001.sse", "{\"error\":{\"message\":\"image input is not "
                              "supported\"}}");
  write_file(sse, "0001.status", "400\n");
  write_file(sse, "0002.sse", TEXT("noted reply", "stop"));
  write_file(sse, "0003.sse", TEXT("later reply", "stop"));
  setenv("CST_REPLAY_SSE", sse, 1);
  setenv("CST_RECORD_SSE", rec, 1);
  cst_llm_trace_reset();

  unsigned char png[32];
  size_t pn = tiny_png(png);
  char err[256];

  cst_agent_config_t cfg;
  memset(&cfg, 0, sizeof(cfg));
  snprintf(cfg.workspace_root, sizeof(cfg.workspace_root), "%s", root);
  snprintf(cfg.session_path, sizeof(cfg.session_path), "%s/blind.jsonl", dir);
  snprintf(cfg.api_base, sizeof(cfg.api_base), "http://replay/v1");
  snprintf(cfg.model, sizeof(cfg.model), "deepseek-v4-pro");
  snprintf(cfg.provider, sizeof(cfg.provider), "deepseek");
  cfg.protocol = CST_PROTO_NATIVE_EOS;
  cfg.max_steps = 4;
  cfg.headless = 1;
  cfg.approval = CST_APPROVAL_AUTO;
  cst_agent_t *a = cst_agent_create(&cfg);
  EXPECT(a != NULL, "blind agent");
  if (!a)
    return;
  EXPECT(cst_agent_sees_images(a) == 0, "deepseek-v4-pro is text only");
  EXPECT(cst_agent_attach_bytes(a, "shot.png", (const char *)png, pn, err,
                                sizeof(err)) == 0,
         "attach to a blind model");
  EXPECT(run(a, "what is in the picture?") == 0, "blind turn");
  cst_agent_destroy(a);

  snprintf(cfg.session_path, sizeof(cfg.session_path), "%s/unknown.jsonl",
           dir);
  snprintf(cfg.model, sizeof(cfg.model), "mystery-model-9");
  snprintf(cfg.provider, sizeof(cfg.provider), "teamorouter");
  a = cst_agent_create(&cfg);
  EXPECT(a != NULL, "unknown-vision agent");
  if (!a)
    return;
  EXPECT(cst_agent_sees_images(a) != 0, "unknown vision: images are tried");
  EXPECT(cst_agent_attach_bytes(a, "shot.png", (const char *)png, pn, err,
                                sizeof(err)) == 0,
         "attach to an unknown model");
  EXPECT(run(a, "describe it") == 0, "400 retried without images");
  EXPECT(run(a, "and now?") == 0, "next turn");
  cst_agent_destroy(a);

  cJSON *r[4];
  for (int i = 0; i < 4; i++) {
    r[i] = load_req(rec, i);
    EXPECT(r[i] != NULL, "vision request recorded");
  }
  if (r[0] && r[1] && r[2] && r[3]) {
    EXPECT(!tools_have(r[0], "view_image"), "blind: no view_image tool");
    EXPECT(!has_image_url(r[0]), "blind: no image part");
    const char *c = str_of(last_msg(r[0]), "content");
    EXPECT(c && strstr(c, "what is in the picture?") &&
               strstr(c, "not sent: this model does not take images"),
           "blind: the image becomes a note");

    EXPECT(tools_have(r[1], "view_image") && has_image_url(r[1]),
           "unknown: first try sends the image");
    EXPECT(!tools_have(r[2], "view_image") && !has_image_url(r[2]),
           "retry: no image, no view_image");
    c = str_of(last_msg(r[2]), "content");
    EXPECT(c && strstr(c, "not sent"), "retry: note instead");
    EXPECT(!has_image_url(r[3]) && extends(r[2], r[3]),
           "images stay off; the next request extends the retry");
  }
  for (int i = 0; i < 4; i++)
    cJSON_Delete(r[i]);
}

int main(void) {
  char root[1024], sse[1100], rec[1100];
  if (test_tmpdir(root, sizeof(root), "cst-prefix") != 0) {
    fprintf(stderr, "cannot create a scratch dir\n");
    return 1;
  }
  snprintf(sse, sizeof(sse), "%s/sse", root);
  snprintf(rec, sizeof(rec), "%s/rec", root);
  os_mkdir_p(sse, 0755);
  for (int i = 0; i < NFIX; i++) {
    char name[32];
    snprintf(name, sizeof(name), "%04d.sse", i);
    write_file(sse, name, FIXTURES[i]);
  }
  write_file(root, "a.txt", "hello\n");
  setenv("CST_REPLAY_SSE", sse, 1);
  setenv("CST_RECORD_SSE", rec, 1);
  setenv("CST_MCP_WAIT_MS", "0", 1);
  setenv("XDG_CONFIG_HOME", root, 1); /* no real models.json or caches */
  unsetenv("CST_CTX_WINDOW");
  unsetenv("CST_NO_CACHE_HINTS");
  unsetenv("CST_PRICE");
  cst_llm_trace_reset();

  cst_agent_config_t cfg;
  memset(&cfg, 0, sizeof(cfg));
  snprintf(cfg.workspace_root, sizeof(cfg.workspace_root), "%s", root);
  snprintf(cfg.session_path, sizeof(cfg.session_path), "%s/s.jsonl", root);
  snprintf(cfg.api_base, sizeof(cfg.api_base), "http://replay/v1");
  snprintf(cfg.model, sizeof(cfg.model), "deepseek-flash");
  snprintf(cfg.provider, sizeof(cfg.provider), "deepseek");
  cfg.protocol = CST_PROTO_NATIVE_EOS;
  cfg.max_steps = 8;
  cfg.headless = 1;
  cfg.approval = CST_APPROVAL_AUTO;
  cst_agent_t *a = cst_agent_create(&cfg);
  EXPECT(a != NULL, "agent created");
  if (!a)
    return 1;

  EXPECT(run(a, "read a.txt") == 0, "turn 1");
  cst_agent_config_mut(a)->plan_mode = 1;
  EXPECT(run(a, "now write b.txt") == 0, "turn 2 (plan)");
  cst_agent_config_mut(a)->plan_mode = 0;
  cst_agent_config_mut(a)->approval = CST_APPROVAL_YOLO;
  EXPECT(run(a, "a long answer please") == 0, "turn 3 (continued)");
  EXPECT(run(a, "bye") == 0, "turn 4");

  char bpath[1100];
  snprintf(bpath, sizeof(bpath), "%s/b.txt", root);
  EXPECT(access(bpath, F_OK) != 0, "Plan mode refused write_file");

  cJSON *req[NFIX];
  for (int i = 0; i < NFIX; i++) {
    req[i] = load_req(rec, i);
    EXPECT(req[i] != NULL, "request recorded");
  }
  if (req[NFIX - 1]) {
    char *tools0 = print(cJSON_GetObjectItemCaseSensitive(req[0], "tools"));
    char *sys0 = print(cJSON_GetArrayItem(msgs_of(req[0]), 0));
    EXPECT(tools0 && strstr(tools0, "write_file"),
           "write tools stay in the schema");
    for (int i = 0; i < NFIX; i++) {
      char *tools = print(cJSON_GetObjectItemCaseSensitive(req[i], "tools"));
      char *sys = print(cJSON_GetArrayItem(msgs_of(req[i]), 0));
      EXPECT(tools && tools0 && strcmp(tools, tools0) == 0,
             "tools[] identical in every request");
      EXPECT(sys && sys0 && strcmp(sys, sys0) == 0,
             "system prompt identical in every request");
      const char *tc = str_of(req[i], "tool_choice");
      EXPECT(tc && strcmp(tc, "auto") == 0, "tool_choice stays auto");
      free(tools);
      free(sys);
      if (i + 1 < NFIX && !extends(req[i], req[i + 1])) {
        fprintf(stderr, "  request %d does not extend request %d\n", i + 1, i);
        nfail++;
      }
    }
    free(tools0);
    free(sys0);

    EXPECT(any_msg_has(req[1], "reasoning_content", "look at the file"),
           "DeepSeek reasoning passed back");
    EXPECT(strstr(str_of(last_msg(req[2]), "content"), "Plan mode is now ON"),
           "Plan mode announced by a note");
    EXPECT(any_msg_has(req[3], "content", "Plan mode is on"),
           "write_file denied with a Plan mode error");
    const char *t3 = str_of(last_msg(req[4]), "content");
    EXPECT(t3 && strstr(t3, "Plan mode is now OFF") &&
               strstr(t3, "Approval mode is now yolo"),
           "Build mode and approval change announced");
    const cJSON *pm = last_msg(req[5]);
    EXPECT(is_prefix_msg(pm) && strcmp(str_of(pm, "role"), "assistant") == 0 &&
               strcmp(str_of(pm, "content"), "part one") == 0,
           "cut-off reply continued through an assistant prefix");
    EXPECT(any_msg_has(req[6], "content", "part one part two"),
           "continued reply stored as one message");
  }
  for (int i = 0; i < NFIX; i++)
    cJSON_Delete(req[i]);

  cst_usage_t s, turn;
  cst_agent_usage(a, &s, &turn);
  EXPECT(s.requests == NFIX, "every response counted");
  EXPECT(s.prompt == 1000L * NFIX && s.cached == 900L * NFIX &&
             s.miss == 100L * NFIX,
         "hit and miss tokens summed");
  EXPECT(s.priced && s.cost > 0 && strcmp(s.currency, "¥") == 0,
         "DeepSeek cost estimated");
  EXPECT(turn.requests == 1, "last turn usage");

  cst_agent_destroy(a);
  test_vision(root);
  test_rmrf(root);
  if (nfail) {
    fprintf(stderr, "test_prefix_stable: %d failed\n", nfail);
    return 1;
  }
  printf("test_prefix_stable: OK\n");
  return 0;
}