import os
import shutil

import sentencepiece as spm
from transformers import PreTrainedTokenizer


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

    def __init__(self, vocab_file, **kwargs):
        self.vocab_file = vocab_file
        self.sp_model = spm.SentencePieceProcessor(model_file=vocab_file)
        bos_token = kwargs.pop("bos_token", "<s>")
        eos_token = kwargs.pop("eos_token", "</s>")
        unk_token = kwargs.pop("unk_token", "<unk>")
        pad_token = kwargs.pop("pad_token", "<pad>")
        super().__init__(bos_token=bos_token, eos_token=eos_token, unk_token=unk_token, pad_token=pad_token, **kwargs)

    @property
    def vocab_size(self):
        return self.sp_model.vocab_size()

    def get_vocab(self):
        return {self.sp_model.id_to_piece(index): index for index in range(self.vocab_size)}

    def _tokenize(self, text):
        return self.sp_model.encode(text, out_type=str)

    def _convert_token_to_id(self, token):
        return self.sp_model.piece_to_id(token)

    def _convert_id_to_token(self, index):
        return self.sp_model.id_to_piece(index)

    def convert_tokens_to_string(self, tokens):
        return self.sp_model.decode(tokens)

    def save_vocabulary(self, save_directory, filename_prefix=None):
        filename = ((filename_prefix + "-") if filename_prefix else "") + "tokenizer.model"
        destination = os.path.join(save_directory, filename)
        if os.path.abspath(self.vocab_file) != os.path.abspath(destination):
            shutil.copyfile(self.vocab_file, destination)
        return (destination,)