/* llm.cpp — libcurl HTTP client, Fibonacci backoff,
 * streaming SSE parser and the llm() dispatch function. */
#include "bantam.hpp"
#include <array>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <ctime>
#include <curl/curl.h>
#include <memory>
#include <unistd.h>

/* ----------------------------- memory callback ------------------------ */
static size_t mem_cb(char *ptr, size_t sz, size_t nm, void *ud) {
  std::string *m = (std::string *)ud;
  size_t total = sz * nm;
  m->append((const char *)ptr, total);
  return total;
}

static int xferinfo_cb(void *clientp, curl_off_t, curl_off_t, curl_off_t, curl_off_t) {
  volatile int *interrupted = (volatile int *)clientp;
  return (interrupted && *interrupted) ? 1 : 0;
}

/* ----------------------------- curl handle & headers ------------------ */
static CURL *g_shared_curl = nullptr;

static CURL *get_shared_curl() {
  if (!g_shared_curl) {
    g_shared_curl = curl_easy_init();
  } else {
    curl_easy_reset(g_shared_curl);
  }
  return g_shared_curl;
}

void http_cleanup() {
  if (g_shared_curl) {
    curl_easy_cleanup(g_shared_curl);
    g_shared_curl = nullptr;
  }
}

struct CurlHeaders {
  curl_slist *list = nullptr;
  CurlHeaders() = default;
  ~CurlHeaders() { if (list) curl_slist_free_all(list); }
  CurlHeaders(const CurlHeaders &) = delete;
  CurlHeaders &operator=(const CurlHeaders &) = delete;
  void append(std::string_view s) {
    std::string tmp(s);
    curl_slist *next = curl_slist_append(list, tmp.c_str());
    if (next) list = next;
  }
};

/* ----------------------------- proxy + headers ------------------------ */
static void set_proxy(CURL *c) {
  const char *socks = std::getenv("SOCKS_PROXY");
  if (!socks || !*socks) socks = std::getenv("socks_proxy");
  if (socks && *socks) {
    std::string buf;
    if (std::strstr(socks, "://")) buf = socks;
    else buf = "socks5://" + std::string(socks);
    curl_easy_setopt(c, CURLOPT_PROXY, buf.c_str());
  }
}

static void apply_common_curl_opts(CURL *c, volatile int *interrupted) {
  curl_easy_setopt(c, CURLOPT_NOSIGNAL, 1L);
  curl_easy_setopt(c, CURLOPT_NOPROGRESS, 0L);
  curl_easy_setopt(c, CURLOPT_XFERINFOFUNCTION, xferinfo_cb);
  curl_easy_setopt(c, CURLOPT_XFERINFODATA, (void *)interrupted);
  curl_easy_setopt(c, CURLOPT_TCP_KEEPALIVE, 1L);
  curl_easy_setopt(c, CURLOPT_TCP_KEEPIDLE, 30L);
  curl_easy_setopt(c, CURLOPT_TCP_KEEPINTVL, 15L);
  curl_easy_setopt(c, CURLOPT_USERAGENT, "Mozilla/5.0 Bantam/nightly-cpp");
  set_proxy(c);
}

static void apply_llm_headers(CurlHeaders &h, const Cfg &cfg) {
  h.append("Content-Type: application/json");
  /* Suppress Expect: 100-continue so large POST bodies are sent immediately
   * without waiting for a 100 Continue round-trip or hanging behind proxies. */
  h.append("Expect:");
  if (!cfg.api_key.empty() && cfg.api_key != "-") {
    h.append(sfmt("Authorization: Bearer %s", cfg.api_key.c_str()));
  }
}

