import json
from pathlib import Path

from transformers import PreTrainedTokenizer


class SupernovaTokenizer(PreTrainedTokenizer):

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

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

    def __init__(
        self,
        vocab_file,
        bos_token="<BOS>",
        eos_token="<EOS>",
        pad_token="<PAD>",
        user_token="<USER>",
        assistant_token="<ASSISTANT>",
        **kwargs,
    ):

        with open(
            vocab_file,
            "r",
            encoding="utf-8",
        ) as f:
            vocab = json.load(f)

        self.vocab = {
            str(token): int(idx)
            for token, idx in vocab.items()
        }

        self.ids_to_tokens = {
            int(idx): str(token)
            for token, idx in self.vocab.items()
        }

        self.user_token = user_token
        self.assistant_token = assistant_token

        super().__init__(
            bos_token=bos_token,
            eos_token=eos_token,
            pad_token=pad_token,
            **kwargs,
        )

    @property
    def vocab_size(self):
        return len(self.vocab)

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

    def _tokenize(self, text, **kwargs):
        return list(text)

    def _convert_token_to_id(self, token):

        if token not in self.vocab:
            raise ValueError(
                f"Unknown character/token: {token!r}"
            )

        return self.vocab[token]

    def _convert_id_to_token(self, index):

        index = int(index)

        if index not in self.ids_to_tokens:
            raise ValueError(
                f"Unknown token id: {index}"
            )

        return self.ids_to_tokens[index]

    def convert_tokens_to_string(self, tokens):
        return "".join(tokens)

    def build_inputs_with_special_tokens(
        self,
        token_ids_0,
        token_ids_1=None,
    ):

        result = [self.bos_token_id]
        result += list(token_ids_0)

        if token_ids_1 is not None:
            result += list(token_ids_1)

        result += [self.eos_token_id]

        return result

    def save_vocabulary(
        self,
        save_directory,
        filename_prefix=None,
    ):

        save_directory = Path(save_directory)

        save_directory.mkdir(
            parents=True,
            exist_ok=True,
        )

        filename = "vocab.json"

        if filename_prefix:
            filename = (
                f"{filename_prefix}-vocab.json"
            )

        path = save_directory / filename

        with open(
            path,
            "w",
            encoding="utf-8",
        ) as f:
            json.dump(
                self.vocab,
                f,
                ensure_ascii=False,
                indent=2,
            )

        return (str(path),)