import base64
import unicodedata
from typing import AbstractSet, Collection, Dict, List, Literal, Union

from .tokenizer_base import Tokenizer

PAT_STR = r"""(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"""  # noqa E501
ENDOFTEXT = "<|endoftext|>"
IMSTART = "<|im_start|>"
IMEND = "<|im_end|>"
# as the default behavior is changed to allow special tokens in
# regular texts, the surface forms of special tokens need to be
# as different as possible to minimize the impact
EXTRAS = tuple((f"<|extra_{i}|>" for i in range(205)))
# changed to use actual index to avoid misconfiguration with vocabulary expansion
SPECIAL_START_ID = 151643
SPECIAL_TOKENS = tuple(
    enumerate(
        (
            (
                ENDOFTEXT,
                IMSTART,
                IMEND,
            )
            + EXTRAS
        ),
        start=SPECIAL_START_ID,
    )
)
SPECIAL_TOKENS_SET = set(t for i, t in SPECIAL_TOKENS)


class QwenTokenizer(Tokenizer):
    @staticmethod
    def _load_tiktoken_bpe(tiktoken_bpe_file: str) -> Dict[bytes, int]:
        with open(tiktoken_bpe_file, "rb") as f:
            contents = f.read()
        return {
            base64.b64decode(token): int(rank)
            for token, rank in (line.split() for line in contents.splitlines() if line)
        }

    def __init__(self, vocab_file, errors="replace", extra_vocab_file=None):
        self._errors = errors
        self._vocab_file = vocab_file
        self._extra_vocab_file = extra_vocab_file

        self._mergeable_ranks = QwenTokenizer._load_tiktoken_bpe(vocab_file)  # type: Dict[bytes, int]
        self._special_tokens = {token: index for index, token in SPECIAL_TOKENS}

        # try load extra vocab from file
        if extra_vocab_file is not None:
            used_ids = set(self._mergeable_ranks.values()) | set(self._special_tokens.values())
            extra_mergeable_ranks = self._load_tiktoken_bpe(extra_vocab_file)
            for token, index in extra_mergeable_ranks.items():
                if token in self._mergeable_ranks:
                    continue
                if index in used_ids:
                    continue
                self._mergeable_ranks[token] = index
            # the index may be sparse after this, but don't worry tiktoken.Encoding will handle this
        import tiktoken

        enc = tiktoken.Encoding(
            "Qwen",
            pat_str=PAT_STR,
            mergeable_ranks=self._mergeable_ranks,
            special_tokens=self._special_tokens,
        )
        assert (
            len(self._mergeable_ranks) + len(self._special_tokens) == enc.n_vocab
        ), f"{len(self._mergeable_ranks) + len(self._special_tokens)} != {enc.n_vocab} in encoding"

        self.decoder = {v: k for k, v in self._mergeable_ranks.items()}  # type: dict[int, bytes|str]
        self.decoder.update({v: k for k, v in self._special_tokens.items()})

        self._tokenizer = enc  # type: tiktoken.Encoding

        self.eod_id = self._tokenizer.eot_token
        self.im_start_id = self._special_tokens[IMSTART]
        self.im_end_id = self._special_tokens[IMEND]

    def encode(
        self,
        text: str,
        allowed_special: Union[AbstractSet[str], Literal["all"]] = "all",
        disallowed_special: Union[Collection[str], str] = (),
    ) -> List[int]:
        text = unicodedata.normalize("NFC", text)
        return self._tokenizer.encode(text, allowed_special=allowed_special, disallowed_special=disallowed_special)

    def decode(
        self,
        token_ids: Union[int, List[int]],
        skip_special_tokens: bool = False,
        errors: str | None = None,
        **kwargs,
    ) -> str:
        if isinstance(token_ids, int):
            token_ids = [token_ids]
        if skip_special_tokens:
            token_ids = [i for i in token_ids if i < self.eod_id]
        return self._tokenizer.decode(token_ids, errors=errors or self._errors)
