"""Transformers adapter preserving literal text and the original byte BPE."""
from pathlib import Path
from transformers import PreTrainedTokenizer
from .purrence_core import Sub3MTokenizer


class PurrenceTokenizer(PreTrainedTokenizer):
    vocab_files_names = {"vocab_file": "tokenizer/ordinary-tokenizer.json",
                        "artifact_file": "tokenizer/artifact.json"}
    model_input_names = ["input_ids", "attention_mask"]

    def __init__(self, vocab_file, artifact_file, **kwargs):
        self.native = Sub3MTokenizer(Path(artifact_file).parent)
        self.vocab_file = vocab_file
        self.artifact_file = artifact_file
        self.vocab = self.native.backend.get_vocab()
        self.vocab["<|endoftext|>"] = 2047
        self.reverse_vocab = {v: k for k, v in self.vocab.items()}
        kwargs.setdefault("bos_token", "<|endoftext|>")
        kwargs.setdefault("eos_token", "<|endoftext|>")
        kwargs.setdefault("pad_token", "<|endoftext|>")
        kwargs.setdefault("model_max_length", 8192)
        kwargs["split_special_tokens"] = True
        kwargs["clean_up_tokenization_spaces"] = False
        super().__init__(**kwargs)

    @property
    def vocab_size(self):
        return 2048

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

    def _tokenize(self, text, **kwargs):
        return [self.reverse_vocab[i] for i in self.native.encode(text)]

    def _convert_token_to_id(self, token):
        return self.vocab[token]

    def _convert_id_to_token(self, index):
        return self.reverse_vocab[index]

    def convert_tokens_to_string(self, tokens):
        return self.native.decode([self.vocab[t] for t in tokens])

    def _decode(self, token_ids, skip_special_tokens=False, **kwargs):
        if isinstance(token_ids, int):
            token_ids = [token_ids]
        return self.native.decode([int(i) for i in token_ids], skip_boundaries=skip_special_tokens)

    def save_vocabulary(self, save_directory, filename_prefix=None):
        import shutil
        if filename_prefix:
            raise ValueError("This fixed tokenizer uses canonical filenames")
        target = Path(save_directory) / "tokenizer"
        target.mkdir(parents=True, exist_ok=True)
        paths = []
        for source in [self.vocab_file, self.artifact_file]:
            dest = target / Path(source).name
            if Path(source).resolve() != dest.resolve():
                shutil.copyfile(source, dest)
            paths.append(str(dest))
        return tuple(paths)