// Author: Simon-Pierre Boucher — contact@spboucher.ai
#include "tokenizer/bpe.h"

#include <cstdio>
#include <cstdlib>
#include <fstream>
#include <limits>
#include <sstream>

namespace forge::tok {

namespace {
[[noreturn]] void die(const std::string& msg) {
    std::fprintf(stderr, "forge/tokenizer: %s\n", msg.c_str());
    std::abort();
}
} // namespace

void BPETokenizer::load(const std::string& model_path) {
    std::ifstream in(model_path);
    if (!in) die("cannot open " + model_path);
    std::string header;
    std::getline(in, header);
    if (header != "forgebpe v1") die("bad model header in " + model_path);
    int64_t vocab_size = 0;
    in >> vocab_size;

    vocab_.resize(size_t(vocab_size));
    for (int i = 0; i < 256 && i < vocab_size; ++i) vocab_[size_t(i)] = std::string(1, char(i));

    int32_t id, left, right;
    while (in >> id >> left >> right) {
        merges_[{left, right}] = id;
        vocab_[size_t(id)] = vocab_[size_t(left)] + vocab_[size_t(right)];
    }
    if (int64_t(merges_.size()) != vocab_size - 256)
        die("merge count does not match vocab size in " + model_path);
}

std::vector<int32_t> BPETokenizer::encode(const std::string& text) const {
    std::vector<int32_t> ids;
    ids.reserve(text.size());
    for (unsigned char c : text) ids.push_back(int32_t(c));

    // greedy: repeatedly apply the LOWEST-id (earliest-learned) merge present
    while (ids.size() >= 2) {
        int32_t best_id = std::numeric_limits<int32_t>::max();
        std::pair<int32_t, int32_t> best_pair{-1, -1};
        for (size_t i = 0; i + 1 < ids.size(); ++i) {
            auto it = merges_.find({ids[i], ids[i + 1]});
            if (it != merges_.end() && it->second < best_id) {
                best_id = it->second;
                best_pair = it->first;
            }
        }
        if (best_pair.first < 0) break;
        std::vector<int32_t> next;
        next.reserve(ids.size());
        for (size_t i = 0; i < ids.size();) {
            if (i + 1 < ids.size() && ids[i] == best_pair.first &&
                ids[i + 1] == best_pair.second) {
                next.push_back(best_id);
                i += 2;
            } else {
                next.push_back(ids[i]);
                i += 1;
            }
        }
        ids = std::move(next);
    }
    return ids;
}

std::string BPETokenizer::decode(const std::vector<int32_t>& ids) const {
    std::string out;
    for (int32_t id : ids) {
        if (id >= 0 && size_t(id) < vocab_.size()) out += vocab_[size_t(id)];
    }
    return out;
}

} // namespace forge::tok