"""
BPETokenizer — HuggingFace tokenizers wrapper
──────────────────────────────────────────────
Wraps a pre-trained BPE tokenizer loaded from a JSON file.

SPECIAL TOKENS:
  <pad>  — Padding token for batching sequences of different lengths.
            Not used during standard next-token-prediction training
            (all windows are the same length), but needed for inference batching.
  <unk>  — Unknown token. A well-trained BPE tokenizer should never emit this
            because byte-level BPE can represent any Unicode character.
  <bos>  — Beginning-of-Sequence. Prepended to each document during training
            so the model learns to associate it with "start of generation".
  <eos>  — End-of-Sequence. Appended to each document. The model is trained
            to predict this token when a natural endpoint is reached. During
            inference, generation stops when EOS is produced.

BYTE-LEVEL BPE:
  The tokenizer uses ByteLevel pre-tokenization (from train_tokenizer.py).
  This means every byte of the input is representable as a single token —
  no character is ever truly "unknown". The Ġ prefix (used in GPT-2 / RoBERTa)
  marks word boundaries: a token starting with Ġ follows a space. The clean()
  method converts these back to readable text.
"""

from tokenizers import Tokenizer, decoders


class BPETokenizer:
    """
    Wraps a HuggingFace `Tokenizer` JSON file for easy encode/decode.

    Args:
        tokenizer_path: Path to the saved tokenizer JSON
                        (produced by tokenizer/train_tokenizer.py).
    """

    def __init__(self, tokenizer_path: str = "dataProcessed/tokenizer.json"):
        self.tokenizer = Tokenizer.from_file(tokenizer_path)

        # Ensure ByteLevel decoder is attached so decode() properly maps
        # byte-level tokens (like Ġ for space and Ċ for newline) back to text.
        if self.tokenizer.decoder is None:
            self.tokenizer.decoder = decoders.ByteLevel()

        # Resolve special token IDs once at init — avoids repeated lookups
        self.pad_id = self.tokenizer.token_to_id("<pad>")
        self.unk_id = self.tokenizer.token_to_id("<unk>")
        self.bos_id = self.tokenizer.token_to_id("<bos>")
        self.eos_id = self.tokenizer.token_to_id("<eos>")

    def encode(self, text: str) -> list[int]:
        """Encode a string to a list of integer token IDs."""
        return self.tokenizer.encode(text).ids

    def decode(self, ids: list[int]) -> str:
        """Decode a list of token IDs back to a string."""
        return self.tokenizer.decode(ids)

    def clean(self, text: str) -> str:
        """
        Remove byte-level BPE artefacts and normalise whitespace.
        """
        return text.replace("Ġ", " ").replace("Ċ", "\n").strip()

    @property
    def vocab_size(self) -> int:
        """Total vocabulary size (including special tokens)."""
        return self.tokenizer.get_vocab_size()