/* ----------------------------- http_get ------------------------------- */
int http_get(Cfg &cfg, std::string_view url, std::string &body, long &code,
             std::string &errmsg) {
  CURL *c = get_shared_curl();
  if (!c) { errmsg = "curl init failed"; return -1; }
  std::string url_s(url);
  CurlHeaders headers;
  if (!cfg.api_key.empty() && cfg.api_key != "-") {
    headers.append(sfmt("Authorization: Bearer %s", cfg.api_key.c_str()));
  }
  apply_common_curl_opts(c, &g_interrupted);
  curl_easy_setopt(c, CURLOPT_URL, url_s.c_str());
  curl_easy_setopt(c, CURLOPT_HTTPGET, 1L);
  if (headers.list) curl_easy_setopt(c, CURLOPT_HTTPHEADER, headers.list);
  curl_easy_setopt(c, CURLOPT_WRITEFUNCTION, mem_cb);
  curl_easy_setopt(c, CURLOPT_WRITEDATA, &body);
  long cfg_t = cfg.timeout > 0 ? cfg.timeout : 15;
  long conn_t = (cfg_t < 10) ? cfg_t : 10;
  long tot_t = (cfg_t < 15) ? cfg_t : 15;
  curl_easy_setopt(c, CURLOPT_CONNECTTIMEOUT, conn_t);
  curl_easy_setopt(c, CURLOPT_TIMEOUT, tot_t);
  curl_easy_setopt(c, CURLOPT_LOW_SPEED_LIMIT, 1L);
  curl_easy_setopt(c, CURLOPT_LOW_SPEED_TIME, tot_t);
  CURLcode res = curl_easy_perform(c);
  curl_easy_getinfo(c, CURLINFO_RESPONSE_CODE, &code);
  if (res != CURLE_OK) {
    errmsg = (res == CURLE_ABORTED_BY_CALLBACK || g_interrupted)
               ? "interrupted"
               : curl_easy_strerror(res);
    return -1;
  }
  return 0;
}

/* ----------------------------- token estimate ------------------------- */
static int est_tokens_msgs(const Messages &msgs) {
  int chars = 0;
  for (auto &m : msgs.items) {
    if (!m.content.empty()) chars += (int)m.content.size();
    if (!m.reasoning.empty()) chars += (int)m.reasoning.size();
    for (auto &tc : m.calls) {
      if (!tc.name.empty()) chars += (int)tc.name.size();
      if (!tc.arguments.empty()) chars += (int)tc.arguments.size();
    }
  }
  if (chars == 0) return 0;
  int t = chars / 4;
  return t > 0 ? t : 1;
}
static int est_tokens_msg(const Message &m) {
  int chars = 0;
  if (!m.content.empty()) chars += (int)m.content.size();
  if (!m.reasoning.empty()) chars += (int)m.reasoning.size();
  for (auto &tc : m.calls) {
    if (!tc.name.empty()) chars += (int)tc.name.size();
    if (!tc.arguments.empty()) chars += (int)tc.arguments.size();
  }
  if (chars == 0) return 0;
  int t = chars / 4;
  return t > 0 ? t : 1;
}

/* ----------------------------- streaming parser ----------------------- */
struct StreamCtx {
  std::string content, reas, line_buf, buf, raw;
  std::vector<std::string> tbl;
  Usage last_usage;
  std::vector<ToolCall> tcs;
  std::vector<std::string> tcs_filt_args;
  std::vector<size_t> tcs_filt_pos;
  std::vector<int> tcs_idx;
  size_t live_printed_len = 0;
  int tool_calls_shown = 0;
  int tty = 0;
  int live_idx = -1;
  MdState md;
  int in_reasoning = 0;
  int done = 0;
  int erase_pending = 0;
  volatile int *interrupted = nullptr;
};

static void advance_filtered_args(StreamCtx *sc, int slot, bool flush_all) {
  const std::string &raw_args = sc->tcs[slot].arguments;
  size_t pos = sc->tcs_filt_pos[slot];
  if (pos >= raw_args.size()) return;
  std::string_view rem = std::string_view(raw_args).substr(pos);
  if (!flush_all) {
    /* Hold back an incomplete trailing multi-byte UTF-8 sequence until the
     * next chunk arrives. */
    size_t back = 0;
    while (back < 3 && back < rem.size()) {
      unsigned char b = (unsigned char)rem[rem.size() - 1 - back];
      if ((b & 0x80) == 0) break; /* ASCII */
      if ((b & 0xC0) == 0xC0) {   /* UTF-8 lead byte */
        size_t need = ((b & 0xE0) == 0xC0) ? 2 : ((b & 0xF0) == 0xE0) ? 3 : 4;
        if (back + 1 < need) rem.remove_suffix(back + 1);
        break;
      }
      back++;
    }
  }
  if (!rem.empty()) {
    sc->tcs_filt_args[slot] += filter_text(rem);
    sc->tcs_filt_pos[slot] = pos + rem.size();
  }
}

