"""alpha-er's byte-level BPE tokenizer (vocab 12,288).

GPT-2-style byte-level BPE, with the vocabulary laid out as:

    0   - 255   raw bytes, in byte order, under GPT-2's bytes->unicode mapping
    256 - 258   <|user|>, <|assistant|>, <|end_of_text|>
    259 +       learned merges, where merges[i] produces token 259 + i

Being byte-level means every input encodes — there is no unknown token and no
normalisation step to disagree about. Decoding maps each vocabulary character
back to its byte and interprets the result as UTF-8, so a sequence cut mid
multi-byte character degrades to a replacement character rather than throwing.
"""
from __future__ import annotations
import json
from functools import lru_cache


@lru_cache(1)
def _byte_maps():
    """GPT-2's reversible bytes<->unicode table."""
    bs = (list(range(ord("!"), ord("~") + 1))
          + list(range(ord("\xa1"), ord("\xac") + 1))
          + list(range(ord("\xae"), ord("\xff") + 1)))
    cs = bs[:]
    n = 0
    for b in range(256):
        if b not in bs:
            bs.append(b); cs.append(256 + n); n += 1
    b2u = {b: chr(c) for b, c in zip(bs, cs)}
    return b2u, {v: k for k, v in b2u.items()}


class AlphaErTokenizer:
    def __init__(self, artifacts: dict):
        self.vocab = artifacts["vocab"]
        self.merges = [tuple(m) for m in artifacts["merges"]]
        self.specials = artifacts["specialTokens"]
        self.n_special = len(self.specials)
        self.first_merge_id = 256 + self.n_special
        # rank[(a,b)] = the token id the pair becomes; lower id = earlier merge.
        self.rank = {pair: self.first_merge_id + i for i, pair in enumerate(self.merges)}
        self.special_ids = {s: 256 + i for i, s in enumerate(self.specials)}

    @classmethod
    def from_file(cls, path: str) -> "AlphaErTokenizer":
        return cls(json.load(open(path)))

    def _encode_chunk(self, text: str) -> list[int]:
        b2u, _ = _byte_maps()
        ids = list(text.encode("utf-8"))            # ids 0-255 ARE the bytes
        while len(ids) > 1:
            best, best_at = None, -1
            for i in range(len(ids) - 1):
                r = self.rank.get((ids[i], ids[i + 1]))
                if r is not None and (best is None or r < best):
                    best, best_at = r, i
            if best is None:
                break
            ids[best_at:best_at + 2] = [best]
        return ids

    def encode(self, text: str) -> list[int]:
        """Split on special tokens first so they survive as single ids."""
        parts, out = [text], None
        for s in self.specials:
            nxt = []
            for p in parts:
                if isinstance(p, int):
                    nxt.append(p); continue
                bits = p.split(s)
                for i, bit in enumerate(bits):
                    if i: nxt.append(self.special_ids[s])
                    if bit: nxt.append(bit)
            parts = nxt
        out = []
        for p in parts:
            out.extend([p] if isinstance(p, int) else self._encode_chunk(p))
        return out

    def decode(self, ids) -> str:
        _, u2b = _byte_maps()
        buf, text = bytearray(), []
        for i in ids:
            tok = self.vocab[i]
            if i < 256 + self.n_special and i >= 256:        # a special token
                text.append(buf.decode("utf-8", "replace")); buf = bytearray()
                text.append(tok)
                continue
            for ch in tok:
                buf.append(u2b.get(ch, ord("?") if ord(ch) < 256 else 63))
        text.append(buf.decode("utf-8", "replace"))
        return "".join(text)