/* json.cpp — request payload building and non-streaming response parsing. */
#include "bantam.hpp"
#include <nlohmann/json.hpp>

using json = nlohmann::json;

/* 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 offset and del_bytes are both omitted, the entire file is overwritten with the new content; otherwise content is written at the given byte offset, optionally deleting bytes first.",
    "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"},
      "offset":{"type":"integer","description":"Byte offset to start writing from. If omitted together with del_bytes, the whole file is overwritten instead. Defaults to 0 (start of file)."},
      "del_bytes":{"type":"integer","description":"Bytes to delete starting at offset. If omitted together with offset, the whole file is overwritten instead."}
    },"required":["filePath","content"]}}}
])JSON";

static void serialize_message(json &obj, const Message &m) {
  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;
      if (!tc.id.empty()) t["id"] = tc.id;
      t["type"] = "function";
      json fn;
      fn["name"] = tc.name.empty() ? "" : tc.name;
      fn["arguments"] = tc.arguments.empty() ? "{}" : tc.arguments;
      t["function"] = fn;
      arr.push_back(t);
    }
    obj["tool_calls"] = arr;
  }
  if (!m.tool_call_id.empty()) obj["tool_call_id"] = m.tool_call_id;
}

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();
  for (auto &m : msgs.items) {
    json mo;
    serialize_message(mo, m);
    arr.push_back(mo);
  }
  root["messages"] = arr;
  root["stream"] = stream;
  if (use_tools) {
    try {
      json tools = json::parse(TOOLS_JSON);
      root["tools"] = tools;
    } catch (...) {}
  }
  if (stream) {
    json so = json::object();
    so["include_usage"] = true;
    root["stream_options"] = so;
  }
  /* Forward non-internal raw params, parsed as JSON when possible. */
  for (auto &kv : cfg.raw) {
    if (internal_key(kv.first)) continue;
    try {
      json jv = json::parse(kv.second);
      root[kv.first] = jv;
    } catch (...) {
      root[kv.first] = kv.second;
    }
  }
  return root.dump(-1, ' ', false, json::error_handler_t::replace);
}

static std::string json_str(const json &j, const char *k, std::string_view def = "") {
  if (j.is_object() && j.contains(k) && j[k].is_string())
    return j[k].get<std::string>();
  return std::string(def);
}

static int json_int(const json &j, const char *k, int def = 0) {
  if (j.is_object() && j.contains(k) && j[k].is_number_integer())
    return j[k].get<int>();
  return def;
}

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