import json
import os
from typing import List, Optional, Tuple

from transformers import PreTrainedTokenizer

VOCAB_FILES_NAMES = {"vocab_file": "vocab.json"}


class TinyGPTTokenizer(PreTrainedTokenizer):
    """Character-level tokenizer: each of the 128 ASCII code points is its
    own token, id == ord(char) — the exact scheme used by decode()/chr(t)
    in the original training script."""

    vocab_files_names = VOCAB_FILES_NAMES
    model_input_names = ["input_ids", "attention_mask"]

    def __init__(self, vocab_file: Optional[str] = None, **kwargs):
        self._vocab = {chr(i): i for i in range(128)}
        self._ids_to_tokens = {i: chr(i) for i in range(128)}
        super().__init__(**kwargs)

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

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

    def _tokenize(self, text: str, **kwargs) -> List[str]:
        return list(text)

    def _convert_token_to_id(self, token: str) -> int:
        return self._vocab.get(token, self._vocab.get(" "))

    def _convert_id_to_token(self, index: int) -> str:
        return self._ids_to_tokens.get(index, " ")

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

    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
        filename = (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
        vocab_path = os.path.join(save_directory, filename)
        with open(vocab_path, "w", encoding="utf-8") as f:
            json.dump(self._vocab, f, ensure_ascii=False, indent=2)
        return (vocab_path,)