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

using json = nlohmann::json;

/* ----------------------------- 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;
}

/* ----------------------------- curl RAII ------------------------------ */
struct CurlHandle {
  CURL *c;
  CurlHandle() : c(curl_easy_init()) {}
  ~CurlHandle() { if (c) curl_easy_cleanup(c); }
  CurlHandle(const CurlHandle &) = delete;
  CurlHandle &operator=(const CurlHandle &) = delete;
  explicit operator bool() const { return c != 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_llm_headers(CurlHeaders &h, const Cfg &cfg) {
  h.append("Content-Type: application/json");
  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) {
  CurlHandle curl;
  if (!curl) { errmsg = "curl init failed"; return -1; }
  std::string url_s(url);
  curl_easy_setopt(curl.c, CURLOPT_URL, url_s.c_str());
  curl_easy_setopt(curl.c, CURLOPT_WRITEFUNCTION, mem_cb);
  curl_easy_setopt(curl.c, CURLOPT_WRITEDATA, &body);
  set_proxy(curl.c);
  long t = cfg.timeout > 0 ? cfg.timeout : 30;
  curl_easy_setopt(curl.c, CURLOPT_CONNECTTIMEOUT, t);
  curl_easy_setopt(curl.c, CURLOPT_TIMEOUT, t);
  curl_easy_setopt(curl.c, CURLOPT_USERAGENT, "Mozilla/5.0 Bantam/nightly-cpp");
  CURLcode res = curl_easy_perform(curl.c);
  curl_easy_getinfo(curl.c, CURLINFO_RESPONSE_CODE, &code);
  if (res != CURLE_OK) {
    errmsg = 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, live_printed, buf, raw;
  std::vector<std::string> tbl;
  Usage last_usage;
  std::vector<ToolCall> tcs;
  std::vector<int> tcs_idx;
  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 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_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);
}

static void live_target(StreamCtx *sc, int slot, bool closed, std::string &out) {
  ToolCall &t = sc->tcs[slot];
  str_appendf(out, "[tool call: %s", t.name.c_str());
  if (closed || !t.arguments.empty()) {
    std::string fa = filter_text(t.arguments);
    str_appendf(out, "(%s", fa.c_str());
    if (closed) out += ")]";
  }
}
static void finish_live(StreamCtx *sc, bool closed) {
  if (sc->live_idx < 0) return;
  std::string target;
  live_target(sc, sc->live_idx, closed, target);
  if (target.size() > sc->live_printed.size()) {
    cprint(target.substr(sc->live_printed.size()), std::array{33});
  }
  sc->live_printed.clear();
  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.clear(); }
  std::string target;
  live_target(sc, cur, false, target);
  if (target.size() > sc->live_printed.size()) {
    cprint(target.substr(sc->live_printed.size()), std::array{33});
    sc->live_printed += target.substr(sc->live_printed.size());
  }
  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, const std::string &data) {
  if (data == "[DONE]") { sc->done = 1; return; }
  try {
    json d = json::parse(data);
    if (!d.is_object()) return;
    if (d.contains("usage") && d["usage"].is_object()) {
      auto u = d["usage"];
      int pt = (u.contains("prompt_tokens") && u["prompt_tokens"].is_number_integer())
               ? u["prompt_tokens"].get<int>() : 0;
      int tt = (u.contains("total_tokens") && u["total_tokens"].is_number_integer())
               ? u["total_tokens"].get<int>() : 0;
      if (pt > 0 || tt > 0) {
        sc->last_usage.prompt_tokens = pt;
        sc->last_usage.completion_tokens =
          (u.contains("completion_tokens") && u["completion_tokens"].is_number_integer())
          ? u["completion_tokens"].get<int>() : 0;
        sc->last_usage.total_tokens = tt;
        if (u.contains("prompt_tokens_details") && u["prompt_tokens_details"].is_object()) {
          auto ptd = u["prompt_tokens_details"];
          if (ptd.contains("cached_tokens") && ptd["cached_tokens"].is_number_integer())
            sc->last_usage.cached_tokens = ptd["cached_tokens"].get<int>();
        }
        if (sc->last_usage.cached_tokens == 0 &&
            u.contains("cached_tokens") && u["cached_tokens"].is_number_integer())
          sc->last_usage.cached_tokens = u["cached_tokens"].get<int>();
      }
    }
    if (d.contains("model") && d["model"].is_string())
      sc->last_usage.model = d["model"].get<std::string>();
    if (!d.contains("choices") || !d["choices"].is_array() || d["choices"].empty()) return;
    auto ch = d["choices"][0];
    if (!ch.is_object() || !ch.contains("delta") || !ch["delta"].is_object()) return;
    auto delta = ch["delta"];

    std::string rc;
    if (delta.contains("reasoning_content") && delta["reasoning_content"].is_string())
      rc = delta["reasoning_content"].get<std::string>();
    if (rc.empty() && delta.contains("reasoning") && delta["reasoning"].is_string())
      rc = delta["reasoning"].get<std::string>();
    if (rc.empty() && delta.contains("thought") && delta["thought"].is_string())
      rc = delta["thought"].get<std::string>();
    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;
    if (delta.contains("content") && delta["content"].is_string())
      ct = delta["content"].get<std::string>();
    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;
      while (true) {
        auto nl = sc->line_buf.find('\n');
        if (nl == std::string::npos) break;
        std::string cur = sc->line_buf.substr(0, nl);
        sc->line_buf.erase(0, nl + 1);
        if (!sc->md.in_code && !is_table_line(cur)) {
          flush_tbl(sc);
          render_md_line_print(cur, sc->md);
        } else {
          sc->tbl.push_back(cur);
        }
      }
    }
    if (delta.contains("tool_calls") && delta["tool_calls"].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 : delta["tool_calls"]) {
        if (!tc.is_object()) continue;
        int idx = (tc.contains("index") && tc["index"].is_number_integer())
                  ? tc["index"].get<int>() : 0;
        std::string id, name, args;
        if (tc.contains("id") && tc["id"].is_string()) id = tc["id"].get<std::string>();
        if (tc.contains("function") && tc["function"].is_object()) {
          auto fn = tc["function"];
          if (fn.contains("name") && fn["name"].is_string()) name = fn["name"].get<std::string>();
          if (fn.contains("arguments") && fn["arguments"].is_string()) args = fn["arguments"].get<std::string>();
        }
        stream_upsert_tc(sc, idx, id, name, args);
      }
      step_live(sc);
    }
  } catch (...) {
    return;
  }
}

