/* qxmx_tok: load the tokenizer from a Bonsai GGUF and exercise encode/decode.
 *
 *   ./qxmx_tok model.gguf              # built-in round-trip tests
 *   ./qxmx_tok model.gguf "text"       # encode + decode one string
 *   ./qxmx_tok model.gguf -f file.txt  # encode file, print token count
 *
 * Validation: built-in suite checks round-trip (decode(encode(x))==x) on
 * English, code, CJK, emoji, mixed; and prints the first-N token ids for
 * eyeball comparison against llama.cpp's llama-tokenize. */
#include "gguf.h"
#include "qxmx_model.h"
#include "qxmx_tokenizer.h"

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

int main(int argc, char** argv) {
    if (argc < 2) { std::fprintf(stderr, "usage: %s model.gguf [text | -f file]\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; }

    qx::tokenizer tok;
    if (!tok.load(&m.gguf, err, sizeof err)) { std::fprintf(stderr, "tok load: %s\n", err); return 1; }
    std::printf("tokenizer: %zu tokens, %zu merges, bos=%u eos=%u pad=%u add_bos=%d\n",
                tok.tokens().size(), tok.merges().size(), tok.bos_id, tok.eos_id, tok.pad_id, tok.add_bos);

    auto enc_dec = [&](const std::string& s) {
        std::vector<int32_t> ids;
        tok.encode(s, &ids);
        std::string back = tok.decode(ids.data(), ids.size());
        bool ok = (back == s);
        std::printf("[%s] %zu toks  \"%s\" -> \"%s\"\n",
                    ok ? "ok" : "FAIL", ids.size(),
                    (s.size() < 50 ? s : (s.substr(0, 47) + "...")).c_str(),
                    (back.size() < 50 ? back : (back.substr(0, 47) + "...")).c_str());
        if (!ok) {
            std::printf("  ids:");
            for (size_t i = 0; i < ids.size() && i < 30; i++) std::printf(" %d", ids[i]);
            if (ids.size() > 30) std::printf(" ...");
            std::printf("\n");
        }
        return ok;
    };

    /* single-string mode */
    if (argc >= 3 && std::strcmp(argv[2], "-f") != 0) {
        std::vector<int32_t> ids;
        tok.encode(argv[2], &ids);
        std::printf("ids (%zu):", ids.size());
        for (int32_t id : ids) std::printf(" %d", id);
        std::printf("\n");
        std::string back = tok.decode(ids.data(), ids.size());
        std::printf("round-trip: \"%s\"\n", back.c_str());
        qx_model_close(&m);
        return 0;
    }
    /* file mode */
    if (argc >= 4 && std::strcmp(argv[2], "-f") == 0) {
        std::ifstream f(argv[3]);
        std::stringstream ss; ss << f.rdbuf();
        std::string s = ss.str();
        std::vector<int32_t> ids;
        tok.encode(s, &ids);
        std::printf("%zu bytes -> %zu tokens\n", s.size(), ids.size());
        qx_model_close(&m);
        return 0;
    }

    /* built-in round-trip suite */
    int fails = 0;
    fails += !enc_dec("Hello, world!");
    fails += !enc_dec("The quick brown fox jumps over the lazy dog.");
    fails += !enc_dec("  leading and trailing spaces  ");
    fails += !enc_dec("line1\nline2\n\nline4");
    fails += !enc_dec("tab\there");
    fails += !enc_dec("code: for (int i=0; i<n; i++) { a[i] = b[i]*c; }");
    fails += !enc_dec("numbers: 12345 67890 3.14159 1e-9");
    fails += !enc_dec("punct! @#$%^&*() {}[]|\\:;\"'<>?,./~`");
    fails += !enc_dec("Chinese: 你好世界，机器学习。");
    fails += !enc_dec("Japanese: こんにちは世界");
    fails += !enc_dec("Korean: 안녕하세요 세계");
    fails += !enc_dec("Emoji: 😀🎉🚀💯 👨‍👩‍👧‍👦");
    fails += !enc_dec("Mixed: Hello 世界 🌍 123 code();");
    fails += !enc_dec("contractions: don't can't won't I'm we're they've");
    fails += !enc_dec("");
    fails += !enc_dec(" ");
    /* Invalid UTF-8: llama.cpp substitutes U+FFFD (lossy by design). Verify
     * the tokenizer doesn't crash and produces a replacement char. */
    {
        std::string bad; bad += (char)0xff; bad += " test";
        std::vector<int32_t> ids;
        tok.encode(bad, &ids);
        std::string back = tok.decode(ids.data(), ids.size());
        std::printf("[%s] invalid-utf8 -> %zu toks, round-trip starts with U+FFFD\n",
                    (!back.empty() && (uint8_t)back[0] == 0xef) ? "ok" : "FAIL",
                    ids.size());
        if (back.empty() || (uint8_t)back[0] != 0xef) fails++;
    }

    std::printf("\n%d failures\n", fails);
    qx_model_close(&m);
    return fails ? 1 : 0;
}