/* json.cpp — native RFC 8259 JSON parser/serializer, request building,
 * and non-streaming response parsing. */
#include "bantam.hpp"
#include <cerrno>
#include <charconv>
#include <climits>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>

/* ----------------------------- Json methods --------------------------- */
Json::Json() = default;
Json::~Json() = default;
Json::Json(const Json &) = default;
Json::Json(Json &&) noexcept = default;
Json &Json::operator=(const Json &) = default;
Json &Json::operator=(Json &&) noexcept = default;

Json::Json(std::nullptr_t) : type(JType::Null) {}
Json::Json(bool b) : type(JType::Bool), bval(b) {}
Json::Json(int n) : type(JType::Int), ival(n), dval(n) {}
Json::Json(int64_t n) : type(JType::Int), ival(n), dval((double)n) {}
Json::Json(double d) : type(JType::Float), ival((int64_t)d), dval(d) {}
Json::Json(const char *s) : type(JType::Str), sval(s ? s : "") {}
Json::Json(std::string s) : type(JType::Str), sval(std::move(s)) {}
Json::Json(std::string_view s) : type(JType::Str), sval(s) {}

Json Json::object() {
  Json j;
  j.type = JType::Obj;
  return j;
}
Json Json::array() {
  Json j;
  j.type = JType::Arr;
  return j;
}

static const Json kNullJson{};

bool Json::empty() const {
  if (type == JType::Arr) return aval.empty();
  if (type == JType::Obj) return oval.empty();
  if (type == JType::Str) return sval.empty();
  return true;
}

size_t Json::size() const {
  if (type == JType::Arr) return aval.size();
  if (type == JType::Obj) return oval.size();
  if (type == JType::Str) return sval.size();
  return 0;
}

bool Json::contains(std::string_view k) const {
  if (type != JType::Obj) return false;
  for (auto &kv : oval) {
    if (kv.first == k) return true;
  }
  return false;
}

Json &Json::operator[](std::string_view k) {
  if (type != JType::Obj) {
    *this = Json::object();
  }
  for (auto &kv : oval) {
    if (kv.first == k) return kv.second;
  }
  oval.emplace_back(std::string(k), Json{});
  return oval.back().second;
}

const Json &Json::operator[](std::string_view k) const {
  if (type != JType::Obj) return kNullJson;
  for (auto &kv : oval) {
    if (kv.first == k) return kv.second;
  }
  return kNullJson;
}

Json &Json::at(size_t i) {
  return aval[i];
}

const Json &Json::at(size_t i) const {
  if (type != JType::Arr || i >= aval.size()) return kNullJson;
  return aval[i];
}

void Json::push_back(Json v) {
  if (type != JType::Arr) {
    *this = Json::array();
  }
  aval.push_back(std::move(v));
}

std::string Json::str(std::string_view def) const {
  if (type == JType::Str) return sval;
  return std::string(def);
}

int64_t Json::as_int64(int64_t def) const {
  if (type == JType::Int) return ival;
  return def;
}

int Json::as_int(int def) const {
  if (type != JType::Int) return def;
  if (ival > INT_MAX) return INT_MAX;
  if (ival < INT_MIN) return INT_MIN;
  return (int)ival;
}

/* ----------------------------- JSON serializer ------------------------ */
static void dump_str(std::string &out, std::string_view s) {
  out += '"';
  const unsigned char *p = (const unsigned char *)s.data();
  size_t avail = s.size();
  while (avail > 0) {
    /* Fast path: scan contiguous unescaped printable ASCII. */
    size_t run = 0;
    while (run < avail) {
      unsigned char b = p[run];
      if (b >= 0x20 && b < 0x80 && b != '"' && b != '\\') run++;
      else break;
    }
    if (run > 0) {
      out.append((const char *)p, run);
      p += run;
      avail -= run;
      if (avail == 0) break;
    }
    unsigned char c = p[0];
    if (c < 0x80) {
      p++;
      avail--;
      switch (c) {
        case '"': out += "\\\""; break;
        case '\\': out += "\\\\"; break;
        case '\b': out += "\\b"; break;
        case '\f': out += "\\f"; break;
        case '\n': out += "\\n"; break;
        case '\r': out += "\\r"; break;
        case '\t': out += "\\t"; break;
        default: {
          char ubuf[7];
          std::snprintf(ubuf, sizeof ubuf, "\\u%04x", (unsigned)c);
          out.append(ubuf, 6);
          break;
        }
      }
    } else {
      size_t pl = 0;
      int cp = utf8_decode(p, avail, &pl);
      if (cp == 0xFFFD && pl == 1) {
        out += "\xEF\xBF\xBD";
        p++;
        avail--;
      } else {
        out.append((const char *)p, pl);
        p += pl;
        avail -= pl;
      }
    }
  }
  out += '"';
}

