// Author: Simon-Pierre Boucher — contact@spboucher.ai
//
// Tokenizer tests: decode(encode(x)) == x for byte-exact round-tripping,
// and — when tools/ has produced token ids for the same text — that the C++
// greedy lowest-id encoder agrees with the vectorized Python encoder used to
// build the .bin files. The test writes its own tiny .model so it needs no
// data prep; the Python-agreement half runs only if FORGE_TOK_MODEL and
// FORGE_TOK_IDS are set (wired up by tests/tokenizer_agreement.py).
#include "tokenizer/bpe.h"

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

namespace {

int g_failures = 0;

void expect(bool cond, const char* what) {
    if (cond) {
        std::printf("  ok: %s\n", what);
    } else {
        std::printf("  FAIL: %s\n", what);
        ++g_failures;
    }
}

// Minimal hand-built vocab: merges chosen so encoding "aaabdaaabac" exercises
// repeated-pair overlap handling (minbpe's worked example).
void write_test_model(const std::string& path) {
    std::ofstream out(path);
    out << "forgebpe v1\n";
    out << 259 << "\n";
    out << "256 97 97\n";  // 'aa'
    out << "257 256 97\n"; // 'aaa'
    out << "258 257 98\n"; // 'aaab'
}

} // namespace

int main() {
    const std::string model_path = "test_tok.model";
    write_test_model(model_path);

    forge::tok::BPETokenizer tok;
    tok.load(model_path);
    expect(tok.vocab_size() == 259, "vocab size");

    // "aaabdaaabac": 'aaab' merges twice, leaving d/a/c as raw bytes
    const std::string text = "aaabdaaabac";
    std::vector<int32_t> ids = tok.encode(text);
    expect(ids == std::vector<int32_t>({258, 100, 258, 97, 99}), "merge order");
    expect(tok.decode(ids) == text, "round-trip (merged)");

    // Overlap: "aaaa" applies merge 256 ('aa') to BOTH non-overlapping pairs
    // in one round, so 'aaa' (id 257) never forms — greedy is lowest-id-first,
    // not longest-match. "aaa" does reach 257 (256 then 256+97).
    expect(tok.decode(tok.encode("aaaa")) == "aaaa", "round-trip (overlap)");
    expect(tok.encode("aaaa") == std::vector<int32_t>({256, 256}), "overlap merge");
    expect(tok.encode("aaa") == std::vector<int32_t>({257}), "chained merge");

    // Bytes with no merges, and full 0-255 range including UTF-8 and NUL
    std::string bytes;
    for (int i = 1; i < 256; ++i) bytes += char(i);
    expect(tok.decode(tok.encode(bytes)) == bytes, "round-trip (all byte values)");
    const std::string utf8 = "héllo wörld — ünïcode ✓";
    expect(tok.decode(tok.encode(utf8)) == utf8, "round-trip (utf-8)");
    expect(tok.encode("").empty(), "empty input");

    std::remove(model_path.c_str());

    // Optional: agreement with the Python encoder on real text.
    const char* py_model = std::getenv("FORGE_TOK_MODEL");
    const char* py_ids = std::getenv("FORGE_TOK_IDS");
    const char* py_text = std::getenv("FORGE_TOK_TEXT");
    if (py_model && py_ids && py_text) {
        forge::tok::BPETokenizer real;
        real.load(py_model);
        std::ifstream tf(py_text, std::ios::binary);
        std::stringstream ts;
        ts << tf.rdbuf();
        std::vector<int32_t> mine = real.encode(ts.str());

        std::ifstream idf(py_ids);
        std::vector<int32_t> theirs;
        int32_t v;
        while (idf >> v) theirs.push_back(v);

        char label[128];
        std::snprintf(label, sizeof(label),
                      "python encoder agreement (%zu vs %zu tokens)", mine.size(),
                      theirs.size());
        expect(mine == theirs, label);
        expect(real.decode(mine) == ts.str(), "round-trip (real vocab, real text)");
    } else {
        std::printf("  skip: python agreement (set FORGE_TOK_MODEL/IDS/TEXT)\n");
    }

    if (g_failures) {
        std::printf("\n%d tokenizer test(s) FAILED\n", g_failures);
        return 1;
    }
    std::printf("\nall tokenizer tests passed\n");
    return 0;
}