#include "cst_agent.h"

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

#define TC_SLOTS 8
#define TC_ID_MAX 128
#define TC_NAME_MAX 128
#define TC_ARGS_MAX 65536

typedef struct {
  int used;
  int index;
  char id[TC_ID_MAX];
  char name[TC_NAME_MAX];
  char *args;
  size_t args_len;
  size_t args_cap;
} tc_slot_t;

typedef struct {
  cst_llm_on_delta_fn on_delta;
  cst_llm_on_native_tool_fn on_native_tool;
  void *userdata;
  char *buf;
  size_t len;
  size_t cap;
  int in_data;
  int in_reason; /* wrapping reasoning_content as <think> for the agent */
  tc_slot_t slots[TC_SLOTS];
} sse_acc_t;

static void emit_raw(sse_acc_t *a, const char *s, size_t n) {
  if (a && a->on_delta && s && n)
    a->on_delta(s, n, a->userdata);
}

static const char *find_json_string_key(const char *payload, const char *key) {
  /* Find "key" that is a full JSON key, not a suffix of another key
   * (e.g. avoid matching content inside audio_content / reasoning_content). */
  size_t klen = strlen(key);
  const char *p = payload;
  while ((p = strstr(p, key)) != NULL) {
    if (p > payload) {
      unsigned char prev = (unsigned char)p[-1];
      if (prev == '_' || isalnum(prev)) {
        p += klen;
        continue;
      }
    }
    return p;
  }
  return NULL;
}

/* Decode one JSON string value starting after the opening quote into out. */
static size_t decode_json_string(const char *colon_quote, char *out,
                                 size_t out_sz) {
  size_t o = 0;
  for (const char *p = colon_quote; *p && o + 4 < out_sz; p++) {
    if (*p == '\\' && p[1]) {
      p++;
      if (*p == 'n')
        out[o++] = '\n';
      else if (*p == 'r')
        out[o++] = '\r';
      else if (*p == 't')
        out[o++] = '\t';
      else if (*p == '"' || *p == '\\' || *p == '/')
        out[o++] = *p;
      else if (*p == 'u' && p[1] && p[2] && p[3] && p[4]) {
        unsigned cp = 0;
        int ok = 1;
        for (int i = 1; i <= 4; i++) {
          char h = p[i];
          cp <<= 4;
          if (h >= '0' && h <= '9')
            cp |= (unsigned)(h - '0');
          else if (h >= 'a' && h <= 'f')
            cp |= (unsigned)(h - 'a' + 10);
          else if (h >= 'A' && h <= 'F')
            cp |= (unsigned)(h - 'A' + 10);
          else {
            ok = 0;
            break;
          }
        }
        p += 4;
        if (!ok) {
          out[o++] = '?';
          continue;
        }
        if (cp < 0x80)
          out[o++] = (char)cp;
        else if (cp < 0x800) {
          out[o++] = (char)(0xC0 | (cp >> 6));
          out[o++] = (char)(0x80 | (cp & 0x3F));
        } else {
          out[o++] = (char)(0xE0 | (cp >> 12));
          out[o++] = (char)(0x80 | ((cp >> 6) & 0x3F));
          out[o++] = (char)(0x80 | (cp & 0x3F));
        }
      } else {
        out[o++] = *p;
      }
      continue;
    }
    if (*p == '"')
      break;
    out[o++] = *p;
  }
  return o;
}

static int extract_json_string_value(const char *payload, const char *key,
                                     char *out, size_t out_sz, size_t *out_len) {
  const char *keyp = find_json_string_key(payload, key);
  if (!keyp)
    return 0;
  const char *colon = strchr(keyp + strlen(key), ':');
  if (!colon)
    return 0;
  colon++;
  while (*colon == ' ')
    colon++;
  if (*colon == 'n' && strncmp(colon, "null", 4) == 0)
    return 0;
  if (*colon != '"')
    return 0;
  if (colon[1] == '"')
    return 0; /* empty */
  size_t n = decode_json_string(colon + 1, out, out_sz);
  if (n == 0)
    return 0;
  *out_len = n;
  return 1;
}