static void dump_rec(std::string &out, const Json &j, int indent, int depth) {
  switch (j.type) {
    case JType::Null:
      out += "null";
      break;
    case JType::Bool:
      out += (j.bval ? "true" : "false");
      break;
    case JType::Int: {
      char ibuf[32];
      auto [ptr, ec] = std::to_chars(ibuf, ibuf + sizeof ibuf, j.ival);
      if (ec == std::errc()) out.append(ibuf, (size_t)(ptr - ibuf));
      else str_appendf(out, "%lld", (long long)j.ival);
      break;
    }
    case JType::Float:
      if (!std::isfinite(j.dval)) {
        out += "null";
      } else {
        char fbuf[32];
        int n = std::snprintf(fbuf, sizeof fbuf, "%.15g", j.dval);
        if (n > 0 && (size_t)n < sizeof fbuf) out.append(fbuf, (size_t)n);
        else str_appendf(out, "%.15g", j.dval);
      }
      break;
    case JType::Str:
      dump_str(out, j.sval);
      break;
    case JType::Arr:
      if (j.aval.empty()) {
        out += "[]";
      } else {
        out += '[';
        for (size_t i = 0; i < j.aval.size(); i++) {
          if (i > 0) out += ',';
          if (indent >= 0) {
            out += '\n';
            out.append((size_t)(depth + 1) * (size_t)indent, ' ');
          }
          dump_rec(out, j.aval[i], indent, depth + 1);
        }
        if (indent >= 0) {
          out += '\n';
          out.append((size_t)depth * (size_t)indent, ' ');
        }
        out += ']';
      }
      break;
    case JType::Obj:
      if (j.oval.empty()) {
        out += "{}";
      } else {
        out += '{';
        for (size_t i = 0; i < j.oval.size(); i++) {
          if (i > 0) out += ',';
          if (indent >= 0) {
            out += '\n';
            out.append((size_t)(depth + 1) * (size_t)indent, ' ');
          }
          dump_str(out, j.oval[i].first);
          out += (indent >= 0) ? ": " : ":";
          dump_rec(out, j.oval[i].second, indent, depth + 1);
        }
        if (indent >= 0) {
          out += '\n';
          out.append((size_t)depth * (size_t)indent, ' ');
        }
        out += '}';
      }
      break;
  }
}

std::string Json::dump(int indent) const {
  std::string out;
  out.reserve(256);
  dump_rec(out, *this, indent, 0);
  return out;
}

/* ----------------------------- JSON parser ---------------------------- */
struct JParser {
  std::string_view src;
  size_t pos = 0;

  void skip_ws() {
    while (pos < src.size()) {
      char c = src[pos];
      if (c == ' ' || c == '\t' || c == '\n' || c == '\r') pos++;
      else break;
    }
  }

  static int hex_val(char c) {
    if (c >= '0' && c <= '9') return c - '0';
    if (c >= 'a' && c <= 'f') return c - 'a' + 10;
    if (c >= 'A' && c <= 'F') return c - 'A' + 10;
    return -1;
  }

  bool parse_hex4(uint32_t &out) {
    if (pos + 4 > src.size()) return false;
    uint32_t v = 0;
    for (int i = 0; i < 4; i++) {
      int h = hex_val(src[pos + (size_t)i]);
      if (h < 0) return false;
      v = (v << 4) | (uint32_t)h;
    }
    pos += 4;
    out = v;
    return true;
  }

  static void append_utf8(std::string &s, uint32_t cp) {
    if (cp <= 0x7F) {
      s += (char)cp;
    } else if (cp <= 0x7FF) {
      s += (char)(0xC0 | (cp >> 6));
      s += (char)(0x80 | (cp & 0x3F));
    } else if (cp <= 0xFFFF) {
      s += (char)(0xE0 | (cp >> 12));
      s += (char)(0x80 | ((cp >> 6) & 0x3F));
      s += (char)(0x80 | (cp & 0x3F));
    } else if (cp <= 0x10FFFF) {
      s += (char)(0xF0 | (cp >> 18));
      s += (char)(0x80 | ((cp >> 12) & 0x3F));
      s += (char)(0x80 | ((cp >> 6) & 0x3F));
      s += (char)(0x80 | (cp & 0x3F));
    } else {
      s += "\xEF\xBF\xBD";
    }
  }

