import json
import os
import tiktoken
from transformers import PreTrainedTokenizer


class axiomTokenizer(PreTrainedTokenizer):
    vocab_files_names = {"tokenizer_file": "tokenizer.model"}
    model_input_names = ["input_ids", "attention_mask"]

    def __init__(
        self,
        tokenizer_file=None,
        encoding_name="cl100k_base",
        eos_token="<|endoftext|>",
        pad_token="<|endoftext|>",
        unk_token="<|unk|>",
        **kwargs,
    ):
        if tokenizer_file and os.path.isfile(tokenizer_file):
            with open(tokenizer_file, "r", encoding="utf-8") as f:
                payload = json.load(f)
            encoding_name = payload.get("encoding_name", encoding_name)
        self.encoding_name = encoding_name
        self._enc = tiktoken.get_encoding(self.encoding_name)
        super().__init__(
            eos_token=eos_token,
            pad_token=pad_token,
            unk_token=unk_token,
            **kwargs,
        )

    @property
    def vocab_size(self):
        return int(self._enc.n_vocab)

    def get_vocab(self):
        return {f"<|{i}|>": i for i in range(self.vocab_size)}

    def _tokenize(self, text, **kwargs):
        ids = self._enc.encode_ordinary(text)
        return [f"<|{i}|>" for i in ids]

    def _convert_token_to_id(self, token):
        if token == self.eos_token or token == self.pad_token:
            return int(self._enc.eot_token)
        if token.startswith("<|") and token.endswith("|>"):
            n = token[2:-2]
            if n.isdigit():
                return int(n)
        return int(self._enc.eot_token)

    def _convert_id_to_token(self, index):
        return f"<|{int(index)}|>"

    def convert_tokens_to_string(self, tokens):
        ids = [self._convert_token_to_id(t) for t in tokens]
        return self._enc.decode(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):
        os.makedirs(save_directory, exist_ok=True)
        out_name = "tokenizer.model" if filename_prefix is None else f"{filename_prefix}-tokenizer.model"
        out_path = os.path.join(save_directory, out_name)
        with open(out_path, "w", encoding="utf-8") as f:
            json.dump({"encoding_name": self.encoding_name}, f)
        return (out_path,)