import json
import os
from pathlib import Path
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple, Union
import torch

import sentencepiece

from transformers.tokenization_utils import PreTrainedTokenizer
from transformers.utils import logging


logger = logging.get_logger(__name__)

SPIECE_UNDERLINE = "▁"

VOCAB_FILES_NAMES = {
    "vocab_file": "vocab.json",
    "spm_file": "mitre_spm.model",
    "tokenizer_config_file": "tokenizer_config.json",
}

# follow iso639-2
FAIRSEQ_LANGUAGE_CODES = ["en", "de", "nl", "sv", "da", "af", "fr", "es", "it", "pt", "ro", "ru", "cs", "pl", "bg", "uk", "id", "jv", "ms", "tl", "ja", "zh", "ko", "vi"]

# This is the tokenizer of MITRE.
# This code is modified from transformers.models.m2m_100.tokenization_m2m_100.M2M100Tokenizer
class MitreTokenizer(PreTrainedTokenizer):
    vocab_files_names = VOCAB_FILES_NAMES
    model_input_names = ["input_ids", "attention_mask"]

    prefix_tokens: List[int] = []
    suffix_tokens: List[int] = []

    def __init__(
        self,
        vocab_file,
        spm_file,
        bos_token="<s>",
        eos_token="</s>",
        sep_token="</s>",
        pad_token="<pad>",
        unk_token="<unk>",
        sp_model_kwargs: Optional[Dict[str, Any]] = None,
        **kwargs,
    ) -> None:
        self.sp_model_kwargs = {} if sp_model_kwargs is None else sp_model_kwargs
        fairseq_language_code = FAIRSEQ_LANGUAGE_CODES
        self.lang_code_to_token = {lang_code: f"__{lang_code}__" for lang_code in fairseq_language_code}

        additional_special_tokens = kwargs.pop("additional_special_tokens", [])
        for lang_code in fairseq_language_code:
            token = self.get_lang_token(lang_code)
            if token not in additional_special_tokens:
                additional_special_tokens.append(token)

        self.vocab_file = vocab_file
        self.encoder = load_json(vocab_file)
        self.decoder = {v: k for k, v in self.encoder.items()}
        self.spm_file = spm_file
        self.sp_model = load_spm(spm_file, self.sp_model_kwargs)

        self.encoder_size = len(self.encoder)

        self.lang_token_to_id = {
            self.get_lang_token(lang_code): self.encoder_size + i for i, lang_code in enumerate(fairseq_language_code)
        }
        self.lang_code_to_id = {lang_code: self.encoder_size + i for i, lang_code in enumerate(fairseq_language_code)}
        self.id_to_lang_token = {v: k for k, v in self.lang_token_to_id.items()}
        # default
        self.tgt_lang = "en"

        super().__init__(
            bos_token=bos_token,
            eos_token=eos_token,
            sep_token=sep_token,
            unk_token=unk_token,
            pad_token=pad_token,
            sp_model_kwargs=self.sp_model_kwargs,
            additional_special_tokens=additional_special_tokens,
            **kwargs,
        )

    @property
    def vocab_size(self) -> int:
        return len(self.encoder)

    def get_vocab(self) -> Dict:
        vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size)}
        vocab.update(self.added_tokens_encoder)
        return vocab

    def _tokenize(self, text: str) -> List[str]:
        return self.sp_model.encode(text, out_type=str)

    def _convert_token_to_id(self, token):
        if token in self.lang_token_to_id:
            return self.lang_token_to_id[token]
        return self.encoder.get(token, self.encoder[self.unk_token])

    def _convert_id_to_token(self, index: int) -> str:
        """Converts an index (integer) in a token (str) using the decoder."""
        if index in self.id_to_lang_token:
            return self.id_to_lang_token[index]
        return self.decoder.get(index, self.unk_token)

    def convert_tokens_to_string(self, tokens):
        """Converts a sequence of tokens (string) in a single string."""
        current_sub_tokens = []
        out_string = ""
        for token in tokens:
            # make sure that special tokens are not decoded using sentencepiece model
            if token in self.all_special_tokens:
                out_string += self.sp_model.decode(current_sub_tokens) + token
                current_sub_tokens = []
            else:
                current_sub_tokens.append(token)
        out_string += self.sp_model.decode(current_sub_tokens)
        return out_string.strip()

    def __getstate__(self) -> Dict:
        state = self.__dict__.copy()
        state["sp_model"] = None
        return state

    def __setstate__(self, d: Dict) -> None:
        self.__dict__ = d

        # for backward compatibility
        if not hasattr(self, "sp_model_kwargs"):
            self.sp_model_kwargs = {}

        self.sp_model = load_spm(self.spm_file, self.sp_model_kwargs)
    
    def build_inputs_with_special_tokens(
        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
    ) -> List[int]:
        if token_ids_1 is None:
            return self.prefix_tokens + token_ids_0 + self.suffix_tokens
        return self.prefix_tokens + token_ids_0 + token_ids_1 + self.suffix_tokens
    
    def _switch_to_input_mode(self):
        self.set_tgt_lang_special_tokens(self.tgt_lang)

    def _switch_to_target_mode(self):
        self.clear_lang_special_tokens()

    def clear_lang_special_tokens(self) -> None:
        self.prefix_tokens = []
        self.suffix_tokens = [self.eos_token_id]

    def set_tgt_lang_special_tokens(self, tgt_lang: str) -> None:
        """Reset the special tokens to the target language setting. No prefix and suffix=[eos, tgt_lang_code]."""
        lang_token = self.get_lang_token(tgt_lang)
        self.cur_lang_id = self.lang_token_to_id[lang_token]
        self.prefix_tokens = [self.cur_lang_id]
        self.suffix_tokens = [self.eos_token_id]

    def get_lang_token(self, lang: str) -> str:
        return self.lang_code_to_token[lang]

    def get_lang_id(self, lang: str) -> int:
        lang_token = self.get_lang_token(lang)
        return self.lang_token_to_id[lang_token]
    
    def encode_source_tokens_to_input_ids(self, inputs, target_language="en"):
        """pads + target language id + source tokens id + eos id"""
        self.tgt_lang = target_language
        input_ids = self.__call__(inputs, add_special_tokens=True, padding_side='left', padding=True, return_attention_mask=False, return_tensors="pt")
        return input_ids["input_ids"]
    
    def encode_source_tokens_to_input_ids_with_different_tags(self, inputs_text, target_languages_list: list):
        """
          'encode_source_tokens_to_input_ids' only supports a language tag,
          but sevenral in a batch could have different language tags. 
        """
        self.tgt_lang = "en"
        input_ids = self.__call__(inputs_text, add_special_tokens=True, padding_side='left', padding=True, return_attention_mask=False, return_tensors="pt")["input_ids"]
        _, max_indices = torch.max(input_ids, dim=1)
        input_ids[torch.arange(max_indices.shape[0]), max_indices] = torch.LongTensor([self.lang_token_to_id[self.get_lang_token(lang_code)] for lang_code in target_languages_list])
        return input_ids
    
    def encode_target_tokens_to_labels(self, inputs_text):
        """target tokens id + eos id + pads"""
        input_ids = self.__call__(text_target=inputs_text, add_special_tokens=True, padding_side='right', padding=True, return_attention_mask=False, return_tensors="pt")
        return input_ids["input_ids"]
    
    def encode_target_tokens_to_input_ids(self, inputs_text):
        """eos id + target tokens id + pads, namely, left shifted"""
        input_ids = self.__call__(text_target=inputs_text, add_special_tokens=False, padding_side='right', padding=True, return_attention_mask=False, return_tensors="pt")
        labels_without_eos = input_ids["input_ids"]
        return torch.cat((torch.full((labels_without_eos.size(0), 1), self.eos_token_id), labels_without_eos), dim=1)
    

def load_spm(path: str, sp_model_kwargs: Dict[str, Any]) -> sentencepiece.SentencePieceProcessor:
    spm = sentencepiece.SentencePieceProcessor(**sp_model_kwargs)
    spm.Load(str(path))
    return spm


def load_json(path: str) -> Union[Dict, List]:
    with open(path, "r") as f:
        return json.load(f)


def save_json(data, path: str) -> None:
    with open(path, "w") as f:
        json.dump(data, f, indent=2)

MitreTokenizer.register_for_auto_class("AutoTokenizer")