static int json_brace_balanced(const char *s) {
  if (!s || *s != '{')
    return 0;
  int depth = 0;
  int in_str = 0;
  for (const char *p = s; *p; p++) {
    if (in_str) {
      if (*p == '\\' && p[1]) {
        p++;
        continue;
      }
      if (*p == '"')
        in_str = 0;
      continue;
    }
    if (*p == '"') {
      in_str = 1;
      continue;
    }
    if (*p == '{')
      depth++;
    else if (*p == '}') {
      depth--;
      if (depth == 0)
        return p[1] == 0;
      if (depth < 0)
        return 0;
    }
  }
  return 0;
}

static tc_slot_t *slot_for_index(sse_acc_t *a, int index) {
  for (int i = 0; i < TC_SLOTS; i++) {
    if (a->slots[i].used && a->slots[i].index == index)
      return &a->slots[i];
  }
  for (int i = 0; i < TC_SLOTS; i++) {
    if (!a->slots[i].used) {
      a->slots[i].used = 1;
      a->slots[i].index = index;
      a->slots[i].id[0] = 0;
      a->slots[i].name[0] = 0;
      a->slots[i].args = NULL;
      a->slots[i].args_len = 0;
      a->slots[i].args_cap = 0;
      return &a->slots[i];
    }
  }
  return NULL;
}

static int slot_append_args(tc_slot_t *s, const char *frag, size_t n) {
  if (!s || !frag || !n)
    return 0;
  if (s->args_len + n + 1 > s->args_cap) {
    size_t nc = s->args_cap ? s->args_cap * 2 : 1024;
    while (nc < s->args_len + n + 1)
      nc *= 2;
    if (nc > TC_ARGS_MAX)
      nc = TC_ARGS_MAX;
    if (s->args_len + n + 1 > nc)
      return -1;
    char *nb = realloc(s->args, nc);
    if (!nb)
      return -1;
    s->args = nb;
    s->args_cap = nc;
  }
  memcpy(s->args + s->args_len, frag, n);
  s->args_len += n;
  s->args[s->args_len] = 0;
  return 0;
}

/* Parse one tool_calls[] element object for index/id/name/arguments fragments. */
static void parse_one_tool_call_obj(sse_acc_t *a, const char *obj) {
  if (!obj || *obj != '{')
    return;

  int index = 0;
  const char *ip = strstr(obj, "\"index\"");
  if (ip) {
    const char *colon = strchr(ip, ':');
    if (colon) {
      colon++;
      while (*colon == ' ')
        colon++;
      index = (int)strtol(colon, NULL, 10);
    }
  }

  tc_slot_t *slot = slot_for_index(a, index);
  if (!slot)
    return;

  char tmp[4096];
  size_t tlen = 0;
  if (extract_json_string_value(obj, "\"id\"", tmp, sizeof(tmp), &tlen)) {
    if (tlen >= TC_ID_MAX)
      tlen = TC_ID_MAX - 1;
    memcpy(slot->id, tmp, tlen);
    slot->id[tlen] = 0;
  }

  /* function.name / function.arguments may be nested under "function":{...} */
  const char *fn = strstr(obj, "\"function\"");
  const char *scope = fn ? fn : obj;
  tlen = 0;
  if (extract_json_string_value(scope, "\"name\"", tmp, sizeof(tmp), &tlen)) {
    if (tlen >= TC_NAME_MAX)
      tlen = TC_NAME_MAX - 1;
    memcpy(slot->name, tmp, tlen);
    slot->name[tlen] = 0;
  }
  tlen = 0;
  if (extract_json_string_value(scope, "\"arguments\"", tmp, sizeof(tmp),
                                &tlen)) {
    slot_append_args(slot, tmp, tlen);
  }
}

static void parse_tool_calls(sse_acc_t *a, const char *payload) {
  const char *key = find_json_string_key(payload, "\"tool_calls\"");
  if (!key)
    return;
  const char *p = strchr(key, '[');
  if (!p)
    return;
  p++;
  while (*p) {
    while (*p == ' ' || *p == '\n' || *p == '\r' || *p == '\t' || *p == ',')
      p++;
    if (*p == ']')
      break;
    if (*p != '{')
      break;
    /* find matching close brace for this object */
    int depth = 0;
    int in_str = 0;
    const char *start = p;
    const char *end = NULL;
    for (; *p; p++) {
      if (in_str) {
        if (*p == '\\' && p[1]) {
          p++;
          continue;
        }
        if (*p == '"')
          in_str = 0;
        continue;
      }
      if (*p == '"') {
        in_str = 1;
        continue;
      }
      if (*p == '{')
        depth++;
      else if (*p == '}') {
        depth--;
        if (depth == 0) {
          end = p;
          p++;
          break;
        }
      }
    }
    if (!end)
      break;
    size_t n = (size_t)(end - start + 1);
    if (n >= 8192)
      n = 8191;
    char obj[8192];
    memcpy(obj, start, n);
    obj[n] = 0;
    parse_one_tool_call_obj(a, obj);
  }
}