static void stream_upsert_tc(StreamCtx *sc, int idx, std::string_view id,
                             std::string_view name, std::string_view args) {
  int slot = -1;
  for (size_t i = 0; i < sc->tcs_idx.size(); i++)
    if (sc->tcs_idx[i] == idx) { slot = (int)i; break; }
  if (slot < 0) {
    slot = (int)sc->tcs.size();
    sc->tcs.emplace_back();
    sc->tcs_filt_args.emplace_back();
    sc->tcs_filt_pos.push_back(0);
    sc->tcs_idx.push_back(idx);
  }
  if (!id.empty()) sc->tcs[slot].id = std::string(id);
  if (!name.empty()) sc->tcs[slot].name += std::string(name);
  if (!args.empty()) {
    sc->tcs[slot].arguments += std::string(args);
    advance_filtered_args(sc, slot, false);
  }
}

static void print_live_delta(StreamCtx *sc, int slot, bool closed) {
  if (closed) advance_filtered_args(sc, slot, true);
  const ToolCall &t = sc->tcs[slot];
  const std::string &fa = sc->tcs_filt_args[slot];
  bool has_args = closed || !t.arguments.empty();
  size_t prefix_len = 12 + t.name.size(); /* "[tool call: " + name */
  size_t total_len = prefix_len + (has_args ? (1 + fa.size() + (closed ? 2 : 0)) : 0);
  size_t cur = sc->live_printed_len;
  if (total_len <= cur) return;

  std::string delta;
  delta.reserve(total_len - cur);
  auto emit_part = [&](std::string_view part, size_t part_start) {
    size_t part_end = part_start + part.size();
    if (part_end <= cur) return;
    size_t off = (cur > part_start) ? (cur - part_start) : 0;
    delta.append(part.substr(off));
  };
  size_t off = 0;
  emit_part("[tool call: ", off); off += 12;
  emit_part(t.name, off);         off += t.name.size();
  if (has_args) {
    emit_part("(", off);          off += 1;
    emit_part(fa, off);           off += fa.size();
    if (closed) {
      emit_part(")]", off);       off += 2;
    }
  }
  cprint(delta, std::array{33});
  sc->live_printed_len = total_len;
}

static void finish_live(StreamCtx *sc, bool closed) {
  if (sc->live_idx < 0) return;
  print_live_delta(sc, sc->live_idx, closed);
  sc->live_printed_len = 0;
  sc->live_idx = -1;
  /* The tool call marker must always terminate its own line so the
   * following [tool result: ...] header is separated from it by a
   * newline (in streaming mode the marker is shown here and not
   * reprinted by the caller). */
  printf("\n");
  fflush(stdout);
}
static void step_live(StreamCtx *sc) {
  if (!sc->tty || sc->tcs.empty()) return;
  int cur = (int)sc->tcs.size() - 1;
  if (sc->live_idx != -1 && sc->live_idx != cur) finish_live(sc, true);
  if (sc->live_idx == -1) { sc->live_idx = cur; sc->live_printed_len = 0; }
  print_live_delta(sc, cur, false);
  sc->tool_calls_shown = 1;
}

static void flush_tbl(StreamCtx *sc) {
  if (!sc->tbl.empty()) {
    render_table_print(sc->tbl);
    sc->tbl.clear();
  }
}

static void stream_maybe_erase(StreamCtx *sc) {
  if (sc->erase_pending) {
    if (g_col) printf("\r\033[K");
    sc->erase_pending = 0;
  }
}

