# Adapted for diffusers from multimodal-art-projection/YuE at commit ef1936f2ee39fe8de486a0f47a481c95f8d4da87.
# Licensed under Apache-2.0; see LICENSE.
import base64
import shutil
import unicodedata
from pathlib import Path

import tiktoken
from transformers import PreTrainedTokenizer


class YuE2Tokenizer(PreTrainedTokenizer):
    vocab_files_names = {"vocab_file": "qwen.tiktoken"}
    model_input_names = ["input_ids"]

    def __init__(self, vocab_file, expected_vocab_size=151643, **kwargs):
        self.vocab_file = str(vocab_file)
        self.ranks = {
            base64.b64decode(token): int(rank)
            for token, rank in (line.split() for line in Path(vocab_file).read_bytes().splitlines() if line)
        }
        if len(self.ranks) != expected_vocab_size:
            raise ValueError(f"Expected {expected_vocab_size} ordinary tokens")
        specials = ["<|endoftext|>", "<|im_start|>", "<|im_end|>", "<R>", "<S>", "<X>", "<mask>", "<sep>"]
        specials += [f"<extra_{i}>" for i in range(200)]
        specials[204:206] = ["<abc>", "</abc>"]
        pattern = 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+"
        self.encoding = tiktoken.Encoding(
            "YuE2",
            pat_str=pattern,
            mergeable_ranks=self.ranks,
            special_tokens={s: i + len(self.ranks) for i, s in enumerate(specials)},
        )
        super().__init__(expected_vocab_size=expected_vocab_size, **kwargs)

    @property
    def vocab_size(self):
        return len(self.ranks)

    def get_vocab(self):
        return {base64.b64encode(token).decode(): rank for token, rank in self.ranks.items()}

    def _tokenize(self, text):
        return [base64.b64encode(self.encoding.decode_single_token_bytes(i)).decode() for i in self.encode(text)]

    def _convert_token_to_id(self, token):
        return self.ranks[base64.b64decode(token)]

    def _convert_id_to_token(self, index):
        return base64.b64encode(self.encoding.decode_single_token_bytes(index)).decode()

    def encode(self, text, **kwargs):
        return self.encoding.encode_ordinary(unicodedata.normalize("NFC", text))

    def decode(self, token_ids, **kwargs):
        return self.encoding.decode([int(i) for i in token_ids if 0 <= i < self.encoding.n_vocab], errors="replace")

    def save_vocabulary(self, save_directory, filename_prefix=None):
        target = Path(save_directory) / ((filename_prefix + "-" if filename_prefix else "") + "qwen.tiktoken")
        if target.resolve() != Path(self.vocab_file).resolve():
            shutil.copyfile(self.vocab_file, target)
        return (str(target),)