  bool parse_string(std::string &out) {
    if (pos >= src.size() || src[pos] != '"') return false;
    pos++;
    out.clear();
    const char *base = src.data();
    size_t n = src.size();
    while (pos < n) {
      /* Fast path: scan contiguous unescaped bytes >= 0x20. */
      size_t run_start = pos;
      while (pos < n) {
        unsigned char b = (unsigned char)base[pos];
        if (b == '"' || b == '\\' || b < 0x20) break;
        pos++;
      }
      if (pos > run_start) {
        out.append(base + run_start, pos - run_start);
      }
      if (pos >= n) return false;
      unsigned char c = (unsigned char)base[pos++];
      if (c == '"') return true;
      if (c < 0x20) return false;
      if (pos >= n) return false;
      char esc = base[pos++];
      switch (esc) {
        case '"': out += '"'; break;
        case '\\': out += '\\'; break;
        case '/': out += '/'; break;
        case 'b': out += '\b'; break;
        case 'f': out += '\f'; break;
        case 'n': out += '\n'; break;
        case 'r': out += '\r'; break;
        case 't': out += '\t'; break;
        case 'u': {
          uint32_t cp = 0;
          if (!parse_hex4(cp)) return false;
          if (cp >= 0xD800 && cp <= 0xDBFF) {
            if (pos + 6 <= n && base[pos] == '\\' && base[pos + 1] == 'u') {
              size_t saved = pos;
              pos += 2;
              uint32_t low = 0;
              if (parse_hex4(low) && low >= 0xDC00 && low <= 0xDFFF) {
                cp = 0x10000u + (((cp - 0xD800u) << 10) | (low - 0xDC00u));
              } else {
                pos = saved;
                cp = 0xFFFDu;
              }
            } else {
              cp = 0xFFFDu;
            }
          } else if (cp >= 0xDC00 && cp <= 0xDFFF) {
            cp = 0xFFFDu;
          }
          append_utf8(out, cp);
          break;
        }
        default:
          return false;
      }
    }
    return false;
  }

  bool parse_number(Json &out) {
    size_t start = pos;
    if (src[pos] == '-') pos++;
    if (pos >= src.size()) return false;
    if (src[pos] == '0') {
      pos++;
      if (pos < src.size() && src[pos] >= '0' && src[pos] <= '9') return false;
    } else if (src[pos] >= '1' && src[pos] <= '9') {
      while (pos < src.size() && src[pos] >= '0' && src[pos] <= '9') pos++;
    } else {
      return false;
    }
    bool is_float = false;
    if (pos < src.size() && src[pos] == '.') {
      is_float = true;
      pos++;
      if (pos >= src.size() || src[pos] < '0' || src[pos] > '9') return false;
      while (pos < src.size() && src[pos] >= '0' && src[pos] <= '9') pos++;
    }
    if (pos < src.size() && (src[pos] == 'e' || src[pos] == 'E')) {
      is_float = true;
      pos++;
      if (pos < src.size() && (src[pos] == '+' || src[pos] == '-')) pos++;
      if (pos >= src.size() || src[pos] < '0' || src[pos] > '9') return false;
      while (pos < src.size() && src[pos] >= '0' && src[pos] <= '9') pos++;
    }
    const char *first = src.data() + start;
    const char *last = src.data() + pos;
    if (!is_float) {
      int64_t iv = 0;
      auto [ptr, ec] = std::from_chars(first, last, iv);
      if (ec == std::errc() && ptr == last) {
        out = Json(iv);
        return true;
      }
    }
    char stack_num[64];
    std::string heap_num;
    const char *cstr = nullptr;
    size_t nlen = pos - start;
    if (nlen < sizeof stack_num) {
      std::memcpy(stack_num, first, nlen);
      stack_num[nlen] = '\0';
      cstr = stack_num;
    } else {
      heap_num.assign(first, nlen);
      cstr = heap_num.c_str();
    }
    errno = 0;
    char *end = nullptr;
    double dv = std::strtod(cstr, &end);
    if (end != cstr + nlen || !std::isfinite(dv)) return false;
    out = Json(dv);
    return true;
  }

