from __future__ import annotations

import json
from pathlib import Path

from transformers import PreTrainedTokenizer


CANONICAL_VOCAB = {"<BOS>": 0, "+": 1, "=": 2, **{str(digit): digit + 3 for digit in range(10)}}


class AdditionTokenizer(PreTrainedTokenizer):
    vocab_files_names = {"vocab_file": "vocab.json"}
    model_input_names = ["input_ids", "attention_mask"]

    def __init__(self, vocab_file: str | None = None, **kwargs) -> None:
        if vocab_file is None:
            vocab = dict(CANONICAL_VOCAB)
        else:
            with Path(vocab_file).open("r", encoding="utf-8") as handle:
                vocab = json.load(handle)
        if vocab != CANONICAL_VOCAB:
            raise ValueError("AdditionTokenizer requires the canonical 13-token vocabulary.")
        self._vocab = vocab
        self._ids_to_tokens = {token_id: token for token, token_id in vocab.items()}
        kwargs.pop("bos_token", None)
        kwargs.pop("eos_token", None)
        kwargs.pop("pad_token", None)
        kwargs.pop("unk_token", None)
        super().__init__(
            bos_token="<BOS>",
            eos_token=None,
            pad_token=None,
            unk_token=None,
            **kwargs,
        )

    @property
    def vocab_size(self) -> int:
        return len(self._vocab)

    def get_vocab(self) -> dict[str, int]:
        return dict(self._vocab)

    def _tokenize(self, text: str, **kwargs) -> list[str]:
        compact = "".join(text.split())
        invalid = sorted(set(compact) - set("0123456789+="))
        if invalid:
            raise ValueError(f"Unsupported characters for addition tokenizer: {''.join(invalid)}")
        return list(compact)

    def _convert_token_to_id(self, token: str) -> int:
        try:
            return self._vocab[token]
        except KeyError as exc:
            raise ValueError(f"Unknown addition token: {token!r}") from exc

    def _convert_id_to_token(self, index: int) -> str:
        try:
            return self._ids_to_tokens[index]
        except KeyError as exc:
            raise ValueError(f"Unknown addition token ID: {index}") from exc

    def convert_tokens_to_string(self, tokens: list[str]) -> str:
        return "".join(tokens)

    def build_inputs_with_special_tokens(
        self,
        token_ids_0: list[int],
        token_ids_1: list[int] | None = None,
    ) -> list[int]:
        if token_ids_1 is not None:
            raise ValueError("AdditionTokenizer does not support sequence pairs.")
        return [self.bos_token_id, *token_ids_0]

    def get_special_tokens_mask(
        self,
        token_ids_0: list[int],
        token_ids_1: list[int] | None = None,
        already_has_special_tokens: bool = False,
    ) -> list[int]:
        if already_has_special_tokens:
            return [int(token_id == self.bos_token_id) for token_id in token_ids_0]
        if token_ids_1 is not None:
            raise ValueError("AdditionTokenizer does not support sequence pairs.")
        return [1, *([0] * len(token_ids_0))]

    def create_token_type_ids_from_sequences(
        self,
        token_ids_0: list[int],
        token_ids_1: list[int] | None = None,
    ) -> list[int]:
        if token_ids_1 is not None:
            raise ValueError("AdditionTokenizer does not support sequence pairs.")
        return [0] * (len(token_ids_0) + 1)

    def save_vocabulary(
        self,
        save_directory: str,
        filename_prefix: str | None = None,
    ) -> tuple[str]:
        directory = Path(save_directory)
        directory.mkdir(parents=True, exist_ok=True)
        filename = f"{filename_prefix + '-' if filename_prefix else ''}vocab.json"
        path = directory / filename
        path.write_text(json.dumps(self._vocab, indent=2, sort_keys=True) + "\n", encoding="utf-8")
        return (str(path),)


AdditionTokenizer.register_for_auto_class("AutoTokenizer")