#include "ling3/chat_protocol.h"
#include <fstream>
#include <iostream>

using namespace ling3::chat;
void Check(bool condition, const std::string & message) {
    if (!condition) throw std::runtime_error(message);
}
int main(int argc, char ** argv) {
    if (argc != 2) throw std::runtime_error("usage: ling3-chat-tools-test tests/chat_tools_golden.json");
    // Prompts rendered by transformers from the official chat_template.jinja
    // (tools/render_chat_tools_golden.py); the C++ renderer must match byte for byte.
    std::ifstream file(argv[1]);
    const auto golden = OrderedJson::parse(file);
    auto find = [&](const std::string & name) -> const OrderedJson & {
        for (const auto & c : golden) if (c["name"] == name) return c;
        throw std::runtime_error("missing golden case " + name);
    };
    for (const auto & c : golden)
        Check(Parse(c["request"]).prompt == c["prompt"].get<std::string>(),
              "template mismatch: " + c["name"].get<std::string>());

    auto body = find("tool_loop")["request"];
    const auto & tools = body["tools"];
    const auto prompt = find("tool_loop")["prompt"].get<std::string>();
    Check(Parse(body).tool_calls && Parse(body).parallel_tool_calls, "tools must enable call parsing");
    auto none = body; none["tool_choice"] = "none";
    Check(!Parse(none).tool_calls && Parse(none).prompt == prompt, "tool_choice none still renders tools");
    auto stored = body; stored["store"] = false;
    Check(Parse(stored).prompt == prompt, "store flag rejected");
    auto parts = body;
    parts["messages"][1]["content"] = OrderedJson::array({{{"type", "text"}, {"text", "Fix /tmp/a.txt."}}});
    Check(Parse(parts).prompt == prompt, "text-part content");
    for (const auto & [pointer, value] : std::vector<std::pair<std::string, OrderedJson>>{
        {"/tools", "read_file"}, {"/tools/0/type", "code"}, {"/tool_choice", "required"},
        {"/parallel_tool_calls", "yes"}, {"/store", "no"}, {"/messages/2/tool_calls/0/function/arguments", "{oops"},
        {"/messages/2/tool_calls/0/function/arguments", "[1]"}, {"/messages/3/content", nullptr},
        {"/messages/3/id", "x"}, {"/messages/1/content", OrderedJson::array({{{"type", "image_url"}}})}}) {
        auto bad = body; bad[OrderedJson::json_pointer(pointer)] = value; bool failed = false;
        try { Parse(bad); } catch (const Error &) { failed = true; }
        Check(failed, "invalid tool request accepted: " + pointer);
    }

    // Arguments are typed by the declared schema; a parsed call re-renders to
    // the generated text, so the next turn's prompt keeps the cached prefix.
    auto calls = ParseToolCalls("<tool_call>run\n<arg_key>command</arg_key>\n<arg_value>123</arg_value>"
        "<arg_key>timeout</arg_key>\n<arg_value>30</arg_value><arg_key>env</arg_key>\n<arg_value>{\"A\": [1, 2.5]}"
        "</arg_value><arg_key>dry</arg_key>\n<arg_value>false</arg_value>\n</tool_call>", tools, true);
    Check(calls.ok && calls.calls.size() == 1 && calls.calls[0].name == "run", "call not parsed");
    Check(calls.calls[0].arguments.dump() == R"({"command":"123","timeout":30,"env":{"A":[1,2.5]},"dry":false})",
          "arguments not typed by schema");
    for (const auto * name : {"tool_loop", "unicode_call", "thinking_tool_loop"})
        for (const auto & m : find(name)["request"]["messages"]) {
            if (!m.contains("tool_calls")) continue;
            std::string generated;
            for (const auto & c : m["tool_calls"])
                generated += (generated.empty() ? "" : "\n") + RenderToolCall({c["function"]["name"].get<std::string>(),
                    OrderedJson::parse(c["function"]["arguments"].get<std::string>())});
            Check(prompt.find(generated) != std::string::npos || std::string(name) != "tool_loop", "history render");
            const auto parsed = ParseToolCalls(generated, find(name)["request"]["tools"], true);
            Check(parsed.ok && parsed.calls.size() == m["tool_calls"].size(), "round trip parse");
            std::string again;
            for (const auto & c : parsed.calls) again += (again.empty() ? "" : "\n") + RenderToolCall(c);
            Check(again == generated, std::string("round trip render: ") + name);
        }
    const std::string two = "<tool_call>read_file\n<arg_key>path</arg_key>\n<arg_value>/tmp/a</arg_value>\n</tool_call>\n"
                            "<tool_call>read_file <arg_key>path</arg_key><arg_value>/tmp/b</arg_value></tool_call>";
    Check(ParseToolCalls(two, tools, true).calls.size() == 2, "parallel calls");
    Check(ParseToolCalls(two, tools, false).calls.size() == 1, "parallel_tool_calls=false");
    Check(ParseToolCalls(two, tools, true).calls[1].arguments["path"] == "/tmp/b", "whitespace-free call");
    Check(!ParseToolCalls("<tool_call>unknown\n\n</tool_call>", tools, true).ok, "undeclared tool accepted");
    Check(!ParseToolCalls("<tool_call>read_file\n<arg_key>path</arg_key>\n<arg_value>/tmp", tools, true).ok,
          "truncated call accepted");
    Check(ParseToolCalls("<tool_call>read_file\n\n</tool_call>\nDone.", tools, true).trailing == "Done.",
          "trailing content lost");

    // Streaming: call text and the newline before it are held, never sent as content.
    ToolCallFilter filter(true);
    std::string shown;
    for (const auto * piece : {"Read it.", "\n<tool", "_call>patch\n<arg_key>path</arg_key>",
                               "\n<arg_value>/tmp/a</arg_value>\n</tool_call>"})
        shown += filter.Push(piece);
    shown += filter.Push({}, true);
    Check(shown == "Read it." && filter.captured().starts_with("<tool_call>patch"), "call text streamed");
    Check(filter.Released() == "\n" + filter.captured(), "held newline lost");
    ToolCallFilter plain(true);
    shown = plain.Push("a <tool");
    shown += plain.Push("box>\n");
    shown += plain.Push({}, true);
    Check(shown == "a <toolbox>\n" && plain.captured().empty(), "false tag prefix withheld");
    ToolCallFilter off(false);
    Check(off.Push("<tool_call>x") == "<tool_call>x", "disabled filter held text");
    std::cout << "chat tools: " << golden.size() << " golden prompts, parsing, round trip, streaming passed\n";
}