"""
Wrapper around piece_tokenizer that provides a HuggingFace-like interface.
Used by eval.py and finetune_muon.py.
"""
import os
import json
import piece_tokenizer as pt


# 词表和中文分词词典的文件名,**按优先级排列**。
#
# 首选带出处的名字(与 ../Summer/ 的产出同名)—— 这样从任何一个 checkpoint
# 目录就能看出词表出自哪里。sha256 核过:Summer-Tokenizer.pt 与旧名
# piece.model 逐字节相同,dict.txt 同理,所以改名不改内容。
#
# **旧名必须保留为回退**:已发布的 Ismantic/Qwen3-1.7B-Base-ReTok 和
# Ismantic/Interpreter-Qwen3-1.7B-ReTok 里都是旧名,现存 checkpoint 也是。
# 一刀切改名会让它们全部加载失败,而失败方式是「找不到就回退到别处的分词器」
# —— 正是这个仓库一路在修的那类静默错误。piece_mt.model 是更早的 HYMT 遗留。
VOCAB_NAMES = ["Summer-Tokenizer.pt", "piece.model", "piece_mt.model"]
DICT_NAMES = ["Summer-Tokenizer.dict.txt", "dict.txt"]


def _resolve(model_dir, names, what):
    for n in names:
        p = os.path.join(model_dir, n)
        if os.path.exists(p):
            return p
    raise FileNotFoundError(
        f"{model_dir} 里找不到{what},试过 {names}\n"
        f"checkpoint 目录必须自包含 —— 检查 src/artifacts.py 的五件套"
        f"是否都拷过来了。\n"
        f"注意中文词典**不是可选项**:少了它中文 token id 会变,而 round-trip "
        f"照样正确,从输出完全看不出来。")


class PieceTokenizerWrapper:
    def __init__(self, model_dir):
        """Load tokenizer from a model directory containing piece.model and token_mapping.json."""
        self._tok = pt.Tokenizer()

        model_file = _resolve(model_dir, VOCAB_NAMES, "piece 词表")

        # 中文分词词典是**必需的,不是可选的**。
        #
        # 这里原来的注释写的是「可选,少了只是 O(n^2) 慢」——实测是错的:
        # 少了它中文的 token id 会**变**。同一句话
        #
        #   带词典  … 另一(5066)  种(1066)  语言(5028) …
        #   无词典  … 另(3406)  一种(2612)  语言(5028) …
        #
        # 长度一样、**两边 round-trip 解码都正确**,所以从输出完全看不出来 ——
        # 模型只是悄悄吃到了与训练时不同的 id。这正是「改错了不报错」那一类。
        # Summer 那边踩过同一个坑,它的 prepare/tokenizer.py 也是直接报错。
        cn_dict = _resolve(model_dir, DICT_NAMES, "中文分词词典")
        self._tok.load(model_file, cn_dict)

        # Load token mapping
        mapping_file = os.path.join(model_dir, "token_mapping.json")
        if os.path.exists(mapping_file):
            with open(mapping_file) as f:
                mapping = json.load(f)
            self.pad_token_id = mapping["pad_id"]
            self.bos_token_id = mapping["bos_id"]
            self.eos_token_id = mapping["eos_id"]
            self.user_token_id = mapping.get("user_id")
            self.assistant_token_id = mapping.get("assistant_id")
            self.system_token_id = mapping.get("system_id")
        else:
            # Fallback to piece_to_id lookups
            self.bos_token_id = self._tok.piece_to_id("<s>")
            self.eos_token_id = self._tok.piece_to_id("</s>")
            self.pad_token_id = self._tok.piece_to_id("<pad>")
            self.user_token_id = self._tok.piece_to_id("<user>")
            self.assistant_token_id = self._tok.piece_to_id("<assistant>")
            self.system_token_id = self._tok.piece_to_id("<system>")
            if self.pad_token_id < 0:
                self.pad_token_id = 0

    @property
    def vocab_size(self):
        return self._tok.vocab_size()

    def encode(self, text, add_special_tokens=False):
        ids = self._tok.encode_as_ids(text)
        if add_special_tokens:
            ids = [self.bos_token_id] + ids + [self.eos_token_id]
        return ids

    def decode(self, ids, skip_special_tokens=True):
        if skip_special_tokens:
            special = {self.bos_token_id, self.eos_token_id, self.pad_token_id,
                       self.user_token_id, self.assistant_token_id, self.system_token_id}
            ids = [i for i in ids if i not in special]
        try:
            return self._tok.decode(ids)
        except UnicodeDecodeError:
            # Model emitted byte-fallback piece(s) that don't form valid UTF-8.
            # Per-piece fallback: keep ids that decode cleanly, drop the rest.
            parts = []
            for i in ids:
                try:
                    parts.append(self._tok.id_to_piece(i))
                except UnicodeDecodeError:
                    continue
            return "".join(parts).replace("▁", " ")

    def apply_chat_template(self, messages, tokenize=True, add_generation_prompt=False, **kwargs):
        """Build chat-formatted token sequence from messages."""
        ids = []

        # Check for system message
        start = 0
        if messages and messages[0]["role"] == "system":
            ids.append(self.bos_token_id)
            ids.extend(self._tok.encode_as_ids(messages[0]["content"]))
            ids.append(self.system_token_id)
            start = 1
        else:
            ids.append(self.bos_token_id)

        for msg in messages[start:]:
            if msg["role"] == "user":
                ids.append(self.user_token_id)
                ids.extend(self._tok.encode_as_ids(msg["content"]))
            elif msg["role"] == "assistant":
                ids.append(self.assistant_token_id)
                ids.extend(self._tok.encode_as_ids(msg["content"]))
                ids.append(self.eos_token_id)

        if add_generation_prompt:
            ids.append(self.assistant_token_id)

        if tokenize:
            return ids
        else:
            # Return as string (rarely needed)
            return self._tok.decode(ids)

    def save_pretrained(self, output_dir):
        """Save tokenizer files to directory (for checkpoint saving)."""
        import shutil
        os.makedirs(output_dir, exist_ok=True)
        # Copy piece.model
        src = os.path.join(os.path.dirname(output_dir), "piece.model")
        if os.path.exists(src):
            shutil.copy2(src, os.path.join(output_dir, "piece.model"))
        # Save mapping
        mapping = {
            "bos_id": self.bos_token_id,
            "eos_id": self.eos_token_id,
            "pad_id": self.pad_token_id,
            "user_id": self.user_token_id,
            "assistant_id": self.assistant_token_id,
            "system_id": self.system_token_id,
        }
        with open(os.path.join(output_dir, "token_mapping.json"), "w") as f:
            json.dump(mapping, f, indent=2)