  bool parse_value(Json &out, int depth) {
    if (depth > 256) return false;
    skip_ws();
    if (pos >= src.size()) return false;
    char c = src[pos];
    if (c == 'n') {
      if (src.substr(pos, 4) != "null") return false;
      pos += 4;
      out = Json(nullptr);
      return true;
    }
    if (c == 't') {
      if (src.substr(pos, 4) != "true") return false;
      pos += 4;
      out = Json(true);
      return true;
    }
    if (c == 'f') {
      if (src.substr(pos, 5) != "false") return false;
      pos += 5;
      out = Json(false);
      return true;
    }
    if (c == '"') {
      std::string s;
      if (!parse_string(s)) return false;
      out = Json(std::move(s));
      return true;
    }
    if (c == '[') {
      pos++;
      skip_ws();
      out = Json::array();
      if (pos < src.size() && src[pos] == ']') {
        pos++;
        return true;
      }
      while (true) {
        Json elem;
        if (!parse_value(elem, depth + 1)) return false;
        out.aval.push_back(std::move(elem));
        skip_ws();
        if (pos < src.size() && src[pos] == ',') {
          pos++;
          continue;
        }
        if (pos < src.size() && src[pos] == ']') {
          pos++;
          return true;
        }
        return false;
      }
    }
    if (c == '{') {
      pos++;
      skip_ws();
      out = Json::object();
      if (pos < src.size() && src[pos] == '}') {
        pos++;
        return true;
      }
      while (true) {
        skip_ws();
        std::string key;
        if (!parse_string(key)) return false;
        skip_ws();
        if (pos >= src.size() || src[pos] != ':') return false;
        pos++;
        Json val;
        if (!parse_value(val, depth + 1)) return false;
        out[key] = std::move(val);
        skip_ws();
        if (pos < src.size() && src[pos] == ',') {
          pos++;
          continue;
        }
        if (pos < src.size() && src[pos] == '}') {
          pos++;
          return true;
        }
        return false;
      }
    }
    if (c == '-' || (c >= '0' && c <= '9')) {
      return parse_number(out);
    }
    return false;
  }
};

bool Json::parse(std::string_view src, Json &out) {
  JParser p{src, 0};
  Json tmp;
  if (!p.parse_value(tmp, 0)) return false;
  p.skip_ws();
  if (p.pos != src.size()) return false;
  out = std::move(tmp);
  return true;
}

/* ----------------------------- message helpers ------------------------ */
/* The tool definitions, identical to the Go reference implementation. */
static const char *TOOLS_JSON = R"JSON([
  {"type":"function","function":{
    "name":"bash",
    "description":"Executes a given bash command in a shell session, returning its output and exit code.",
    "parameters":{"type":"object","properties":{
      "command":{"type":"string","description":"The command to execute"},
      "workdir":{"type":"string","description":"The working directory to run the command in. Defaults to the current directory."}
    },"required":["command"]}}},
  {"type":"function","function":{
    "name":"read",
    "description":"Read a file or directory from the local filesystem. For files, returns numbered lines prefixed as `<line>: <content>`. For directories, returns entry names with trailing slashes for subdirectories.",
    "parameters":{"type":"object","properties":{
      "filePath":{"type":"string","description":"The path to the file or directory to read"},
      "offset":{"type":"integer","description":"The line number to start reading from (1-indexed, defaults to 1)"},
      "limit":{"type":"integer","description":"The maximum number of lines to read (defaults to 2000)"}
    },"required":["filePath"]}}},
  {"type":"function","function":{
    "name":"write",
    "description":"Write content to a file. If none of lineOffset, byteOffset, delLines and delBytes are provided, the entire file is overwritten with the new content (the file and any necessary parent directories are created if missing); otherwise content is written at the given line or byte offset, optionally deleting lines or bytes first. Line-based editing is primary.",
    "parameters":{"type":"object","properties":{
      "filePath":{"type":"string","description":"The path to the file to write"},
      "content":{"type":"string","description":"The content to write to the file"},
      "lineOffset":{"type":"integer","description":"1-indexed line number where writing begins (1 = first line). Takes precedence over byteOffset when both are given."},
      "byteOffset":{"type":"integer","description":"Byte offset where writing begins. Used only when lineOffset is not given."},
      "delLines":{"type":"integer","description":"Number of lines to delete starting at the insertion point. Takes precedence over delBytes when both are given."},
      "delBytes":{"type":"integer","description":"Number of bytes to delete starting at the insertion point. Used only when delLines is not given."}
    },"required":["filePath","content"]}}}
])JSON";

