#pragma once
// bpe.h, Qwen3/GPT-2 byte-level BPE tokenizer (CPU-only, no dependencies)
//
// Parses the tokenizer.json fields stored in a GGUF model and produces a
// byte-level BPE encoder/decoder. Arch-specific special tokens (text
// markers, language tags, audio sentinels) are loaded through
// bpe_load_specials_from_keys with a caller-provided list of GGUF KV keys.
// Loads vocab + merges from a GGUF tokenizer payload. Handles byte-level
// encoding, GPT-2 regex pre-tokenizer, BPE merges, and a registry of
// verbatim special tokens (endoftext plus any caller-registered tokens
// such as TTS style markers and language tags).

#include "gguf.h"

#include <cassert>
#include <climits>
#include <cstdio>
#include <cstring>
#include <string>
#include <unordered_map>
#include <vector>

// GPT-2 byte-level encoding table
// Maps byte [0..255] -> Unicode char for BPE vocab keys.
// Printable ASCII stays as-is, control/space bytes get remapped.
static void build_byte_encoder(std::string byte2str[256]) {
    // Standard GPT-2 byte encoder
    int bs[256], cs[256], n = 0, total = 0;
    // Printable ranges that map to themselves
    for (int b = '!'; b <= '~'; b++) {
        bs[total] = b;
        cs[total] = b;
        total++;
    }
    for (int b = 0xA1; b <= 0xAC; b++) {
        bs[total] = b;
        cs[total] = b;
        total++;
    }
    for (int b = 0xAE; b <= 0xFF; b++) {
        bs[total] = b;
        cs[total] = b;
        total++;
    }
    // Remaining bytes get mapped to 256+
    bool used[256] = {};
    for (int i = 0; i < total; i++) {
        used[bs[i]] = true;
    }
    for (int b = 0; b < 256; b++) {
        if (!used[b]) {
            bs[total] = b;
            cs[total] = 256 + n;
            n++;
            total++;
        }
    }
    assert(total == 256);
    // Convert codepoints to UTF-8 strings
    for (int i = 0; i < 256; i++) {
        int  cp = cs[i];
        char buf[4];
        int  len;
        if (cp < 0x80) {
            buf[0] = (char) cp;
            len    = 1;
        } else if (cp < 0x800) {
            buf[0] = (char) (0xC0 | (cp >> 6));
            buf[1] = (char) (0x80 | (cp & 0x3F));
            len    = 2;
        } else {
            buf[0] = (char) (0xE0 | (cp >> 12));
            buf[1] = (char) (0x80 | ((cp >> 6) & 0x3F));
            buf[2] = (char) (0x80 | (cp & 0x3F));
            len    = 3;
        }
        byte2str[bs[i]] = std::string(buf, len);
    }
}

// UTF-8 helpers
static int utf8_codepoint(const char * s, int * advance) {
    unsigned char c = s[0];
    if (c < 0x80) {
        *advance = 1;
        return c;
    }
    if ((c & 0xE0) == 0xC0) {
        *advance = 2;
        return ((c & 0x1F) << 6) | (s[1] & 0x3F);
    }
    if ((c & 0xF0) == 0xE0) {
        *advance = 3;
        return ((c & 0x0F) << 12) | ((s[1] & 0x3F) << 6) | (s[2] & 0x3F);
    }
    if ((c & 0xF8) == 0xF0) {
        *advance = 4;
        return ((c & 0x07) << 18) | ((s[1] & 0x3F) << 12) | ((s[2] & 0x3F) << 6) | (s[3] & 0x3F);
    }
    // invalid lead byte: advance one and return the raw byte to avoid an
    // infinite loop. The Python tokenizer never reaches this path since Python
    // str guarantees valid UTF 8; in C++ the std::string input has no such
    // guarantee, so this branch handles malformed input defensively.
    *advance = 1;
    return c;
}