static int stream_feed(StreamCtx *sc, const char *data, size_t len,
                       volatile int *interrupted) {
  sc->interrupted = interrupted;
  sc->buf.append(data, len);
  while (true) {
    auto nl = sc->buf.find('\n');
    if (nl == std::string::npos) break;
    std::string line = sc->buf.substr(0, nl);
    while (!line.empty() && (line.back() == '\r' || line.back() == ' ' || line.back() == '\t'))
      line.pop_back();
    if (line.rfind("data:", 0) == 0) {
      std::string_view dd = line;
      dd.remove_prefix(5);
      while (!dd.empty() && dd.front() == ' ') dd.remove_prefix(1);
      stream_handle_data(sc, std::string(dd));
    }
    sc->buf.erase(0, nl + 1);
    if (sc->done) break;
    if (*interrupted) return -1;
  }
  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 = sc->content;
  if (!sc->reas.empty()) out.reasoning = sc->reas;
  for (auto &t : sc->tcs) out.calls.push_back(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;
  sc->raw.append((const char *)ptr, total);
  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) {
  for (int i = 0; i < secs * 10; i++) {
    if (*interrupted) return;
    std::this_thread::sleep_for(std::chrono::milliseconds(100));
  }
}

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});
  CurlHandle curl;
  if (!curl) { errmsg = "curl init failed"; return -1; }
  curl_easy_setopt(curl.c, CURLOPT_USERAGENT, "Mozilla/5.0 Bantam/nightly-cpp");
  CurlHeaders headers;
  apply_llm_headers(headers, cfg);

  curl_easy_setopt(curl.c, CURLOPT_URL, url.c_str());
  curl_easy_setopt(curl.c, CURLOPT_POST, 1L);
  curl_easy_setopt(curl.c, CURLOPT_POSTFIELDS, payload.c_str());
  curl_easy_setopt(curl.c, CURLOPT_POSTFIELDSIZE, (long)payload.size());
  curl_easy_setopt(curl.c, CURLOPT_HTTPHEADER, headers.list);
  set_proxy(curl.c);
  curl_easy_setopt(curl.c, CURLOPT_CONNECTTIMEOUT, cfg.timeout > 0 ? cfg.timeout : 30);
  if (!stream) curl_easy_setopt(curl.c, CURLOPT_TIMEOUT, cfg.timeout > 0 ? cfg.timeout : 300);

  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(curl.c, CURLOPT_WRITEFUNCTION, stream_write_cb);
      curl_easy_setopt(curl.c, CURLOPT_WRITEDATA, &sc);
    } else {
      curl_easy_setopt(curl.c, CURLOPT_WRITEFUNCTION, mem_cb);
      curl_easy_setopt(curl.c, CURLOPT_WRITEDATA, &mbody);
    }
    res = curl_easy_perform(curl.c);
    if (*interrupted) {
      if (g_col && (!stream || sc.erase_pending)) printf("\r\033[K");
      errmsg = "interrupted";
      rc = -1;
      break;
    }
    curl_easy_getinfo(curl.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;
}