static void flush_native_tools(sse_acc_t *a) {
  if (!a->on_native_tool)
    return;
  for (int i = 0; i < TC_SLOTS; i++) {
    tc_slot_t *s = &a->slots[i];
    if (!s->used || !s->name[0] || !s->args || !s->args_len)
      continue;
    if (!json_brace_balanced(s->args))
      continue;
    const char *id = s->id[0] ? s->id : "";
    a->on_native_tool(id, s->name, s->args, a->userdata);
    /* Mark consumed so end-of-stream flush is idempotent after [DONE]. */
    free(s->args);
    s->args = NULL;
    s->args_len = 0;
    s->used = 0;
  }
}

static void free_slots(sse_acc_t *a) {
  for (int i = 0; i < TC_SLOTS; i++) {
    free(a->slots[i].args);
    a->slots[i].args = NULL;
    a->slots[i].used = 0;
  }
}

static void emit_data_payload(sse_acc_t *a, const char *payload) {
  /* OpenAI SSE: data: {...}  or data: [DONE] */
  if (strncmp(payload, "[DONE]", 6) == 0) {
    if (a->in_reason) {
      emit_raw(a, "</think>", 8);
      a->in_reason = 0;
    }
    flush_native_tools(a);
    return;
  }

  char out[8192];
  size_t olen = 0;

  /* MiniMax reasoning_split: route reasoning_content through <think> so the
   * agent UI can show it, without polluting the final answer. */
  if (extract_json_string_value(payload, "\"reasoning_content\"", out,
                                sizeof(out), &olen)) {
    if (!a->in_reason) {
      emit_raw(a, "<think>", 7);
      a->in_reason = 1;
    }
    emit_raw(a, out, olen);
  }

  olen = 0;
  if (extract_json_string_value(payload, "\"content\"", out, sizeof(out),
                                &olen)) {
    if (a->in_reason) {
      emit_raw(a, "</think>", 8);
      a->in_reason = 0;
    }
    emit_raw(a, out, olen);
  }

  parse_tool_calls(a, payload);
}

static void process_sse_chunk(sse_acc_t *a, const char *data, size_t n) {
  for (size_t i = 0; i < n; i++) {
    if (a->len + 1 >= a->cap) {
      size_t nc = a->cap ? a->cap * 2 : 4096;
      char *nb = realloc(a->buf, nc);
      if (!nb)
        return;
      a->buf = nb;
      a->cap = nc;
    }
    a->buf[a->len++] = data[i];
    if (data[i] == '\n') {
      a->buf[a->len] = 0;
      /* trim CR */
      if (a->len >= 2 && a->buf[a->len - 2] == '\r')
        a->buf[a->len - 2] = 0;
      else
        a->buf[a->len - 1] = 0;
      if (strncmp(a->buf, "data:", 5) == 0) {
        const char *p = a->buf + 5;
        while (*p == ' ')
          p++;
        emit_data_payload(a, p);
      }
      a->len = 0;
    }
  }
}

static size_t write_cb(char *ptr, size_t size, size_t nmemb, void *userdata) {
  sse_acc_t *a = userdata;
  size_t n = size * nmemb;
  process_sse_chunk(a, ptr, n);
  return n;
}

static int xferinfo_cb(void *clientp, curl_off_t dltotal, curl_off_t dlnow,
                       curl_off_t ultotal, curl_off_t ulnow) {
  (void)dltotal;
  (void)dlnow;
  (void)ultotal;
  (void)ulnow;
  volatile int *flag = (volatile int *)clientp;
  if (flag && *flag != 0)
    return 1; /* abort */
  return 0;
}

