/* M3.3 gate: render ChatML fixtures and eyeball the output + token ids.
 *
 *   ./chat_template_test model.gguf
 *
 * Renders four fixtures via render_chat_template_text + render_chat_template,
 * prints the text (with markers) and the token-id stream.  Cross-check the
 * special-token ids against marker_probe: im_start=248045, im_end=248046,
 * tool_call open/close=248058/248059.  No automated equality against a Jinja
 * interpreter (post-MVP); this is a human-eyeball gate plus a structural
 * assert that the special ids appear where expected.
 */
#include "gguf.h"
#include "qxmx_chat_template.h"
#include "qxmx_model.h"
#include "qxmx_tokenizer.h"

#include <cstdio>
#include <cstring>
#include <string>
#include <vector>

using namespace qx;

static void dump(const char* label, const std::string& text,
                 const std::vector<int32_t>& ids) {
    std::printf("===== %s (%zu chars, %zu tokens) =====\n", label, text.size(), ids.size());
    for (char c : text) {
        if (c == '\n') std::fputs("\\n", stdout);
        else if (c == '\t') std::fputs("\\t", stdout);
        else std::fputc(c, stdout);
    }
    std::printf("\n--- ids:");
    for (size_t i = 0; i < ids.size(); i++) {
        if (i && (i % 16 == 0)) std::printf("\n       ");
        std::printf(" %d", ids[i]);
    }
    std::printf("\n\n");
}

/* Write a fixture's raw rendered text to <dir>/<name>.txt and token ids to
 * <dir>/<name>.ids (one decimal per line). */
static void dump_file(const char* dir, const char* name,
                      const std::string& text, const std::vector<int32_t>& ids) {
    if (!dir) return;
    std::string p = std::string(dir) + "/" + name + ".txt";
    FILE* f = std::fopen(p.c_str(), "wb");
    if (!f) return;
    std::fwrite(text.data(), 1, text.size(), f);
    std::fclose(f);
    std::string q = std::string(dir) + "/" + name + ".ids";
    FILE* g = std::fopen(q.c_str(), "w");
    if (!g) return;
    for (int32_t id : ids) std::fprintf(g, "%d\n", id);
    std::fclose(g);
}

static bool expect_id(const std::vector<int32_t>& ids, int32_t want, const char* where) {
    for (int32_t id : ids) if (id == want) return true;
    std::fprintf(stderr, "FAIL: expected id %d in %s\n", want, where);
    return false;
}

int main(int argc, char** argv) {
    if (argc < 2) { std::fprintf(stderr, "usage: %s model.gguf\n", argv[0]); return 1; }
    qx_model m; char err[256];
    if (qx_model_open(&m, argv[1], err, sizeof err)) { std::fprintf(stderr, "load: %s\n", err); return 1; }
    tokenizer tok;
    if (!tok.load(&m.gguf, err, sizeof err)) { std::fprintf(stderr, "tok: %s\n", err); return 1; }

    bool ok = true;
    /* If given a dump path (argv[2]), write each fixture's rendered TEXT to
     * <path>.F<n>.txt for the jinja crosscheck script to diff byte-for-byte. */
    const char* dumpdir = (argc >= 3) ? argv[2] : nullptr;

    /* F1: plain two-turn chat + generation prompt, thinking disabled. */
    {
        std::vector<chat_message> msgs = {
            {"user", "What is the capital of France?", {}},
        };
        std::vector<chat_tool> no_tools;
        std::string text = render_chat_template_text(msgs, no_tools, true, false);
        std::vector<int32_t> ids = render_chat_template(tok, msgs, no_tools, true, false);
        dump("F1 plain chat + gen (thinking off)", text, ids);
        dump_file(dumpdir, "F1", text, ids);
        ok &= expect_id(ids, 248045, "F1 (im_start)");
        ok &= expect_id(ids, 248046, "F1 (im_end)");
    }

    /* F2: system + tools, generation prompt, thinking enabled. */
    {
        std::vector<chat_tool> tools = {
            {"get_weather",
             "Get the current weather for a city.",
             "{\"type\":\"object\",\"properties\":{\"city\":{\"type\":\"string\"}},\"required\":[\"city\"]}"},
        };
        std::vector<chat_message> msgs = {
            {"system", "You are a helpful assistant.", {}},
            {"user", "What's the weather in Paris?", {}},
        };
        std::string text = render_chat_template_text(msgs, tools, true, true);
        std::vector<int32_t> ids = render_chat_template(tok, msgs, tools, true, true);
        dump("F2 system + tools + gen (thinking on)", text, ids);
        dump_file(dumpdir, "F2", text, ids);
        ok &= expect_id(ids, 248045, "F2 (im_start)");
        ok &= expect_id(ids, 248046, "F2 (im_end)");
        ok &= expect_id(ids, 248058, "F2 (tc_open in tools example)");
        ok &= expect_id(ids, 248059, "F2 (tc_close in tools example)");
    }

    /* F3: multi-turn with tool result. */
    {
        std::vector<chat_tool> tools = {
            {"get_weather", "Get weather.",
             "{\"type\":\"object\",\"properties\":{\"city\":{\"type\":\"string\"}},\"required\":[\"city\"]}"},
        };
        std::vector<chat_message> msgs = {
            {"user", "Weather in Paris?", {}},
            {"assistant", "",
                {{"get_weather", "{\"city\":\"Paris\"}"}}},
            {"tool", "18C, sunny", {}},
            {"user", "Thanks!", {}},
        };
        std::string text = render_chat_template_text(msgs, tools, true, false);
        std::vector<int32_t> ids = render_chat_template(tok, msgs, tools, true, false);
        dump("F3 tool call + result multi-turn", text, ids);
        dump_file(dumpdir, "F3", text, ids);
        ok &= expect_id(ids, 248058, "F3 (tc_open in assistant tool_calls)");
        ok &= expect_id(ids, 248059, "F3 (tc_close in assistant tool_calls)");
    }

    /* F4: two consecutive tool results share one user turn. */
    {
        std::vector<chat_message> msgs = {
            {"user", "Weather in Paris and London?", {}},
            {"assistant", "",
                {{"get_weather", "{\"city\":\"Paris\"}"},
                 {"get_weather", "{\"city\":\"London\"}"}}},
            {"tool", "Paris: 18C", {}},
            {"tool", "London: 15C", {}},
            {"user", "Great.", {}},
        };
        std::vector<chat_tool> tools;
        std::string text = render_chat_template_text(msgs, tools, true, false);
        std::vector<int32_t> ids = render_chat_template(tok, msgs, tools, true, false);
        dump("F4 two parallel tool calls + two results", text, ids);
        dump_file(dumpdir, "F4", text, ids);
    }

    qx_model_close(&m);
    std::printf("%s\n", ok ? "PASS" : "FAIL");
    return ok ? 0 : 1;
}