"""Export a ZENITH BPE tokenizer to the Hugging Face `tokenizers` format.

The fast encoder only executes ZENITH's own vocabulary and merges; parity with
`BPETokenizer.encode` is enforced by tests.
"""

from __future__ import annotations

from pathlib import Path

from tokenizers import AddedToken, Regex, Tokenizer, decoders, models, pre_tokenizers

from zenith.tokenizer.bpe import PRETOKENIZE_PATTERN, SPECIAL_TOKENS, BPETokenizer


def _bytes_to_unicode() -> dict[int, str]:
    """GPT-2's reversible byte -> printable character map, used by HF ByteLevel."""
    bs = list(range(ord("!"), ord("~") + 1)) + list(range(ord("¡"), ord("¬") + 1)) + list(range(ord("®"), ord("ÿ") + 1))
    cs = bs[:]
    n = 0
    for b in range(256):
        if b not in bs:
            bs.append(b)
            cs.append(256 + n)
            n += 1
    return {b: chr(c) for b, c in zip(bs, cs)}


_BYTE_CHARS = _bytes_to_unicode()


def _as_unicode(b: bytes) -> str:
    return "".join(_BYTE_CHARS[x] for x in b)


def export_hf(tok: BPETokenizer, path: Path) -> None:
    vocab: dict[str, int] = {}
    for i in range(256 + len(tok.merges)):
        s = _as_unicode(tok.token_bytes(i))
        if s in vocab:
            raise ValueError(f"token {i} duplicates the bytes of token {vocab[s]}; HF BPE cannot represent this")
        vocab[s] = i
    for s, i in tok.special_ids.items():
        vocab[s] = i
    merges = [(_as_unicode(tok.token_bytes(a)), _as_unicode(tok.token_bytes(b))) for a, b in tok.merges]

    fast = Tokenizer(models.BPE(vocab=vocab, merges=merges))
    fast.pre_tokenizer = pre_tokenizers.Sequence(
        [
            pre_tokenizers.Split(Regex(PRETOKENIZE_PATTERN), behavior="isolated"),
            pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=False),
        ]
    )
    fast.decoder = decoders.ByteLevel()
    fast.add_special_tokens([AddedToken(s, special=True, normalized=False) for s in SPECIAL_TOKENS])
    fast.save(str(path))


def load_fast(path: Path) -> Tokenizer:
    fast = Tokenizer.from_file(str(path))
    # Documents that contain special-token text must encode as ordinary bytes.
    fast.encode_special_tokens = True
    return fast


def encode_fast(fast: Tokenizer, texts: list[str]) -> list[list[int]]:
    return [e.ids for e in fast.encode_batch(texts, add_special_tokens=False)]