void serialize_message_json(Json &obj, const Message &m) {
  obj = Json::object();
  obj["role"] = m.role.empty() ? "" : m.role;
  if (!m.content.empty()) obj["content"] = m.content;
  if (!m.reasoning.empty()) obj["reasoning_content"] = m.reasoning;
  if (!m.calls.empty()) {
    Json arr = Json::array();
    for (auto &tc : m.calls) {
      Json t = Json::object();
      if (!tc.id.empty()) t["id"] = tc.id;
      t["type"] = "function";
      Json fn = Json::object();
      fn["name"] = tc.name.empty() ? "" : tc.name;
      fn["arguments"] = tc.arguments.empty() ? "{}" : tc.arguments;
      t["function"] = std::move(fn);
      arr.push_back(std::move(t));
    }
    obj["tool_calls"] = std::move(arr);
  }
  if (!m.tool_call_id.empty()) obj["tool_call_id"] = m.tool_call_id;
}

void parse_message_json(const Json &obj, Message &m) {
  m = Message{};
  if (!obj.is_object()) return;
  m.role = obj["role"].str("user");
  m.content = obj["content"].str();
  m.reasoning = obj["reasoning_content"].str();
  m.tool_call_id = obj["tool_call_id"].str();
  const Json &tcs = obj["tool_calls"];
  if (tcs.is_array()) {
    for (auto &tc : tcs.aval) {
      if (!tc.is_object()) continue;
      ToolCall t;
      t.id = tc["id"].str();
      const Json &fn = tc["function"];
      if (fn.is_object()) {
        t.name = fn["name"].str();
        t.arguments = fn["arguments"].str();
      }
      m.calls.push_back(std::move(t));
    }
  }
}

/* ----------------------------- request / response --------------------- */
std::string build_request_json(const Cfg &cfg, const Messages &msgs,
                               bool use_tools, bool stream) {
  Json root = Json::object();
  root["model"] = cfg.model;
  root["temperature"] = cfg.temperature;
  Json arr = Json::array();
  arr.aval.reserve(msgs.items.size());
  for (auto &m : msgs.items) {
    Json mo;
    serialize_message_json(mo, m);
    arr.push_back(std::move(mo));
  }
  root["messages"] = std::move(arr);
  root["stream"] = stream;
  if (use_tools) {
    static const Json kTools = []() {
      Json t;
      Json::parse(TOOLS_JSON, t);
      return t;
    }();
    root["tools"] = kTools;
  }
  if (stream) {
    Json so = Json::object();
    so["include_usage"] = true;
    root["stream_options"] = std::move(so);
  }
  /* Forward non-internal raw params, parsed as JSON when possible. */
  for (auto &kv : cfg.raw) {
    if (internal_key(kv.first)) continue;
    Json jv;
    if (Json::parse(kv.second, jv)) {
      root[kv.first] = std::move(jv);
    } else {
      root[kv.first] = kv.second;
    }
  }
  return root.dump();
}

bool parse_nonstream_response(std::string_view body, Message &out, Usage &usage) {
  out = Message{};
  usage = Usage{};
  Json root_mut;
  if (!Json::parse(body, root_mut) || !root_mut.is_object()) return false;
  const Json &root = root_mut;
  const Json &choices = root["choices"];
  if (!choices.is_array() || choices.empty()) return false;
  const Json &ch = choices.at(0);
  if (!ch.is_object()) return false;
  const Json &msg = ch["message"];
  if (!msg.is_object()) return false;
  out.role = msg["role"].str("assistant");
  out.content = msg["content"].str();
  std::string reas = msg["reasoning_content"].str();
  if (reas.empty()) reas = msg["reasoning"].str();
  if (reas.empty()) reas = msg["thought"].str();
  if (!reas.empty()) out.reasoning = std::move(reas);
  const Json &tcs = msg["tool_calls"];
  if (tcs.is_array()) {
    for (auto &tc : tcs.aval) {
      if (!tc.is_object()) continue;
      ToolCall t;
      t.id = tc["id"].str();
      const Json &fn = tc["function"];
      if (fn.is_object()) {
        t.name = fn["name"].str();
        t.arguments = fn["arguments"].str("{}");
      }
      out.calls.push_back(std::move(t));
    }
  }
  const Json &u = root["usage"];
  if (u.is_object()) {
    usage.prompt_tokens = u["prompt_tokens"].as_int();
    usage.completion_tokens = u["completion_tokens"].as_int();
    usage.total_tokens = u["total_tokens"].as_int();
    const Json &ptd = u["prompt_tokens_details"];
    if (ptd.is_object()) {
      usage.cached_tokens = ptd["cached_tokens"].as_int();
    }
    if (usage.cached_tokens == 0) {
      usage.cached_tokens = u["cached_tokens"].as_int();
    }
  }
  usage.model = root["model"].str();
  return true;
}