static void stream_handle_data(StreamCtx *sc, std::string_view data) {
  if (data == "[DONE]") { sc->done = 1; return; }
  Json d_mut;
  if (!Json::parse(data, d_mut) || !d_mut.is_object()) return;
  const Json &d = d_mut;
  const Json &u = d["usage"];
  if (u.is_object()) {
    int pt = u["prompt_tokens"].as_int();
    int tt = u["total_tokens"].as_int();
    if (pt > 0 || tt > 0) {
      sc->last_usage.prompt_tokens = pt;
      sc->last_usage.completion_tokens = u["completion_tokens"].as_int();
      sc->last_usage.total_tokens = tt;
      const Json &ptd = u["prompt_tokens_details"];
      if (ptd.is_object()) {
        int ct_val = ptd["cached_tokens"].as_int();
        if (ct_val > 0) sc->last_usage.cached_tokens = ct_val;
      }
      if (sc->last_usage.cached_tokens == 0) {
        int ct_val = u["cached_tokens"].as_int();
        if (ct_val > 0) sc->last_usage.cached_tokens = ct_val;
      }
    }
  }
  if (d["model"].is_string()) {
    sc->last_usage.model = d["model"].str();
  }
  const Json &choices = d["choices"];
  if (!choices.is_array() || choices.empty()) return;
  const Json &ch = choices.at(0);
  if (!ch.is_object()) return;
  const Json &delta = ch["delta"];
  if (!delta.is_object()) return;

  std::string rc = delta["reasoning_content"].str();
  if (rc.empty()) rc = delta["reasoning"].str();
  if (rc.empty()) rc = delta["thought"].str();
  if (!rc.empty()) {
    stream_maybe_erase(sc);
    if (!sc->in_reasoning) {
      if (!sc->content.empty()) {
        flush_tbl(sc);
        if (!sc->line_buf.empty()) { render_md_line_print(sc->line_buf, sc->md); sc->line_buf.clear(); }
        printf("\n");
      }
      cprint("--- reasoning start ---\n", std::array{36});
      sc->in_reasoning = true;
    }
    cprint(rc, std::array{2});
    sc->reas += rc;
  }
  std::string ct = delta["content"].str();
  if (!ct.empty()) {
    stream_maybe_erase(sc);
    if (sc->in_reasoning) { cprint("\n--- reasoning end ---\n", std::array{36}); sc->in_reasoning = false; }
    sc->content += ct;
    sc->line_buf += ct;
    size_t lpos = 0;
    while (true) {
      auto nl = sc->line_buf.find('\n', lpos);
      if (nl == std::string::npos) break;
      std::string_view cur_sv = std::string_view(sc->line_buf).substr(lpos, nl - lpos);
      lpos = nl + 1;
      /* Inside a code block every line is rendered verbatim as code; only
       * outside a code block do table-shaped lines get accumulated for
       * box-drawing table rendering. */
      if (sc->md.in_code || !is_table_line(cur_sv)) {
        flush_tbl(sc);
        render_md_line_print(cur_sv, sc->md);
      } else {
        sc->tbl.emplace_back(cur_sv);
      }
    }
    if (lpos > 0) sc->line_buf.erase(0, lpos);
  }
  const Json &tcs = delta["tool_calls"];
  if (tcs.is_array()) {
    stream_maybe_erase(sc);
    if (sc->in_reasoning) { cprint("\n--- reasoning end ---\n", std::array{36}); sc->in_reasoning = false; }
    if (sc->live_idx == -1) {
      flush_tbl(sc);
      if (!sc->line_buf.empty()) { render_md_line_print(sc->line_buf, sc->md); sc->line_buf.clear(); }
    }
    for (auto &tc : tcs.aval) {
      if (!tc.is_object()) continue;
      int idx = tc["index"].as_int();
      std::string id = tc["id"].str();
      std::string name, args;
      const Json &fn = tc["function"];
      if (fn.is_object()) {
        name = fn["name"].str();
        args = fn["arguments"].str();
      }
      stream_upsert_tc(sc, idx, id, name, args);
    }
    step_live(sc);
  }
}

static int stream_feed(StreamCtx *sc, const char *data, size_t len,
                       volatile int *interrupted) {
  sc->interrupted = interrupted;
  sc->buf.append(data, len);
  size_t pos = 0;
  while (true) {
    auto nl = sc->buf.find('\n', pos);
    if (nl == std::string::npos) break;
    std::string_view line = std::string_view(sc->buf).substr(pos, nl - pos);
    pos = nl + 1;
    while (!line.empty() && (line.back() == '\r' || line.back() == ' ' || line.back() == '\t'))
      line.remove_suffix(1);
    if (line.starts_with("data:")) {
      line.remove_prefix(5);
      while (!line.empty() && line.front() == ' ') line.remove_prefix(1);
      stream_handle_data(sc, line);
    }
    if (sc->done) break;
    if (*interrupted) {
      sc->buf.erase(0, pos);
      return -1;
    }
  }
  if (pos > 0) sc->buf.erase(0, pos);
  return 0;
}