static int json_escape_append(char **buf, size_t *len, size_t *cap,
                              const char *s) {
  if (!buf || !len || !cap)
    return -1;
  if (!*buf || *cap == 0) {
    *cap = 4096;
    *buf = malloc(*cap);
    if (!*buf)
      return -1;
    *len = 0;
    (*buf)[0] = 0;
  }
  if (!s)
    s = "";
  for (; *s; s++) {
    char tmp[8];
    const char *piece = tmp;
    size_t pn = 1;
    unsigned char c = (unsigned char)*s;
    if (c == '"' || c == '\\') {
      tmp[0] = '\\';
      tmp[1] = (char)c;
      pn = 2;
    } else if (c == '\n') {
      piece = "\\n";
      pn = 2;
    } else if (c == '\r') {
      piece = "\\r";
      pn = 2;
    } else if (c == '\t') {
      piece = "\\t";
      pn = 2;
    } else if (c < 0x20) {
      snprintf(tmp, sizeof(tmp), "\\u%04x", c);
      pn = 6;
    } else {
      tmp[0] = (char)c;
      pn = 1;
    }
    if (*len + pn + 1 > *cap) {
      size_t nc = *cap ? *cap * 2 : 4096;
      while (nc < *len + pn + 1)
        nc *= 2;
      char *nb = realloc(*buf, nc);
      if (!nb)
        return -1;
      *buf = nb;
      *cap = nc;
    }
    memcpy(*buf + *len, piece, pn);
    *len += pn;
  }
  (*buf)[*len] = 0;
  return 0;
}

static int body_ensure(char **body, size_t *blen, size_t *bcap, size_t need) {
  if (*blen + need + 1 <= *bcap)
    return 0;
  size_t nc = *bcap ? *bcap * 2 : 4096;
  while (nc < *blen + need + 1)
    nc *= 2;
  char *nb = realloc(*body, nc);
  if (!nb)
    return -1;
  *body = nb;
  *bcap = nc;
  return 0;
}

static int body_append_raw(char **body, size_t *blen, size_t *bcap,
                           const char *s) {
  size_t n = strlen(s);
  if (body_ensure(body, blen, bcap, n) != 0)
    return -1;
  memcpy(*body + *blen, s, n);
  *blen += n;
  (*body)[*blen] = 0;
  return 0;
}

static const char *role_api(cst_role_t r) {
  switch (r) {
  case CST_ROLE_SYSTEM:
    return "system";
  case CST_ROLE_USER:
    return "user";
  case CST_ROLE_ASSISTANT:
    return "assistant";
  case CST_ROLE_TOOL:
    return "tool";
  }
  return "user";
}

static const char *TOOLS_JSON =
    ",\"tools\":["
    "{\"type\":\"function\",\"function\":{\"name\":\"read_file\","
    "\"description\":\"Read a text file under the workspace\","
    "\"parameters\":{\"type\":\"object\",\"properties\":{"
    "\"path\":{\"type\":\"string\"}},\"required\":[\"path\"]}}},"
    "{\"type\":\"function\",\"function\":{\"name\":\"run_shell\","
    "\"description\":\"Run a shell command in the workspace root\","
    "\"parameters\":{\"type\":\"object\",\"properties\":{"
    "\"cmd\":{\"type\":\"string\"}},\"required\":[\"cmd\"]}}},"
    "{\"type\":\"function\",\"function\":{\"name\":\"write_file\","
    "\"description\":\"Write full file contents under the workspace\","
    "\"parameters\":{\"type\":\"object\",\"properties\":{"
    "\"path\":{\"type\":\"string\"},\"content\":{\"type\":\"string\"}},"
    "\"required\":[\"path\",\"content\"]}}},"
    "{\"type\":\"function\",\"function\":{\"name\":\"apply_patch\","
    "\"description\":\"Apply a unified diff or Update File patch\","
    "\"parameters\":{\"type\":\"object\",\"properties\":{"
    "\"patch\":{\"type\":\"string\"}},\"required\":[\"patch\"]}}}"
    "]";

