"""Hugging Face adapter for Limen0.2B's SuperBPE tokenizer.

Install the Rust-backed tokenizer package before loading this tokenizer:

    pip install "git+https://github.com/UniversalComputingResearch/fastboundlessbpe.git@perf/tokenid-training"
"""

from __future__ import annotations

from pathlib import Path

from transformers import PreTrainedTokenizer

try:
    from boundlessbpe import FastTokenizer, RUST_AVAILABLE
    from boundlessbpe.vocabulary import Vocabulary
except ImportError as exc:  # pragma: no cover - depends on the consumer environment
    raise ImportError(
        "Limen0.2B requires the Rust-backed `boundlessbpe` package. Install it with: "
        'pip install "git+https://github.com/UniversalComputingResearch/fastboundlessbpe.git@perf/tokenid-training"'
    ) from exc


class SuperwordTokenizer(PreTrainedTokenizer):
    """Exact inference adapter for the SuperBPE model used in pretraining."""

    model_input_names = ["input_ids", "attention_mask"]
    vocab_files_names = {"superword_model_file": "superword.model"}

    def __init__(self, superword_model_file: str = "superword.model", **kwargs):
        if not RUST_AVAILABLE or FastTokenizer is None:
            raise RuntimeError(
                "`boundlessbpe` is installed without its Rust extension. Reinstall the "
                "package from https://github.com/UniversalComputingResearch/fastboundlessbpe/tree/perf/tokenid-training."
            )

        model_file = Path(superword_model_file)
        if not model_file.is_absolute():
            model_file = Path(kwargs.pop("name_or_path", ".")) / model_file
        self.superword_model_file = str(model_file)

        self._fast = FastTokenizer()
        self._fast.load(str(model_file))
        with model_file.open("r", encoding="utf-8") as model_handle:
            header = model_handle.readline().strip()
            if not header.startswith("BoundlessBPE v2 "):
                raise ValueError(f"Unsupported SuperBPE model header: {header!r}")
            self._vocabulary = Vocabulary.load(model_handle)

        self._special_tokens = dict(self._vocabulary.special_tokens)
        self._inverse_special_tokens = dict(self._vocabulary.inverse_special_tokens)
        self._vocab = {
            token.decode("utf-8", errors="replace"): int(token_id)
            for token, token_id in self._vocabulary.token_to_id.items()
        }
        self._vocab.update(self._special_tokens)

        model_max_length = int(kwargs.pop("model_max_length", 1024))
        for key in (
            "pad_token",
            "bos_token",
            "eos_token",
            "unk_token",
        ):
            kwargs.pop(key, None)
        super().__init__(
            pad_token="<|pad|>",
            bos_token="<|bos|>",
            eos_token="<|endoftext|>",
            unk_token="<|unk|>",
            model_max_length=model_max_length,
            **kwargs,
        )

    def get_vocab(self):
        return dict(self._vocab)

    @property
    def vocab_size(self):
        return int(self._fast.get_vocab_size(with_added_tokens=False))

    def _id_to_token(self, token_id: int) -> str:
        token = self._vocabulary.id_to_token.get(int(token_id))
        if token is not None:
            return token.decode("utf-8", errors="replace")
        return self._inverse_special_tokens.get(int(token_id), "<|unk|>")

    def _tokenize(self, text, **kwargs):
        return [self._id_to_token(token_id) for token_id in self._fast.encode_ordinary(text)]

    def _convert_token_to_id(self, token):
        return self._vocab.get(token, self._special_tokens["<|unk|>"])

    def _convert_id_to_token(self, index):
        return self._id_to_token(int(index))

    def encode(self, text, text_pair=None, add_special_tokens=False, **kwargs):
        if text_pair is not None:
            text = text + text_pair
        if add_special_tokens:
            return list(self._fast.encode(text, allowed_special="all"))
        return list(self._fast.encode_ordinary(text))

    def decode(self, token_ids, skip_special_tokens=True, **kwargs):
        if isinstance(token_ids, int):
            token_ids = [token_ids]
        if skip_special_tokens:
            token_ids = [
                token_id
                for token_id in token_ids
                if int(token_id) not in self._inverse_special_tokens
            ]
        return self._fast.decode(list(token_ids))

    def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
        if token_ids_1 is None:
            return list(token_ids_0)
        return list(token_ids_0) + list(token_ids_1)

    def save_vocabulary(self, save_directory, filename_prefix=None):
        target = Path(save_directory) / (filename_prefix or "")
        target = target.with_name(target.name + "superword.model")
        target.write_bytes(Path(self.superword_model_file).read_bytes())
        return (str(target),)