static void stream_finish(StreamCtx *sc) {
  flush_tbl(sc);
  if (!sc->line_buf.empty()) {
    std::string cur = sc->line_buf;
    if (!sc->md.in_code && is_table_line(cur)) {
      sc->tbl.push_back(cur);
      flush_tbl(sc);
    } else {
      render_md_line_print(cur, sc->md);
    }
    sc->line_buf.clear();
  }
  if (sc->live_idx >= 0) finish_live(sc, true);
  if (sc->in_reasoning) { cprint("\n--- reasoning end ---\n", std::array{36}); sc->in_reasoning = false; }
}

static void stream_result(StreamCtx *sc, Message &out, Usage &usage) {
  out = Message{};
  out.role = "assistant";
  out.calls_shown = sc->tool_calls_shown;
  if (!sc->content.empty()) out.content = std::move(sc->content);
  if (!sc->reas.empty()) out.reasoning = std::move(sc->reas);
  for (auto &t : sc->tcs) out.calls.push_back(std::move(t));
  usage = sc->last_usage;
  if (usage.total_tokens == 0 && usage.prompt_tokens > 0)
    usage.total_tokens = usage.prompt_tokens + usage.completion_tokens;
}

/* ----------------------------- streaming write cb --------------------- */
static size_t stream_write_cb(char *ptr, size_t sz, size_t nm, void *ud) {
  StreamCtx *sc = (StreamCtx *)ud;
  size_t total = sz * nm;
  if (*sc->interrupted) return 0;
  /* Cap raw error buffer at 4096 bytes; http_err_body only uses 4000 bytes. */
  if (sc->raw.size() < 4096) {
    size_t rem = 4096 - sc->raw.size();
    sc->raw.append((const char *)ptr, total < rem ? total : rem);
  }
  int r = stream_feed(sc, ptr, total, sc->interrupted);
  if (r != 0 || *sc->interrupted) return 0;
  return total;
}

/* ----------------------------- llm() ---------------------------------- */
static void sleep_interruptible(int secs, volatile int *interrupted) {
  struct timespec ts{0, 100000000L};
  for (int i = 0; i < secs * 10; i++) {
    if (*interrupted) return;
    nanosleep(&ts, nullptr);
  }
}

static std::string http_err_body(const StreamCtx *sc, const std::string *m, long code) {
  std::string_view b = sc ? std::string_view(sc->raw) : (m ? std::string_view(*m) : "");
  b = utf8_truncate(b, 4000);
  return sfmt("HTTP %ld: %.*s", code, (int)b.size(), b.data());
}

