import os
import shutil

import sentencepiece as spm
from transformers import PreTrainedTokenizer

VOCAB_FILES_NAMES = {"vocab_file": "spm_16384.model"}


class G0NanoTokenizer(PreTrainedTokenizer):
    """Wrapper SentencePiece + tokens de chat additionnels (>= 16384).

    Les tokens de chat (``<|user|>``, ``<|assistant|>``, ``<|end|>``,
    ``<|system|>``) sont geres par le mecanisme "added tokens" standard de
    ``transformers`` (passe via ``additional_special_tokens``) plutot que par
    une logique maison : HF les decoupe avant tokenization et les reinsere
    correctement au decode, et leur assigne des ids sequentiels a partir de
    ``len(self)`` au moment du ``__init__`` -- ce qui reproduit exactement le
    mapping fige dans ``training/chat_format.py`` (16384..16387), du moment
    que la liste est fournie dans le meme ordre.
    """

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

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

        super().__init__(
            bos_token=bos_token,
            eos_token=eos_token,
            unk_token=unk_token,
            pad_token=pad_token,
            additional_special_tokens=additional_special_tokens or [],
            **kwargs,
        )

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

    def get_vocab(self):
        vocab = {self.sp_model.id_to_piece(i): i for i in range(self.vocab_size)}
        vocab.update(self.added_tokens_encoder)
        return vocab

    def _tokenize(self, text, **kwargs):
        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) if tokens else ""

    def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
        return [self.bos_token_id] + token_ids_0

    def save_vocabulary(self, save_directory, filename_prefix=None):
        out_name = (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
        out_path = os.path.join(save_directory, out_name)
        if os.path.abspath(self.vocab_file) != os.path.abspath(out_path):
            shutil.copyfile(self.vocab_file, out_path)
        return (out_path,)