int cst_llm_chat_stream(const cst_llm_request_t *req, char *err, size_t err_sz) {
  if (err && err_sz)
    err[0] = 0;
  if (!req || !req->api_base || !req->model) {
    if (err)
      snprintf(err, err_sz, "bad request");
    return -1;
  }

  char *body = NULL;
  size_t blen = 0;
  size_t bcap = 8192;
  body = malloc(bcap);
  if (!body)
    return -1;
  /* MiniMax: keep think+answer in content (no reasoning_split). Split mode
   * often leaves content empty while the useful text is only in
   * reasoning_content, which made the agent exit with blank turns. */
  blen = (size_t)snprintf(body, bcap,
                          "{\"model\":\"%s\",\"stream\":true,\"messages\":[",
                          req->model);

  for (int i = 0; i < req->nmsgs; i++) {
    const cst_message_t *m = &req->msgs[i];
    char prefix[256];
    int first = i ? 0 : 1;
    if (m->role == CST_ROLE_ASSISTANT && m->tool_calls_json &&
        m->tool_calls_json[0]) {
      /* OpenCode/OpenAI: tool-only assistant uses content:null */
      int empty =
          !m->content || !m->content[0] ||
          (m->content[0] == '\0');
      if (empty) {
        snprintf(prefix, sizeof(prefix),
                 "%s{\"role\":\"assistant\",\"content\":null,\"tool_calls\":",
                 first ? "" : ",");
        if (body_append_raw(&body, &blen, &bcap, prefix) != 0 ||
            body_append_raw(&body, &blen, &bcap, m->tool_calls_json) != 0 ||
            body_append_raw(&body, &blen, &bcap, "}") != 0) {
          free(body);
          return -1;
        }
      } else {
        snprintf(prefix, sizeof(prefix),
                 "%s{\"role\":\"assistant\",\"content\":\"", first ? "" : ",");
        size_t pl = strlen(prefix);
        if (blen + pl + 8 > bcap) {
          bcap = (blen + pl) * 2 + 1024;
          body = realloc(body, bcap);
          if (!body)
            return -1;
        }
        memcpy(body + blen, prefix, pl);
        blen += pl;
        body[blen] = 0;
        if (json_escape_append(&body, &blen, &bcap, m->content) != 0) {
          free(body);
          return -1;
        }
        if (body_append_raw(&body, &blen, &bcap, "\",\"tool_calls\":") != 0 ||
            body_append_raw(&body, &blen, &bcap, m->tool_calls_json) != 0 ||
            body_append_raw(&body, &blen, &bcap, "}") != 0) {
          free(body);
          return -1;
        }
      }
      continue;
    }
    if (m->role == CST_ROLE_TOOL && m->tool_call_id[0]) {
      /* Include name when known (OpenAI-compatible gateways accept it). */
      if (m->tool_name[0]) {
        snprintf(prefix, sizeof(prefix),
                 "%s{\"role\":\"%s\",\"tool_call_id\":\"%s\",\"name\":\"%s\","
                 "\"content\":\"",
                 first ? "" : ",", role_api(m->role), m->tool_call_id,
                 m->tool_name);
      } else {
        snprintf(prefix, sizeof(prefix),
                 "%s{\"role\":\"%s\",\"tool_call_id\":\"%s\",\"content\":\"",
                 first ? "" : ",", role_api(m->role), m->tool_call_id);
      }
    } else {
      snprintf(prefix, sizeof(prefix), "%s{\"role\":\"%s\",\"content\":\"",
               first ? "" : ",", role_api(m->role));
    }
    size_t pl = strlen(prefix);
    if (blen + pl + 8 > bcap) {
      bcap = (blen + pl) * 2 + 1024;
      body = realloc(body, bcap);
      if (!body)
        return -1;
    }
    memcpy(body + blen, prefix, pl);
    blen += pl;
    body[blen] = 0;
    /* Keep full tool output in session/TUI; shrink only the wire payload so
     * large write diffs don't blow MiniMax/JSON limits. */
    {
      const char *raw = m->content ? m->content : "";
      char *scratch = NULL;
      const char *send = raw;
      enum { API_TOOL_CAP = 3500 };
      if (m->role == CST_ROLE_TOOL && strlen(raw) > API_TOOL_CAP) {
        scratch = malloc(API_TOOL_CAP + 80);
        if (scratch) {
          snprintf(scratch, API_TOOL_CAP + 80,
                   "%.*s\n…(truncated for model; full diff in UI/session)",
                   API_TOOL_CAP, raw);
          send = scratch;
        }
      }
      if (json_escape_append(&body, &blen, &bcap, send) != 0) {
        free(scratch);
        free(body);
        return -1;
      }
      free(scratch);
    }
    if (blen + 3 > bcap) {
      bcap *= 2;
      body = realloc(body, bcap);
      if (!body)
        return -1;
    }
    body[blen++] = '"';
    body[blen++] = '}';
    body[blen] = 0;
  }
  if (body_append_raw(&body, &blen, &bcap, "]") != 0) {
    free(body);
    return -1;
  }
  if (req->enable_tools) {
    if (body_append_raw(&body, &blen, &bcap, TOOLS_JSON) != 0) {
      free(body);
      return -1;
    }
    if (req->tool_choice_none) {
      if (body_append_raw(&body, &blen, &bcap, ",\"tool_choice\":\"none\"") !=
          0) {
        free(body);
        return -1;
      }
    } else {
      if (body_append_raw(&body, &blen, &bcap, ",\"tool_choice\":\"auto\"") !=
          0) {
        free(body);
        return -1;
      }
    }
  }
  if (body_append_raw(&body, &blen, &bcap, "}") != 0) {
    free(body);
    return -1;
  }

  char url[1024];
  snprintf(url, sizeof(url), "%s/chat/completions", req->api_base);

  CURL *curl = curl_easy_init();
  if (!curl) {
    free(body);
    if (err)
      snprintf(err, err_sz, "curl init failed");
    return -1;
  }
  sse_acc_t acc = {0};
  acc.on_delta = req->on_delta;
  acc.on_native_tool = req->on_native_tool;
  acc.userdata = req->userdata;

  struct curl_slist *hdrs = NULL;
  hdrs = curl_slist_append(hdrs, "Content-Type: application/json");
  char auth[600];
  if (req->api_key && req->api_key[0]) {
    snprintf(auth, sizeof(auth), "Authorization: Bearer %s", req->api_key);
    hdrs = curl_slist_append(hdrs, auth);
  }

  curl_easy_setopt(curl, CURLOPT_URL, url);
  curl_easy_setopt(curl, CURLOPT_HTTPHEADER, hdrs);
  curl_easy_setopt(curl, CURLOPT_POSTFIELDS, body);
  curl_easy_setopt(curl, CURLOPT_WRITEFUNCTION, write_cb);
  curl_easy_setopt(curl, CURLOPT_WRITEDATA, &acc);
  curl_easy_setopt(curl, CURLOPT_TIMEOUT, 300L);

  if (req->cancel_flag) {
    curl_easy_setopt(curl, CURLOPT_NOPROGRESS, 0L);
    curl_easy_setopt(curl, CURLOPT_XFERINFOFUNCTION, xferinfo_cb);
    curl_easy_setopt(curl, CURLOPT_XFERINFODATA, (void *)req->cancel_flag);
  }

  CURLcode rc = curl_easy_perform(curl);
  long http = 0;
  curl_easy_getinfo(curl, CURLINFO_RESPONSE_CODE, &http);
  if (acc.in_reason) {
    emit_raw(&acc, "</think>", 8);
    acc.in_reason = 0;
  }
  /* Stream ended without [DONE] — still flush completed native tools. */
  if (rc == CURLE_OK)
    flush_native_tools(&acc);
  curl_slist_free_all(hdrs);
  curl_easy_cleanup(curl);

  if (rc != CURLE_OK) {
    if (err) {
      if (rc == CURLE_ABORTED_BY_CALLBACK)
        snprintf(err, err_sz, "cancelled");
      else
        snprintf(err, err_sz, "curl: %s", curl_easy_strerror(rc));
    }
    free_slots(&acc);
    free(acc.buf);
    free(body);
    return -1;
  }
  if (http >= 400) {
    if (err) {
      /* show a snippet of body if present */
      if (acc.buf && acc.len) {
        acc.buf[acc.len < 200 ? acc.len : 200] = 0;
        snprintf(err, err_sz, "http %ld: %.180s", http, acc.buf);
      } else {
        snprintf(err, err_sz, "http %ld", http);
      }
    }
    free_slots(&acc);
    free(acc.buf);
    free(body);
    return -1;
  }
  free_slots(&acc);
  free(acc.buf);
  free(body);
  return 0;
}