int llm(Cfg &cfg, Messages &msgs, bool use_tools, Message &out, Usage &usage,
        std::string &errmsg, volatile int *interrupted) {
  out = Message{};
  usage = Usage{};
  if (*interrupted) { errmsg = "interrupted"; return -1; }

  bool stream = (cfg.force_stream >= 0)
                 ? (cfg.force_stream != 0)
                 : (cfg.stream != 0);
  std::string payload = build_request_json(cfg, msgs, use_tools, stream);
  std::string url = cfg.endpoint;
  while (!url.empty() && url.back() == '/') url.pop_back();
  url += "/chat/completions";

  int fib[] = {1, 1, 2, 3, 5, 8, 13, 21, 34};
  int nf = 9;

  std::string pend = cstr("...requesting...", std::array{1, 2});
  CURL *c = get_shared_curl();
  if (!c) { errmsg = "curl init failed"; return -1; }
  CurlHeaders headers;
  apply_llm_headers(headers, cfg);

  apply_common_curl_opts(c, interrupted);
  curl_easy_setopt(c, CURLOPT_URL, url.c_str());
  curl_easy_setopt(c, CURLOPT_POST, 1L);
  curl_easy_setopt(c, CURLOPT_POSTFIELDS, payload.c_str());
  curl_easy_setopt(c, CURLOPT_POSTFIELDSIZE, (long)payload.size());
  curl_easy_setopt(c, CURLOPT_HTTPHEADER, headers.list);
  long t = (long)(cfg.timeout > 0 ? cfg.timeout : 300);
  long conn_t = (t < 20) ? t : 20;
  curl_easy_setopt(c, CURLOPT_CONNECTTIMEOUT, conn_t);
  curl_easy_setopt(c, CURLOPT_LOW_SPEED_LIMIT, 1L);
  curl_easy_setopt(c, CURLOPT_LOW_SPEED_TIME, t);
  if (!stream) curl_easy_setopt(c, CURLOPT_TIMEOUT, t);

  int attempt, rc = -1;
  for (attempt = 0; attempt <= nf; attempt++) {
    if (*interrupted) { errmsg = "interrupted"; rc = -1; break; }
    if (g_col) { printf("\r%s", pend.c_str()); fflush(stdout); }
    else printf("%s\n", pend.c_str());

    StreamCtx sc;
    std::string mbody;
    CURLcode res; long code = 0;
    if (stream) {
      sc.tty = is_terminal(STDOUT_FILENO);
      sc.interrupted = interrupted;
      sc.erase_pending = 1;
      curl_easy_setopt(c, CURLOPT_WRITEFUNCTION, stream_write_cb);
      curl_easy_setopt(c, CURLOPT_WRITEDATA, &sc);
    } else {
      curl_easy_setopt(c, CURLOPT_WRITEFUNCTION, mem_cb);
      curl_easy_setopt(c, CURLOPT_WRITEDATA, &mbody);
    }
    res = curl_easy_perform(c);
    if (*interrupted || res == CURLE_ABORTED_BY_CALLBACK) {
      *interrupted = 1;
      if (g_col && (!stream || sc.erase_pending)) printf("\r\033[K");
      errmsg = "interrupted";
      rc = -1;
      break;
    }
    curl_easy_getinfo(c, CURLINFO_RESPONSE_CODE, &code);
    int transport_err = (res != CURLE_OK);
    int is4xx = (!transport_err && code >= 400 && code < 500 && code != 408 && code != 429);

    if (g_col && !stream) printf("\r\033[K");

    if (!transport_err && code < 400) {
      if (stream) {
        stream_finish(&sc);
        if (sc.erase_pending) { if (g_col) printf("\r\033[K"); sc.erase_pending = 0; }
        stream_result(&sc, out, usage);
      } else {
        parse_nonstream_response(mbody, out, usage);
      }
      rc = 0;
      break;
    }

    /* error path */
    if (stream) {
      if (transport_err) errmsg = curl_easy_strerror(res);
      else errmsg = http_err_body(&sc, nullptr, code);
    } else {
      if (transport_err) errmsg = curl_easy_strerror(res);
      else errmsg = http_err_body(nullptr, &mbody, code);
    }
    if (is4xx) { rc = -1; break; }
    if (attempt < nf) {
      if (g_col) printf("\r\033[K");
      std::string msg = sfmt("[network error: %s, retrying in %ds...]", errmsg.c_str(), fib[attempt]);
      cprint(msg, std::array{31});
      printf("\n");
      sleep_interruptible(fib[attempt], interrupted);
    }
  }

  if (rc == 0) {
    /* The gateway may report zero usage (it does not count tokens). In that
     * case fall back to a rough bytes/4 estimate so the user still sees a
     * number, but never let a bogus estimate override real gateway numbers. */
    if (usage.prompt_tokens == 0) {
      usage.prompt_tokens = est_tokens_msgs(msgs);
      usage.completion_tokens = est_tokens_msg(out);
      usage.total_tokens = usage.prompt_tokens + usage.completion_tokens;
    }
    /* Context size = what the model actually processed this turn
     * (prompt + completion). Prefer real gateway numbers; only fall back to
     * the estimate when none were provided. The old code set this to the
     * prompt-only estimate and ignored both the completion and any real
     * prompt_tokens the gateway returned. */
    int ctx = usage.prompt_tokens + usage.completion_tokens;
    if (ctx > usage.context_tokens) usage.context_tokens = ctx;
  }
  return rc;
}