#include "qxmx_tokenizer.h"

#include <algorithm>
#include <cstring>
#include <queue>
#include <unordered_map>

#include "unicode.h"

namespace qx {

/* ---- helpers for reading GGUF string arrays (length-prefixed, 8B len) ---- */
static const char* skip_str(const char* p, std::string* out) {
    uint64_t n;
    std::memcpy(&n, p, 8);
    p += 8;
    out->assign(p, n);
    return p + n;
}

bool tokenizer::load(const gguf_file* g, char* err, size_t errlen) {
    const gguf_kv* k_tokens = nullptr;
    const gguf_kv* k_type   = nullptr;
    const gguf_kv* k_merges = nullptr;
    for (uint64_t i = 0; i < g->n_kv; i++) {
        char key[96];
        std::snprintf(key, sizeof key, "%.*s", (int)g->kv[i].key.len, g->kv[i].key.ptr);
        if      (std::strcmp(key, "tokenizer.ggml.tokens")     == 0) k_tokens = &g->kv[i];
        else if (std::strcmp(key, "tokenizer.ggml.token_type") == 0) k_type   = &g->kv[i];
        else if (std::strcmp(key, "tokenizer.ggml.merges")     == 0) k_merges = &g->kv[i];
    }
    if (!k_tokens || !k_type) {
        std::snprintf(err, errlen, "missing tokenizer.ggml.tokens/token_type");
        return false;
    }

    /* tokens: array of N length-prefixed strings */
    const char* p = (const char*)k_tokens->data;
    tokens_.resize(k_tokens->arr_len);
    for (uint64_t i = 0; i < k_tokens->arr_len; i++) p = skip_str(p, &tokens_[i]);

    /* token_type: array of N int32 */
    type_.assign((const int32_t*)k_type->data,
                 (const int32_t*)k_type->data + k_type->arr_len);

    /* merges: array of N length-prefixed "left right" strings; split on first
     * space to get (left, right) rank pairs (rank = array index). */
    if (k_merges) {
        p = (const char*)k_merges->data;
        merges_storage_.reserve(k_merges->arr_len);
        merge_ranks_.reserve(k_merges->arr_len * 2);
        for (uint64_t i = 0; i < k_merges->arr_len; i++) {
            std::string m;
            p = skip_str(p, &m);
            auto sp = m.find(' ');
            if (sp == std::string::npos) {
                std::snprintf(err, errlen, "merge %llu has no space", (unsigned long long)i);
                return false;
            }
            auto pair = std::make_pair(m.substr(0, sp), m.substr(sp + 1));
            merge_ranks_[pair] = (int)i;
            merges_storage_.push_back(std::move(pair));
        }
    }

    /* text -> id index (sorted for binary search). */
    token_index_.reserve(tokens_.size());
    for (size_t i = 0; i < tokens_.size(); i++) {
        if (type_[i] == 6 /* BYTE */ || tokens_[i].empty()) continue;
        token_index_.push_back({tokens_[i], (int32_t)i});
    }
    std::sort(token_index_.begin(), token_index_.end());

    gguf_u32(g, "tokenizer.ggml.bos_token_id", &bos_id);
    gguf_u32(g, "tokenizer.ggml.eos_token_id", &eos_id);
    gguf_u32(g, "tokenizer.ggml.padding_token_id", &pad_id);
    const gguf_kv* kb = gguf_find(g, "tokenizer.ggml.add_bos_token");
    if (kb && kb->type == GGUF_BOOL) add_bos = *(const uint8_t*)kb->data;
    const gguf_kv* ke = gguf_find(g, "tokenizer.ggml.add_eos_token");
    if (ke && ke->type == GGUF_BOOL) add_eos = *(const uint8_t*)ke->data;
    return true;
}

/* ---- BPE primitives ----------------------------------------------------- */

static int32_t token_id_for(const tokenizer* t, const std::string& s) {
    const auto& idx = t->token_index();
    auto it = std::lower_bound(idx.begin(), idx.end(), std::make_pair(s, 0),
        [](const std::pair<std::string, int32_t>& a, const std::pair<std::string, int32_t>& b) {
            return a.first < b.first;
        });
    return (it != idx.end() && it->first == s) ? it->second : -1;
}

/* rank of merging (left+right), or -1 if no such merge exists (O(1) hash). */
static int merge_rank(const tokenizer* t, const std::string& left,
                      const std::string& right) {
    const auto& mr = t->merge_ranks();
    auto it = mr.find({left, right});
    return it != mr.end() ? it->second : -1;
}

/* One symbol in the doubly-linked merge list for a word. */
struct sym {
    int32_t prev, next;
    std::string text;
    bool alive;
};

struct bigram {
    int rank;       /* merge rank (lower merges first) */
    int left, right;
    std::string text;  /* left+right, cached for staleness check */
    bool operator>(const bigram& o) const {
        if (rank != o.rank) return rank > o.rank;
        return left > o.left;  /* tie-break: stable */
    }
};

/* Greedy lowest-rank BPE merge over one pre-tokenized word. */
static void bpe_word(const tokenizer* t, const std::string& word,
                     std::vector<int32_t>* out) {
    /* Special-case: whole word is already a vocab token (ignore_merges path). */
    int32_t whole = token_id_for(t, word);
    if (whole >= 0) { out->push_back(whole); return; }

    /* Split word into single-symbol list.  GPT-2 BPE operates on the byte-
     * encoded form (unicode_regex_split already did byte->utf8 encoding when
     * called with byte_encode=true), so each UTF-8 character is one symbol. */
    std::vector<sym> syms;
    size_t off = 0;
    while (off < word.size()) {
        size_t cl = std::min(word.size() - off, (size_t)unicode_len_utf8(word[off]));
        syms.push_back({(int)syms.size() - 1, (int)syms.size() + 1,
                        word.substr(off, cl), true});
        off += cl;
    }
    if (syms.empty()) return;
    syms.back().next = -1;

    std::priority_queue<bigram, std::vector<bigram>, std::greater<bigram>> q;
    auto add_bg = [&](int li, int ri) {
        if (li < 0 || ri < 0 || li >= (int)syms.size() || ri >= (int)syms.size()) return;
        if (!syms[li].alive || !syms[ri].alive) return;
        int r = merge_rank(t, syms[li].text, syms[ri].text);
        if (r < 0) return;
        q.push({r, li, ri, syms[li].text + syms[ri].text});
    };
    for (int i = 0; i + 1 < (int)syms.size(); i++) add_bg(i, i + 1);

    while (!q.empty()) {
        bigram b = q.top(); q.pop();
        if (!syms[b.left].alive || !syms[b.right].alive) continue;
        if (syms[b.left].text + syms[b.right].text != b.text) continue;  /* stale */

        /* merge right into left */
        syms[b.left].text += syms[b.right].text;
        syms[b.right].alive = false;
        syms[b.left].next = syms[b.right].next;
        if (syms[b.right].next >= 0) syms[syms[b.right].next].prev = b.left;

        add_bg(syms[b.left].prev, b.left);
        add_bg(b.left, syms[b.left].next);
    }

    /* Emit surviving symbols; look up each in vocab.  Byte-level BPE guarantees
     * every merged symbol is a vocab token, but fall back to per-byte tokens. */
    int i = 0;
    while (i != -1) {
        if (syms[i].alive) {
            int32_t id = token_id_for(t, syms[i].text);
            if (id >= 0) {
                out->push_back(id);
            } else {
                for (char c : syms[i].text) {
                    int32_t bid = token_id_for(t, std::string(1, c));
                    if (bid >= 0) out->push_back(bid);
                }
            }
        }
        i = syms[i].next;
    }
}

/* ---- public encode/decode ---------------------------------------------- */

size_t tokenizer::encode(const std::string& text, std::vector<int32_t>* out,
                         bool add_special) const {
    if (add_special && add_bos) out->push_back(bos_id);

    /* qwen35 pre-tokenizer + GPT-2 byte encoding, both done by the vendored
     * unicode_regex_split (the regex matches the qwen35 case in llama-vocab.cpp). */
    static const std::vector<std::string> qwen35_re = {
        "(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])|"
        "[^\\r\\n\\p{L}\\p{N}]?[\\p{L}\\p{M}]+|\\p{N}|"
        " ?[^\\s\\p{L}\\p{M}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
    };
    auto words = unicode_regex_split(text, qwen35_re, /*byte_encode=*/true);
    for (const auto& w : words) bpe_word(this, w, out);

    if (add_special && add_eos) out->push_back(eos_id);
    return out->size();
}

std::string tokenizer::decode(const int32_t* tokens, size_t n,
                              bool skip_special) const {
    std::string out;
    out.reserve(n * 4);
    for (size_t i = 0; i < n; i++) {
        int32_t id = tokens[i];
        if (id < 0 || (size_t)id >= tokens_.size()) continue;
        if (skip_special) {
            uint8_t tt = type_[id];
            /* 1=NORMAL, 4=USER: emit.  2=UNKNOWN,3=CONTROL,5=UNUSED: skip. */
            if (tt != 1 && tt != 4) {
                if (tt == 6 /* BYTE */) { /* still emit byte tokens */ }
                else continue;
            }
        }
        /* tokens are stored in byte-encoded form; unicode_utf8_to_byte reverses
         * one UTF-8 char (one encoded byte) at a time. */
        const std::string& s = tokens_[id];
        for (size_t j = 0; j < s.size(); ) {
            size_t cl = std::min(s.size() - j, (size_t)unicode_len_utf8(s[j]));
            uint8_t b = unicode_utf8_to_byte(s.substr(j, cl));
            out.push_back((char)b);
            j += cl;
        }
    }
    return out;
}

int32_t tokenizer::find_by_text(const std::string& text) const {
    /* token_index_ excludes BYTE tokens (which never match multi-char text)
     * and empty tokens, so a direct lookup is correct for special markers. */
    const auto& idx = token_index_;
    auto it = std::lower_bound(idx.begin(), idx.end(), std::make_pair(text, 0),
        [](const std::pair<std::string, int32_t>& a,
           const std::pair<std::string, int32_t>& b) {
            return a.first < b.first;
        });
    return (it != idx.end() && it->first == text) ? it->second : -1;
}

/* M2.4: incremental UTF-8-safe detokenizer.  See qxmx_tokenizer.h.
 * The byte-decode of one token mirrors tokenizer::decode's inner loop; the
 * only addition is the UTF-8 boundary check on the flush. */

static void decode_token_bytes(const tokenizer& tok, int32_t id,
                               bool skip_special, std::string* out) {
    const auto& tokens = tok.tokens();
    const auto& types  = tok.token_types();
    if (id < 0 || (size_t)id >= tokens.size()) return;
    if (skip_special) {
        uint8_t tt = types[id];
        if (tt != 1 && tt != 4 && tt != 6) return;  /* skip CONTROL/UNKNOWN/UNUSED */
    }
    const std::string& s = tokens[id];
    for (size_t j = 0; j < s.size(); ) {
        size_t cl = std::min(s.size() - j, (size_t)unicode_len_utf8(s[j]));
        uint8_t b = unicode_utf8_to_byte(s.substr(j, cl));
        out->push_back((char)b);
        j += cl;
    }
}

/* Flush whole UTF-8 codepoints from the front of buf, leaving any trailing
 * partial sequence buffered.  Returns the flushed text. */
static std::string flush_utf8(std::string& buf) {
    std::string out;
    size_t i = 0;
    while (i < buf.size()) {
        size_t need = (size_t)unicode_len_utf8(buf[i]);
        if (need == 0) need = 1;   /* invalid lead byte — emit 1, resync */
        if (i + need > buf.size()) break;  /* partial codepoint — stop */
        out.append(buf, i, need);
        i += need;
    }
    if (i > 0) buf.erase(0, i);
    return out;
}

std::string incremental_decoder::push(int32_t token_id) {
    decode_token_bytes(tok_, token_id, skip_special_, &buf_);
    return flush_utf8(buf_);
}

std::string incremental_decoder::finish() {
    /* No more tokens coming; flush whatever remains (best-effort).  A
     * well-formed stream ends on a codepoint boundary, so this is usually
     * empty.  If a partial remains, emit it raw rather than dropping bytes. */
    std::string out;
    out.swap(buf_);
    return out;
}

}  // namespace qx