"""Byte-level BPE tokenizer for rurtech.ai MoE.

Small, dependency-free, and self-contained: trains a byte-level BPE on the
corpus and serializes to JSON. Byte-level means it never emits <unk> — any
input is representable. Special tokens: <pad> <bos> <eos>.
"""

from __future__ import annotations

import json
from collections import Counter
from pathlib import Path

SPECIAL = ["<pad>", "<bos>", "<eos>"]


class ByteBPETokenizer:
    def __init__(self, merges=None, vocab=None):
        self.merges: list[tuple[int, int]] = merges or []
        self.vocab: dict[int, bytes] = vocab or {}
        self.pad_id = 0
        self.bos_id = 1
        self.eos_id = 2

    # -- training ----------------------------------------------------------

    @classmethod
    def train(cls, text: str, vocab_size: int) -> "ByteBPETokenizer":
        # Base vocab: 3 specials + 256 byte values.
        vocab: dict[int, bytes] = {i: tok.encode() for i, tok in enumerate(SPECIAL)}
        for b in range(256):
            vocab[len(SPECIAL) + b] = bytes([b])

        ids = [len(SPECIAL) + b for b in text.encode("utf-8")]
        merges: list[tuple[int, int]] = []
        next_id = len(vocab)

        while next_id < vocab_size:
            pairs = Counter(zip(ids, ids[1:]))
            if not pairs:
                break
            (a, b), count = pairs.most_common(1)[0]
            if count < 2:
                break
            merges.append((a, b))
            vocab[next_id] = vocab[a] + vocab[b]
            ids = _merge(ids, a, b, next_id)
            next_id += 1

        return cls(merges=merges, vocab=vocab)

    # -- encode / decode ---------------------------------------------------

    def encode(self, text: str, add_bos=False, add_eos=False) -> list[int]:
        ids = [len(SPECIAL) + b for b in text.encode("utf-8")]
        for i, (a, b) in enumerate(self.merges):
            ids = _merge(ids, a, b, len(SPECIAL) + 256 + i)
        if add_bos:
            ids = [self.bos_id] + ids
        if add_eos:
            ids = ids + [self.eos_id]
        return ids

    def decode(self, ids: list[int]) -> str:
        out = b""
        for i in ids:
            if i in (self.pad_id, self.bos_id, self.eos_id):
                continue
            out += self.vocab.get(i, b"")
        return out.decode("utf-8", errors="replace")

    @property
    def vocab_size(self) -> int:
        return len(self.vocab)

    # -- serialization -----------------------------------------------------

    def save(self, path: Path) -> None:
        data = {
            "special_tokens": SPECIAL,
            "merges": self.merges,
            "vocab": {str(k): list(v) for k, v in self.vocab.items()},
        }
        Path(path).write_text(json.dumps(data), encoding="utf-8")

    @classmethod
    def load(cls, path: Path) -> "ByteBPETokenizer":
        data = json.loads(Path(path).read_text(encoding="utf-8"))
        vocab = {int(k): bytes(v) for k, v in data["vocab"].items()}
        merges = [tuple(m) for m in data["merges"]]
        return cls(merges=merges, vocab=vocab)


def _merge(ids: list[int], a: int, b: int, new_id: int) -> list[int]:
    out, i = [], 0
    while i < len(ids):
        if i < len(ids) - 1 and ids[i] == a and ids[i + 1] == b:
            out.append(new_id)
            i += 2
        else:
            out.append(ids[i])
            i += 1
    return out