// Unicode category checks (simplified but covers Latin + common scripts)
static bool is_letter(int cp) {
    if ((cp >= 'A' && cp <= 'Z') || (cp >= 'a' && cp <= 'z')) {
        return true;
    }
    if (cp < 0x80) {
        return false;
    }
    // Latin Extended: U+00C0-U+00D6, U+00D8-U+00F6, U+00F8-U+01BF + Latin Extended-A/B
    if (cp >= 0xC0 && cp <= 0x024F && cp != 0xD7 && cp != 0xF7) {
        return true;
    }
    // Common CJK, Cyrillic, Greek, Arabic, etc., treat as letters
    if (cp >= 0x0370 && cp <= 0x1FFF) {
        return true;  // Greek, Cyrillic, Armenian, etc.
    }
    if (cp >= 0x2C00 && cp <= 0x2DFF) {
        return true;  // Georgian, etc.
    }
    if (cp >= 0x3040 && cp <= 0x9FFF) {
        return true;  // CJK
    }
    if (cp >= 0xAC00 && cp <= 0xD7AF) {
        return true;  // Korean
    }
    if (cp >= 0xF900 && cp <= 0xFAFF) {
        return true;  // CJK compatibility
    }
    if (cp >= 0x10000) {
        return true;  // SMP, mostly letters/symbols
    }
    return false;
}

static bool is_digit(int cp) {
    return cp >= '0' && cp <= '9';
}

static bool is_whitespace(int cp) {
    return cp == ' ' || cp == '\t' || cp == '\n' || cp == '\r' || cp == 0x0B || cp == 0x0C || cp == 0xA0 ||
           cp == 0x2000 || cp == 0x2001 || cp == 0x2002 || cp == 0x200B;
}

static bool is_newline(int cp) {
    return cp == '\n' || cp == '\r';
}

