from __future__ import annotations
import json
import os
from typing import Dict, List, Optional
from transformers import PreTrainedTokenizer
import torch


class ChessTokenizer(PreTrainedTokenizer):
    """
    符合评估脚本要求的 Chess Tokenizer。
    1. 词表大小为 144 (4 special + 12 pieces + 64 from_sq + 64 to_sq)。
    2. Decode 结果为紧凑格式（如 "WPe2e4"），确保 evaluate.py 的切片 [2:4] 和 [4:6] 正确。
    3. 区分起始格和目标格语义。
    """

    model_input_names = ["input_ids", "attention_mask"]
    vocab_files_names = {"vocab_file": "vocab.json"}

    PAD_TOKEN = "[PAD]"
    BOS_TOKEN = "[BOS]"
    EOS_TOKEN = "[EOS]"
    UNK_TOKEN = "[UNK]"

    def __init__(self, vocab_file: Optional[str] = None, vocab: Optional[Dict[str, int]] = None, **kwargs):
        special_tokens = [self.PAD_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN, self.UNK_TOKEN]
        
        # 必须使用大写，以匹配 evaluate.py 生成的棋谱
        self.colors_pieces = [f'{c}{p}' for c in ['W','B'] for p in ['P','N','B','R','Q','K']]  # 12个
        self.squares = [f'{f}{r}' for r in '12345678' for f in 'abcdefgh']  # 64个

        if vocab is not None:
            self._vocab = vocab
        elif vocab_file is not None and os.path.exists(vocab_file):
            with open(vocab_file, "r", encoding="utf-8") as f:
                self._vocab = json.load(f)
        else:
            # 构建 144 大小的词表
            self._vocab = {t: i for i, t in enumerate(special_tokens)} # 0-3
            
            # 4-15: Piece tokens
            for cp in self.colors_pieces:
                self._vocab[cp] = len(self._vocab)
            
            # 16-79: From Square tokens (内部带后缀防止重名)
            for sq in self.squares:
                self._vocab[f"{sq}_f"] = len(self._vocab)
            
            # 80-143: To Square tokens
            for sq in self.squares:
                self._vocab[f"{sq}_t"] = len(self._vocab)

        self._ids_to_tokens = {v: k for k, v in self._vocab.items()}

        super().__init__(
            pad_token=self.PAD_TOKEN,
            bos_token=self.BOS_TOKEN,
            eos_token=self.EOS_TOKEN,
            unk_token=self.UNK_TOKEN,
            **kwargs,
        )

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

    def get_vocab(self) -> Dict[str, int]:
        return dict(self._vocab)

    def _tokenize(self, text: str) -> List[str]:
        """将 WPe2e4 拆分为三个 token"""
        tokens = []
        # 处理可能的空格分隔（如历史棋谱）
        moves = text.strip().split()
        for move in moves:
            # 过滤特殊 token 字符串
            if move in [self.PAD_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN, self.UNK_TOKEN]:
                tokens.append(move)
                continue
            
            if len(move) >= 6:
                cp = move[:2]         # 例如 "WP"
                from_sq = move[2:4] + "_f" # 例如 "e2_f"
                to_sq = move[4:6] + "_t"   # 例如 "e4_t"
                tokens.extend([cp, from_sq, to_sq])
        return tokens

    def _convert_token_to_id(self, token: str) -> int:
        return self._vocab.get(token, self._vocab[self.UNK_TOKEN])

    def _convert_id_to_token(self, index: int) -> str:
        token = self._ids_to_tokens.get(index, self.UNK_TOKEN)
        # 如果是特殊 Token，返回空字符串，避免干扰 decode 结果
        if token in [self.PAD_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN, self.UNK_TOKEN]:
            return ""
        # 去掉内部后缀
        return token.replace("_f", "").replace("_t", "")

    def convert_tokens_to_string(self, tokens: List[str]) -> str:
        """
        核心修复：确保拼接结果符合 evaluate.py 的 6 位切片要求
        """
        # 1. 过滤掉 None 或空字符串
        clean_tokens = [t for t in tokens if t and t.strip()]
        
        # 2. 拼接原始字符
        raw_res = "".join(clean_tokens)
        
        # 3. 逻辑补全：
        # 老师的脚本期待的是 [Piece(2)][From(2)][To(2)]
        # 如果当前已经凑够了 3 个组件（比如 WP, e2, e4），raw_res 长度就是 6
        # 如果只凑了 2 个组件（比如 WP, e2），长度是 4
        
        # 特别注意：如果 tokens 只有 1 个且长度 >= 6（说明是一次性生成的全量 move）
        if len(raw_res) >= 6:
            # 这种情况下直接返回，满足 if len(token_str) >= 6: break
            return raw_res
            
        return raw_res

    def decode(self, token_ids, skip_special_tokens=True, **kwargs) -> str:
        """
        覆盖父类的 decode，增加对老师脚本的长度伪装
        """
        # 将输入统一转为 list，防止 Tensor 报错
        if hasattr(token_ids, "tolist"):
            ids = token_ids.tolist()
        elif isinstance(token_ids, (int, torch.LongTensor, torch.IntTensor)):
            ids = [int(token_ids)]
        else:
            ids = token_ids

        # 将 ID 转回 token
        tokens = [self._convert_id_to_token(i) for i in ids]
        
        # 调用你写好的拼接逻辑
        decoded_str = self.convert_tokens_to_string(tokens)
        
        return decoded_str

    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> tuple:
        if not os.path.isdir(save_directory):
            os.makedirs(save_directory, exist_ok=True)
        vocab_file = os.path.join(
            save_directory,
            (filename_prefix + "-" if filename_prefix else "") + "vocab.json"
        )
        with open(vocab_file, "w", encoding="utf-8") as f:
            json.dump(self._vocab, f, ensure_ascii=False, indent=2)
        return (vocab_file,)

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs) -> "ChessTokenizer":
        vocab_file = os.path.join(pretrained_model_name_or_path, "vocab.json")
        if not os.path.exists(vocab_file):
            return cls() # 如果没有文件则初始化默认的
        with open(vocab_file, "r", encoding="utf-8") as f:
            vocab = json.load(f)
        return cls(vocab=vocab, **kwargs)