"""
MBG 1.0 — src/tokenizer_bpe.py
Byte-Pair-Encoding tokenizer (English-only) using the `tokenizers` library.

Rung L1 upgrade from the L0 char-level tokenizer. The tokenizer is trained on
the curated corpus, saved as a JSON vocab artifact (kept in snapshot), and
wraps a minimal encode/decode interface compatible with the training loop.
"""

from __future__ import annotations

import os
from typing import List

from tokenizers import Tokenizer, models, pre_tokenizers, trainers, decoders


class BpeTokenizer:
    """Thin wrapper around a HuggingFace `tokenizers` BPE."""

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

    def __init__(self, vocab_size: int = 512) -> None:
        self._target_vocab = vocab_size
        self.tok = Tokenizer(models.BPE())
        self.tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
        self.tok.decoder = decoders.ByteLevel()

    def train(self, texts: List[str]) -> None:
        trainer = trainers.BpeTrainer(
            vocab_size=self._target_vocab,
            special_tokens=self.SPECIAL,
            show_progress=False,
        )
        self.tok.train_from_iterator(texts, trainer=trainer)

    def save(self, path: str) -> None:
        os.makedirs(os.path.dirname(path), exist_ok=True)
        self.tok.save(path)

    def load(self, path: str) -> None:
        self.tok = Tokenizer.from_file(path)

    @property
    def vocab_size(self) -> int:
        return self.tok.get_vocab_size()

    def _special_id(self, tok: str) -> int:
        try:
            return self.tok.token_to_id(tok)
        except Exception:
            return 0

    def encode(self, text: str, add_bos_eos: bool = True) -> List[int]:
        ids = self.tok.encode(text).ids
        out = []
        if add_bos_eos:
            out.append(self._special_id("<bos>"))
        out.extend(ids)
        if add_bos_eos:
            out.append(self._special_id("<eos>"))
        return out

    def decode(self, ids: List[int]) -> str:
        return self.tok.decode(ids)

    def pad_id(self) -> int:
        return self._special_id("<pad>")