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

class ChessTokenizer(PreTrainedTokenizer):
    model_input_names = ["input_ids", "attention_mask"]
    PAD_TOKEN, BOS_TOKEN, EOS_TOKEN, UNK_TOKEN = "[PAD]", "[BOS]", "[EOS]", "[UNK]"
    
    def __init__(self, vocab_file=None, **kwargs):
        self._vocab = self._create_vocab()
        self._ids_to_tokens = {v: k for k, v in self._vocab.items()}
        
        # On définit les tokens spéciaux
        kwargs.setdefault("pad_token", self.PAD_TOKEN)
        kwargs.setdefault("bos_token", self.BOS_TOKEN) # ID 1
        kwargs.setdefault("eos_token", self.EOS_TOKEN)
        kwargs.setdefault("unk_token", self.UNK_TOKEN)
        super().__init__(**kwargs)

    def _create_vocab(self) -> Dict[str, int]:
        vocab = {self.PAD_TOKEN: 0, self.BOS_TOKEN: 1, self.EOS_TOKEN: 2, self.UNK_TOKEN: 3}
        idx = 4
        for f in "abcdefgh":
            for r in "12345678":
                vocab[f + r] = idx; idx += 1
        for p in ["q", "r", "b", "n"]: vocab[p] = idx; idx += 1
        return vocab
    
    @property
    def vocab_size(self) -> int: return len(self._vocab)
    def get_vocab(self) -> Dict[str, int]: return self._vocab
    
    def _tokenize(self, text: str) -> List[str]:
        text = text.replace(" ", "")
        import re
        moves = re.findall(r'[a-h][1-8][a-h][1-8][qrbn]?', text)
        tokens = []
        for move in moves:
            tokens.append(move[:2])
            tokens.append(move[2:4])
            if len(move) > 4: tokens.append(move[4])
        return tokens
        
    def _convert_token_to_id(self, token: str) -> int: return self._vocab.get(token, self._vocab.get(self.UNK_TOKEN))
    def _convert_id_to_token(self, index: int) -> str: return self._ids_to_tokens.get(index, self.UNK_TOKEN)
    
    def convert_tokens_to_string(self, tokens: List[str]) -> str:
        text = "".join(tokens)
        for special in [self.PAD_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN, self.UNK_TOKEN]:
            text = text.replace(special, "")
        return text.strip()

    # --- LA CORRECTION "STARTER" ---
    # Cette fonction force l'ajout du BOS token (ID 1) au début de tout input
    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.bos_token_id is not None else []
        eos_token_id = [self.eos_token_id] if self.eos_token_id is not None else []
        
        # On ajoute BOS au début !
        return bos_token_id + token_ids_0

    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> tuple:
        with open(os.path.join(save_directory, "vocab.json"), "w") as f: json.dump(self._vocab, f)
        return (os.path.join(save_directory, "vocab.json"),)
        
    @classmethod
    def build_vocab_from_dataset(cls, *args, **kwargs): return cls()