// GPT-2 pre-tokenizer regex (manual implementation)
// Pattern: (?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}|
//          \s?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+
// Splits text into non-overlapping chunks (on original text, not byte-encoded).
static std::vector<std::string> gpt2_pre_tokenize(const std::string & text) {
    std::vector<std::string> chunks;
    const char *             s   = text.c_str();
    int                      len = (int) text.size();
    int                      i   = 0;

    while (i < len) {
        int adv;
        int cp = utf8_codepoint(s + i, &adv);

        // Rule 1: Contractions 's 't 're 've 'm 'll 'd
        if ((cp == '\'' || cp == 0x2019) && i + adv < len) {
            const char * rest      = s + i + adv;
            int          rlen      = len - i - adv;
            auto         try_match = [&](const char * suffix, int slen) -> bool {
                if (rlen >= slen) {
                    // case-insensitive compare
                    for (int k = 0; k < slen; k++) {
                        char c1 = rest[k], c2 = suffix[k];
                        if (c1 >= 'A' && c1 <= 'Z') {
                            c1 = (char) (c1 + 32);
                        }
                        if (c1 != c2) {
                            return false;
                        }
                    }
                    // next char should NOT be a letter
                    if (rlen > slen) {
                        int a2;
                        int cp2 = utf8_codepoint(rest + slen, &a2);
                        if (is_letter(cp2)) {
                            return false;
                        }
                    }
                    chunks.push_back(std::string(s + i, adv + slen));
                    i += adv + slen;
                    return true;
                }
                return false;
            };
            if (try_match("ll", 2)) {
                continue;
            }
            if (try_match("re", 2)) {
                continue;
            }
            if (try_match("ve", 2)) {
                continue;
            }
            if (try_match("s", 1)) {
                continue;
            }
            if (try_match("t", 1)) {
                continue;
            }
            if (try_match("m", 1)) {
                continue;
            }
            if (try_match("d", 1)) {
                continue;
            }
        }

        // Rule 2: [^\r\n\p{L}\p{N}]?\p{L}+
        if (is_letter(cp)) {
            int start = i;
            i += adv;
            while (i < len) {
                int a2;
                int cp2 = utf8_codepoint(s + i, &a2);
                if (!is_letter(cp2)) {
                    break;
                }
                i += a2;
            }
            chunks.push_back(std::string(s + start, i - start));
            continue;
        }
        if (!is_newline(cp) && !is_letter(cp) && !is_digit(cp) && !is_whitespace(cp)) {
            // Non-letter/number/space, check if followed by letters
            int start = i;
            int after = i + adv;
            if (after < len) {
                int a2;
                int cp2 = utf8_codepoint(s + after, &a2);
                if (is_letter(cp2)) {
                    i = after + a2;
                    while (i < len) {
                        int a3;
                        int cp3 = utf8_codepoint(s + i, &a3);
                        if (!is_letter(cp3)) {
                            break;
                        }
                        i += a3;
                    }
                    chunks.push_back(std::string(s + start, i - start));
                    continue;
                }
            }
        }

        // Rule 3: \p{N}+  (digits, consume consecutively)
        if (is_digit(cp)) {
            int start = i;
            while (i < len && is_digit((unsigned char) s[i])) {
                i++;
            }
            // GPT-2 regex matches single \p{N}, let's match one at a time
            // to be safe, but in practice consecutive digits usually merge anyway.
            // The regex is \p{N} (single digit), so split each digit:
            for (int j = start; j < i; j++) {
                chunks.push_back(std::string(s + j, 1));
            }
            continue;
        }

        // Rule 5: \s*[\r\n]+  (newlines with optional leading whitespace)
        if (is_newline(cp)) {
            int start = i;
            while (i < len && is_newline((unsigned char) s[i])) {
                i++;
            }
            chunks.push_back(std::string(s + start, i - start));
            continue;
        }

        // Rule 6: whitespace handling
        // Regex order: \s+(?!\S) first (trailing whitespace), then \s+ as fallback
        // \s+(?!\S) backtracks: consumes whitespace NOT followed by non-whitespace
        // This peels off leading spaces, leaving the last space to combine with the next word
        if (is_whitespace(cp)) {
            int start  = i;
            // Find end of whitespace run
            int ws_end = i + adv;
            while (ws_end < len && is_whitespace((unsigned char) s[ws_end]) && !is_newline((unsigned char) s[ws_end])) {
                ws_end++;
            }
            // Check what follows the whitespace run
            bool followed_by_non_ws =
                (ws_end < len && !is_whitespace((unsigned char) s[ws_end]) && !is_newline((unsigned char) s[ws_end]));
            if (followed_by_non_ws && ws_end - start > 1) {
                // \s+(?!\S) matches all but the last space
                // Leave one space for the next iteration to combine with word
                int trailing = ws_end - 1;
                chunks.push_back(std::string(s + start, trailing - start));
                i = trailing;
                continue;
            }
            // Single space followed by word: combine space + word as one chunk
            i = start + adv;
            if (i < len) {
                int a2;
                int cp2 = utf8_codepoint(s + i, &a2);
                if (is_letter(cp2)) {
                    i += a2;
                    while (i < len) {
                        int a3;
                        int cp3 = utf8_codepoint(s + i, &a3);
                        if (!is_letter(cp3)) {
                            break;
                        }
                        i += a3;
                    }
                    chunks.push_back(std::string(s + start, i - start));
                    continue;
                }
                if (is_digit(cp2)) {
                    chunks.push_back(std::string(s + start, i - start));
                    continue;
                }
                if (!is_whitespace(cp2) && !is_newline(cp2)) {
                    int pstart = start;
                    while (i < len) {
                        int a3;
                        int cp3 = utf8_codepoint(s + i, &a3);
                        if (is_whitespace(cp3) || is_letter(cp3) || is_digit(cp3)) {
                            break;
                        }
                        i += a3;
                    }
                    while (i < len && is_newline((unsigned char) s[i])) {
                        i++;
                    }
                    chunks.push_back(std::string(s + pstart, i - pstart));
                    continue;
                }
            }
            // Trailing whitespace (end of string or before newline), consume all
            i = ws_end;
            while (i < len) {
                int a2;
                int cp2 = utf8_codepoint(s + i, &a2);
                if (!is_whitespace(cp2)) {
                    break;
                }
                i += a2;
            }
            chunks.push_back(std::string(s + start, i - start));
            continue;
        }

        // Rule 4: [^\s\p{L}\p{N}]+[\r\n]*  (punctuation/symbols)
        {
            int start = i;
            i += adv;
            while (i < len) {
                int a2;
                int cp2 = utf8_codepoint(s + i, &a2);
                if (is_whitespace(cp2) || is_letter(cp2) || is_digit(cp2) || is_newline(cp2)) {
                    break;
                }
                i += a2;
            }
            // trailing newlines
            while (i < len && is_newline((unsigned char) s[i])) {
                i++;
            }
            chunks.push_back(std::string(s + start, i - start));
        }
    }
    return chunks;
}

// BPE tokenizer struct
struct BPETokenizer {
    std::unordered_map<std::string, int> vocab;          // token_str -> id
    std::unordered_map<std::string, int> merges;         // "a b" -> rank
    std::string                          byte2str[256];  // byte -> GPT-2 UTF-8 string
    int                                  eos_id;         // <|endoftext|>
    int                                  n_vocab;
    std::vector<std::string>             id_to_str;      // id -> token_str (reverse vocab)

    // Registered special tokens. Each (str, id) pair is matched verbatim in
    // bpe_encode and emitted as a single id, bypassing the BPE merge passes.
    std::vector<std::pair<std::string, int>> specials;
};

