"""
tokenizer.py - بناء وتحميل Tokenizer مخصص للوحة XO.
"""

from tokenizers import Tokenizer
from tokenizers.models import WordLevel
from tokenizers import pre_tokenizers, decoders
from transformers import PreTrainedTokenizerFast


def build_tokenizer(save_path="./xo_tokenizer"):
    """
    ينشئ Tokenizer على مستوى المحارف (character-level) مع الرموز الخاصة:
    <pad>, <eos>, <unk>.
    ثم يحفظه ويعيده.
    """
    special_tokens = ["<pad>", "<eos>", "<unk>"]
    # جميع الرموز الممكنة في اللوحة + أرقام (للاستخدام المستقبلي)
    vocab_chars = ['.', 'X', 'O'] + [str(i) for i in range(10)]

    # بناء القاموس (كل رمز -> معرف)
    vocab = {tok: i for i, tok in enumerate(special_tokens + vocab_chars)}

    # إنشاء tokenizer أساسي
    base_tokenizer = Tokenizer(WordLevel(vocab=vocab, unk_token="<unk>"))
    base_tokenizer.pre_tokenizer = pre_tokenizers.WhitespaceSplit()
    base_tokenizer.decoder = decoders.ByteLevel()

    # تغليفه ليتوافق مع مكتبة transformers
    tokenizer = PreTrainedTokenizerFast(
        tokenizer_object=base_tokenizer,
        bos_token=None,
        eos_token="<eos>",
        unk_token="<unk>",
        pad_token="<pad>",
    )

    tokenizer.save_pretrained(save_path)
    print(f"✅ تم بناء Tokenizer وحفظه في {save_path}")
    return tokenizer


def load_tokenizer(tokenizer_path="./xo_tokenizer"):
    """
    يحمّل Tokenizer محفوظاً مسبقاً.
    """
    tokenizer = PreTrainedTokenizerFast.from_pretrained(tokenizer_path)
    print(f"📂 تم تحميل Tokenizer من {tokenizer_path}")
    return tokenizer