import io
import json
import os
from shutil import copyfile
from typing import Any, Dict, IO, List, Optional, Tuple

import pkg_resources
import sentencepiece as spm
from pytrie import StringTrie
from transformers.tokenization_utils import AddedToken, PreTrainedTokenizer
from transformers.utils import logging

logger = logging.get_logger(__name__)

VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt"}

PRETRAINED_VOCAB_FILES_MAP = {
    "vocab_file": {},
    "tokenizer_file": {},
}
PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES = {}


class CPM9GTokenizer(PreTrainedTokenizer):
    """
    CPM9G 分词器类。用于基于字节对编码的分词。

    参数:
        path (str, 可选): 词汇表文件的路径。
    """

    vocab_files_names = VOCAB_FILES_NAMES
    pretrained_vocab_files_map = PRETRAINED_VOCAB_FILES_MAP
    max_model_input_sizes = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
    model_input_names = ["input_ids", "attention_mask"]

    def __init__(
        self,
        vocab_file: Optional[str] = None,
        unk_token: str = "<unk>",
        bos_token: str = "<s>",
        eos_token: str = "</s>",
        pad_token: Optional[str] = None,
        sp_model_kwargs: Optional[Dict[str, Any]] = None,
        add_bos_token: bool = True,
        add_eos_token: bool = False,
        clean_up_tokenization_spaces: bool = False,
        **kwargs,
    ):
        self.sp_model_kwargs = sp_model_kwargs or {}
        self.vocab_file = vocab_file
        self.add_bos_token = add_bos_token
        self.add_eos_token = add_eos_token

        self.unk_token = unk_token
        self.bos_token = bos_token
        self.eos_token = eos_token
        self.pad_token = pad_token

        self.byte_list: List[str] = (
            [f"<0x0{hex(i).upper()[2:]}>" for i in range(0x10)] +
            [f"<0x{hex(i).upper()[2:]}>" for i in range(0x10, 0x100)]
        )

        self._special_token_set = set([self.unk_token, self.bos_token, self.eos_token] + self.byte_list)

        if vocab_file:
            if 'vocab.txt' not in vocab_file:
                all_tokens = self.load_vocab(io.FileIO(vocab_file + VOCAB_FILES_NAMES['vocab_file'], "rb"))
            else:
                all_tokens = self.load_vocab(io.FileIO(vocab_file, "rb"))
        else:
            all_tokens = self.load_vocab(io.FileIO(VOCAB_FILES_NAMES['vocab_file'], "rb"))

        self.encoder: Dict[str, int] = {}
        self._special_encoder: Dict[str, int] = {}
        for token, token_id in all_tokens.items():
            if token in self._special_token_set:
                self._special_encoder[token] = token_id
            else:
                self.encoder[token] = token_id

        self.decoder = {v: k for k, v in self.encoder.items()}
        self._byte_decoder = {self._special_encoder[token]: i for i, token in enumerate(self.byte_list)}

        self._max_word_len = max([len(x) for x in self.encoder.keys()])

        self._len_word_first = {}
        for x in self.encoder.keys():
            if not x[0] in self._len_word_first:
                self._len_word_first[x[0]] = 1
            if len(x) > self._len_word_first[x[0]]:
                self._len_word_first[x[0]] = len(x)
        self.tencoder = StringTrie(self.encoder)

        super().__init__(
            bos_token=AddedToken(bos_token, lstrip=False, rstrip=False),
            eos_token=AddedToken(eos_token, lstrip=False, rstrip=False),
            unk_token=AddedToken(unk_token, lstrip=False, rstrip=False),
            pad_token=AddedToken(pad_token, lstrip=False, rstrip=False) if pad_token else None,
            add_bos_token=add_bos_token,
            add_eos_token=add_eos_token,
            sp_model_kwargs=self.sp_model_kwargs,
            clean_up_tokenization_spaces=clean_up_tokenization_spaces,
            **kwargs,
        )

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

    def __setstate__(self, d: Dict[str, Any]) -> None:
        self.__dict__ = d

    def load_vocab(self, fp: IO[bytes]) -> Dict[str, int]:
        """
        加载词汇表文件到字典中。

        参数:
            fp (IO[bytes]): 词汇表文件指针。

        返回:
            Dict[str, int]: 词汇表字典。
        """
        vocab: Dict[str, int] = {}
        reader = io.TextIOWrapper(fp, encoding="utf-8")
        for token in reader.readlines():
            token = token.strip()
            if len(token) == 0:
                continue
            token = json.loads(token)
            vocab[token] = len(vocab)
        return vocab

    @property
    def vocab_size(self) -> int:
        """返回词汇表大小"""
        return len(self.encoder) + len(self._special_encoder)

    @property
    def eos_id(self):
        return self._special_encoder[self.eos_token]

    @property
    def bos_id(self):
        return self._special_encoder[self.bos_token]

    @property
    def unk_id(self):
        return self._special_encoder[self.unk_token]

    def get_vocab(self) -> Dict[str, int]:
        """返回词汇表作为字典"""
        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]:
        """返回分词后的字符串"""
        output_tokens: List[str] = []
        st = 0
        while st < len(text):
            piece = self.get_piece(text[st:])
            output_tokens.append(piece)
            st += len(piece)
        return output_tokens

    def _convert_token_to_id(self, token: str) -> int:
        """使用词汇表将标记（字符串）转换为 id"""
        return self.encoder.get(token, self.unk_id)

    def _convert_id_to_token(self, index: int) -> str:
        """使用词汇表将索引（整数）转换为标记（字符串）"""
        return self.decoder.get(index, self.unk_token)

    def convert_tokens_to_string(self, tokens: List[str]) -> str:
        """将标记序列（字符串）转换为单个字符串"""
        current_sub_tokens: List[str] = []
        out_string = ""
        prev_is_special = False
        for i, token in enumerate(tokens):
            if token in self._special_token_set:
                if not prev_is_special and i != 0:
                    out_string += " "
                out_string += self.decode(current_sub_tokens) + token
                prev_is_special = True
                current_sub_tokens = []
            else:
                current_sub_tokens.append(token)
                prev_is_special = False
        out_string += self.sp_model.decode(current_sub_tokens)
        return out_string

    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
        """
        保存词汇表和特殊标记文件到目录。

        参数:
            save_directory (str): 要保存词汇表的目录。

        返回:
            Tuple[str]: 保存的文件路径。
        """
        if not os.path.isdir(save_directory):
            raise ValueError(f"Vocabulary path ({save_directory}) should be a directory")

        out_vocab_file = os.path.join(
            save_directory,
            (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"],
        )

        if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file) and os.path.isfile(self.vocab_file):
            copyfile(self.vocab_file, out_vocab_file)
        elif not os.path.isfile(self.vocab_file):
            with open(out_vocab_file, "wb") as fi:
                fi.write(self.sp_model.serialized_model_proto())

        return (out_vocab_file, )

    def build_inputs_with_special_tokens(
        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
    ) -> List[int]:
        bos_token_id = [self.bos_token_id] if self.add_bos_token else []
        eos_token_id = [self.eos_token_id] if self.add_eos_token else []

        output = bos_token_id + token_ids_0 + eos_token_id

        if token_ids_1 is not None:
            output = output + bos_token_id + token_ids_1 + eos_token_id

        return output

    def get_special_tokens_mask(
        self, 
        token_ids_0: List[int], 
        token_ids_1: Optional[List[int]] = None, 
        already_has_special_tokens: bool = False
    ) -> List[int]:
        """
        获取从未添加特殊标记的标记列表中检索到的序列 id。
        在使用分词器的 `prepare_for_model` 方法添加特殊标记时调用此方法。

        参数:
            token_ids_0 (List[int]): id 列表。
            token_ids_1 (List[int], 可选): 序列对的可选第二 id 列表。
            already_has_special_tokens (bool, 可选, 默认值为 False): 
                标记列表是否已使用模型的特殊标记进行格式化。

        返回:
            List[int]: 一个包含整数（0 或 1）的列表。1 表示特殊标记，0 表示序列标记。
        """
        if already_has_special_tokens:
            return super().get_special_tokens_mask(
                token_ids_0=token_ids_0,
                token_ids_1=token_ids_1,
                already_has_special_tokens=True,
            )

        bos_token_id = [1] if self.add_bos_token else []
        eos_token_id = [1] if self.add_eos_token else []

        if token_ids_1 is None:
            return bos_token_id + ([0] * len(token_ids_0)) + eos_token_id
        return bos_token_id + ([0] * len(token_ids_0)) + eos_token_id + bos_token_id + ([0] * len(token_ids_1)) + eos_token_id

    def create_token_type_ids_from_sequences(
        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
    ) -> List[int]:
        """
        从传递的两个序列创建掩码，用于序列对分类任务。

        参数:
            token_ids_0 (List[int]): id 列表。
            token_ids_1 (List[int], 可选): 序列对的可选第二 id 列表。

        返回:
            List[int]: 根据给定序列的标记类型 id 列表。
        """
        bos_token_id = [self.bos_token_id] if self.add_bos_token else []
        eos_token_id = [self.eos_token_id] if self.add_eos_token else []

        output = [0] * len(bos_token_id + token_ids_0 + eos_token_id)

        if token_ids_1 is not None:
            output += [1] * len(bos_token_id + token_ids_1 + eos_token_id)

        return output

    def get_piece(self, text: str) -> str:
        """
        获取文本中的分词片段。

        参数:
            text (str): 输入文本。

        返回:
            str: 分词片段。
        """
        if text[0] in self._len_word_first:
            text = text[: self._len_word_first[text[0]]]
            len_text = len(text)
            for i in range(len(text)):
                sub = text[: len_text - i]
                if sub in self.encoder:
                    return sub
        return text[0]


    def encode(self, text: str, add_special_tokens: bool) -> List[int]:
        """
        将文本编码为 ID 列表。

        参数:
            text (str): 输入文本。

        返回:
            List[int]: 编码后的 ID 列表。
        """
        #if len(text) > 20480:
        #    return [0 for _ in range(20480)]
        ret = []
        for x in self._tokenize(text):
            if x in self.encoder:
                ret.append(self.encoder[x])
            else:
                ret.extend(self._encode_unicode(x))
        return ret


    def decode_all(self, tokens: List[int]):
        """Decode ids into a string."""
        ret = []
        st = 0

        while st < len(tokens):
            if tokens[st] in self.decoder:
                ret.append(self.decoder[tokens[st]])
                st += 1
            elif tokens[st] in self._byte_decoder:
                if (
                    st + 3 < len(tokens)
                    and tokens[st + 1] in self._byte_decoder
                    and tokens[st + 2] in self._byte_decoder
                    and tokens[st + 3] in self._byte_decoder
                ):
                    first_id = self._byte_decoder[tokens[st]]
                    plane_id = self._byte_decoder[tokens[st + 1]]
                    row_id = self._byte_decoder[tokens[st + 2]]
                    cell_id = self._byte_decoder[tokens[st + 3]]
                    ret.append(
                        int.to_bytes(first_id << 24 | plane_id << 16 | row_id << 8 | cell_id, 4, "big").decode("utf-8")
                    )
                    st += 4
                elif (
                    st + 2 < len(tokens)
                    and tokens[st + 1] in self._byte_decoder
                    and tokens[st + 2] in self._byte_decoder
                ):
                    plane_id = self._byte_decoder[tokens[st]]
                    row_id = self._byte_decoder[tokens[st + 1]]
                    cell_id = self._byte_decoder[tokens[st + 2]]
                    ret.append(int.to_bytes(plane_id << 16 | row_id << 8 | cell_id, 3, "big").decode("utf-8"))
                    st += 3
                elif st + 1 < len(tokens) and tokens[st + 1] in self._byte_decoder:
                    row_id = self._byte_decoder[tokens[st]]
                    cell_id = self._byte_decoder[tokens[st + 1]]
                    ret.append(int.to_bytes(row_id << 8 | cell_id, 2, "big").decode("utf-8"))
                    st += 2
                else:
                    cell_id = self._byte_decoder[tokens[st]]
                    ret.append(int.to_bytes(cell_id, 1, "big").decode("utf-8"))
                    st += 1
            elif tokens[st] == self.eos_id:
                ret.append(self.eos_token)
                st += 1
            elif tokens[st] == self.bos_id:
                ret.append(self.bos_token)
                st += 1
            else:
                ret.append(self.unk_token)
                st += 1
        return "".join(ret)
    
    def decode(self, tokens: List[int], skip_special_tokens: bool) -> str:
        """
        将 ID 列表解码为字符串。

        参数:
            tokens (List[int]): ID 列表。

        返回:
            str: 解码后的字符串。
        """
        ret = []
        st = 0

        while st < len(tokens):
            if tokens[st] in self._byte_decoder:
                if (
                    st + 3 < len(tokens)
                    and tokens[st + 1] in self._byte_decoder
                    and tokens[st + 2] in self._byte_decoder
                    and tokens[st + 3] in self._byte_decoder
                ):
                    first_id = self._byte_decoder[tokens[st]]
                    plane_id = self._byte_decoder[tokens[st + 1]]
                    row_id = self._byte_decoder[tokens[st + 2]]
                    cell_id = self._byte_decoder[tokens[st + 3]]
                    ret.append(
                        int.to_bytes(first_id << 24 | plane_id << 16 | row_id << 8 | cell_id, 4, "big").decode("utf-8")
                    )
                    st += 4
                elif (
                    st + 2 < len(tokens)
                    and tokens[st + 1] in self._byte_decoder
                    and tokens[st + 2] in self._byte_decoder
                ):
                    plane_id = self._byte_decoder[tokens[st]]
                    row_id = self._byte_decoder[tokens[st + 1]]
                    cell_id = self._byte_decoder[tokens[st + 2]]
                    ret.append(int.to_bytes(plane_id << 16 | row_id << 8 | cell_id, 3, "big").decode("utf-8"))
                    st += 3
                elif st + 1 < len(tokens) and tokens[st + 1] in self._byte_decoder:
                    row_id = self._byte_decoder[tokens[st]]
                    cell_id = self._byte_decoder[tokens[st + 1]]
                    ret.append(int.to_bytes(row_id << 8 | cell_id, 2, "big").decode("utf-8"))
                    st += 2
                else:
                    cell_id = self._byte_decoder[tokens[st]]
                    ret.append(int.to_bytes(cell_id, 1, "big").decode("utf-8"))
                    st += 1
            elif tokens[st] == self.eos_id:
                ret.append(self.eos_token)
                st += 1
            elif tokens[st] == self.bos_id:
                ret.append(self.bos_token)
                st += 1
            else:
                ret.append(tokens[st])
                st += 1
            #else:
            #    ret.append(self.unk_token)
            #    st += 1
        return ''.join([str(item) for item in ret])

    def _encode_unicode(self, token: str) -> List[int]:
        """
        将 Unicode 编码包装到一个辅助函数中。

        参数:
            token (str): 要编码的标记。

        返回:
            List[int]: 编码后的 ID 列表。
        """
        ids = []
        utf8_id = token.encode("utf-8")
        for _id in utf8_id:
            ids.append(self._special_encoder[self.byte_list[_id]])
        return ids

    def next_token(self, text: str) -> Tuple[str, List[int]]:
        """
        快速获取下一个匹配的标记。

        参数:
            text (str): 输入文本。

        返回:
            Tuple[str, List[int]]: 匹配的标记及其 ID 列表。
        """
        token, token_id = self.tencoder.longest_prefix_item(text, (None, None))
        if token is None:
            token = text[0]
            token_ids = self._encode_unicode(token)
        else:
            token_ids = [token_id]
        return token, token_ids