// Register a special token. Strings already registered are skipped (no dup).
static void bpe_add_special(BPETokenizer * tok, const std::string & str, int id) {
    for (const auto & sp : tok->specials) {
        if (sp.first == str) {
            return;
        }
    }
    tok->specials.emplace_back(str, id);
}

// Load tokenizer from GGUF KV (tokenizer.ggml.tokens + tokenizer.ggml.merges)
static bool load_bpe_from_gguf(BPETokenizer * tok, const char * gguf_path) {
    build_byte_encoder(tok->byte2str);

    struct gguf_init_params gp  = { true, NULL };
    struct gguf_context *   ctx = gguf_init_from_file(gguf_path, gp);
    if (!ctx) {
        fprintf(stderr, "[BPE] Failed to open %s\n", gguf_path);
        return false;
    }

    int64_t tok_key = gguf_find_key(ctx, "tokenizer.ggml.tokens");
    int64_t mrg_key = gguf_find_key(ctx, "tokenizer.ggml.merges");
    if (tok_key < 0 || mrg_key < 0) {
        fprintf(stderr, "[BPE] Tokenizer not found in %s\n", gguf_path);
        gguf_free(ctx);
        return false;
    }

    int n_tokens = (int) gguf_get_arr_n(ctx, tok_key);
    int n_merges = (int) gguf_get_arr_n(ctx, mrg_key);

    for (int i = 0; i < n_tokens; i++) {
        const char * s             = gguf_get_arr_str(ctx, tok_key, (size_t) i);
        tok->vocab[std::string(s)] = i;
    }

    for (int i = 0; i < n_merges; i++) {
        const char * s              = gguf_get_arr_str(ctx, mrg_key, (size_t) i);
        tok->merges[std::string(s)] = i;
    }

    gguf_free(ctx);

    tok->n_vocab = (int) tok->vocab.size();

    tok->id_to_str.resize(tok->n_vocab);
    for (auto & kv : tok->vocab) {
        if (kv.second >= 0 && kv.second < tok->n_vocab) {
            tok->id_to_str[kv.second] = kv.first;
        }
    }

    // Resolve eos_id from the vocab itself rather than hard-coding 151643.
    // Falls back to -1 if the standard sentinel is absent.
    auto eos_it = tok->vocab.find("<|endoftext|>");
    tok->eos_id = (eos_it != tok->vocab.end()) ? eos_it->second : -1;
    if (tok->eos_id >= 0) {
        bpe_add_special(tok, "<|endoftext|>", tok->eos_id);
    }

    fprintf(stderr, "[BPE] Loaded from GGUF: %d vocab, %d merges, eos_id=%d\n", tok->n_vocab, n_merges, tok->eos_id);
    return true;
}

// Read arch-specific special tokens from a caller-provided list of GGUF KV
// keys. Each key holds a u32 vocab id, mapped back to its vocab string and
// registered through bpe_add_special. The endoftext sentinel is already
// registered by load_bpe_from_gguf, so callers should not list it here.
static bool bpe_load_specials_from_keys(BPETokenizer *       tok,
                                        const char *         gguf_path,
                                        const char * const * keys,
                                        int                  n_keys) {
    struct gguf_init_params gp  = { true, NULL };
    struct gguf_context *   ctx = gguf_init_from_file(gguf_path, gp);
    if (!ctx) {
        fprintf(stderr, "[BPE] Failed to open %s for specials\n", gguf_path);
        return false;
    }

    int n_added = 0;
    for (int i = 0; i < n_keys; i++) {
        int64_t k = gguf_find_key(ctx, keys[i]);
        if (k < 0) {
            fprintf(stderr, "[BPE] WARNING: missing %s in GGUF\n", keys[i]);
            continue;
        }
        int id = (int) gguf_get_val_u32(ctx, k);
        if (id < 0 || id >= tok->n_vocab) {
            fprintf(stderr, "[BPE] WARNING: %s id=%d out of vocab range\n", keys[i], id);
            continue;
        }
        const std::string & s = tok->id_to_str[id];
        if (s.empty()) {
            fprintf(stderr, "[BPE] WARNING: %s id=%d has empty vocab string\n", keys[i], id);
            continue;
        }
        bpe_add_special(tok, s, id);
        n_added++;
    }

    gguf_free(ctx);
    fprintf(stderr, "[BPE] Registered %d arch special tokens (total specials=%zu)\n", n_added, tok->specials.size());
    return true;
}

// Byte-level encode: raw text bytes -> GPT-2 BPE string
static std::string byte_level_encode(const BPETokenizer * tok, const std::string & text) {
    std::string out;
    for (unsigned char c : text) {
        out += tok->byte2str[c];
    }
    return out;
}

// BPE merge algorithm
// Input: list of symbols (strings). Merges pairs by priority.
static std::vector<std::string> bpe_merge(const std::unordered_map<std::string, int> & merge_rank,
                                          const std::vector<std::string> &             symbols) {
    if (symbols.size() <= 1) {
        return symbols;
    }

    std::vector<std::string> work = symbols;

    while (work.size() > 1) {
        // Find the pair with lowest rank (highest priority)
        int best_rank = INT_MAX;
        int best_pos  = -1;
        for (int i = 0; i < (int) work.size() - 1; i++) {
            std::string key = work[i] + " " + work[i + 1];
            auto        it  = merge_rank.find(key);
            if (it != merge_rank.end() && it->second < best_rank) {
                best_rank = it->second;
                best_pos  = i;
            }
        }
        if (best_pos < 0) {
            break;  // no more merges
        }

        // Merge the pair
        std::string merged = work[best_pos] + work[best_pos + 1];
        work[best_pos]     = merged;
        work.erase(work.begin() + best_pos + 1);
    }
    return work;
}

// Encode a single pre-tokenized chunk -> token ids
static void encode_chunk(const BPETokenizer * tok, const std::string & chunk, std::vector<int> & ids) {
    // Byte-level encode
    std::string encoded = byte_level_encode(tok, chunk);

    // Split into individual UTF-8 characters (each is a BPE symbol)
    std::vector<std::string> symbols;
    const char *             s   = encoded.c_str();
    int                      len = (int) encoded.size();
    int                      i   = 0;
    while (i < len) {
        int adv;
        utf8_codepoint(s + i, &adv);
        symbols.push_back(std::string(s + i, adv));
        i += adv;
    }

    // Apply BPE merges
    std::vector<std::string> merged = bpe_merge(tok->merges, symbols);

    // Look up in vocab
    for (const auto & piece : merged) {
        auto it = tok->vocab.find(piece);
        if (it != tok->vocab.end()) {
            ids.push_back(it->second);
        } else {
            // Fallback: encode each byte individually (should not happen with byte-level BPE)
            fprintf(stderr, "[BPE] WARNING: unknown token '%s'\n", piece.c_str());
            for (unsigned char c : piece) {
                auto it2 = tok->vocab.find(std::string(1, c));
                if (it2 != tok->vocab.end()) {
                    ids.push_back(it2->second);
                }
            }
        }
    }
}

// Full encode: text -> token ids.
// Walks the text from left to right, matching any registered special token
// verbatim. For each segment between specials, runs the GPT-2 byte-level
// pre-tokenizer + BPE merges. The endoftext sentinel is auto-registered as
// a special by load_bpe_from_gguf, so existing call sites that embed
// "<|endoftext|>" in the input text keep working.
// add_eos = true appends the eos_id at the end (post-processor behavior).
static std::vector<int> bpe_encode(const BPETokenizer * tok, const std::string & text, bool add_eos = true) {
    std::vector<int> ids;

    auto encode_segment = [&](const std::string & seg) {
        if (seg.empty()) {
            return;
        }
        auto chunks = gpt2_pre_tokenize(seg);
        for (const auto & chunk : chunks) {
            encode_chunk(tok, chunk, ids);
        }
    };

    size_t pos = 0;
    while (pos < text.size()) {
        // Find the leftmost occurrence of any registered special token.
        size_t best_pos = std::string::npos;
        int    best_idx = -1;
        for (size_t i = 0; i < tok->specials.size(); i++) {
            size_t p = text.find(tok->specials[i].first, pos);
            if (p != std::string::npos && p < best_pos) {
                best_pos = p;
                best_idx = (int) i;
            }
        }

        if (best_idx < 0) {
            encode_segment(text.substr(pos));
            break;
        }

        if (best_pos > pos) {
            encode_segment(text.substr(pos, best_pos - pos));
        }
        const auto & sp = tok->specials[(size_t) best_idx];
        ids.push_back(sp.second);
        pos = best_pos + sp.first.size();
    }

    if (add_eos && tok->eos_id >= 0) {
        ids.push_back(tok->eos_id);
    }
    return ids;
}