"""
play.py - تشغيل نموذج XO المُدرّب للعب ضده.
الاستخدام: python play.py
يتطلب: xo_model/ , xo_tokenizer/ (يتم إنشاؤها بعد train.py)
"""

import torch
from model import load_model        # دالة تحميل النموذج
from tokenizer import load_tokenizer  # دالة تحميل الـ tokenizer


def display_board(board_str):
    """عرض اللوحة بشكل جميل."""
    b = list(board_str)
    print("\n")
    for i in range(0, 9, 3):
        print(" " + " | ".join(b[i:i+3]))
        if i < 6:
            print("---+---+---")
    print()


def check_win(board, player):
    """التحقق من فوز لاعب معين (X أو O)."""
    wins = [
        (0,1,2), (3,4,5), (6,7,8),
        (0,3,6), (1,4,7), (2,5,8),
        (0,4,8), (2,4,6)
    ]
    return any(all(board[i] == player for i in combo) for combo in wins)


def model_move(board_str, model, tokenizer):
    """
    إعطاء النموذج لوحة وإرجاع رقم الخانة (0-8) التي سيختارها.
    """
    inputs = tokenizer(board_str, return_tensors="pt")
    with torch.no_grad():
        logits = model(**inputs).logits.squeeze()

    # إخفاء الخانات الممتلئة
    for i, cell in enumerate(board_str):
        if cell != '.':
            logits[i] = -float('inf')

    return torch.argmax(logits).item()


def main():
    # تحميل النموذج والـ tokenizer
    model = load_model("./xo_model")
    tokenizer = load_tokenizer("./xo_tokenizer")
    model.eval()

    print("\n🎮 مرحباً بك في لعبة XO ضد نموذج Transformer!")
    print("أنت (O) | النموذج (X) | النموذج يبدأ\n")
    print("ترقيم الخانات:")
    print(" 0 | 1 | 2 ")
    print("---+---+---")
    print(" 3 | 4 | 5 ")
    print("---+---+---")
    print(" 6 | 7 | 8 \n")

    board = '.' * 9
    turn = 'X'  # النموذج يبدأ

    while True:
        display_board(board)

        if turn == 'X':
            print("🤖 النموذج يفكر...")
            move = model_move(board, model, tokenizer)
            print(f"النموذج اختار الخانة {move}")
        else:
            # دور اللاعب
            while True:
                try:
                    move = int(input("📍 دورك - أدخل رقم الخانة (0-8): "))
                    if move not in range(9):
                        print("❌ خارج النطاق! أدخل رقم من 0 إلى 8.")
                        continue
                    if board[move] != '.':
                        print("❌ الخانة مشغولة، اختر خانة فارغة.")
                        continue
                    break
                except ValueError:
                    print("❌ أدخل رقماً صحيحاً.")

        # تنفيذ الحركة
        board_list = list(board)
        board_list[move] = turn
        board = ''.join(board_list)

        # التحقق من نهاية اللعبة
        if check_win(board, turn):
            display_board(board)
            winner = "النموذج" if turn == 'X' else "أنت"
            print(f"🏆 {winner} فاز!")
            break

        if '.' not in board:
            display_board(board)
            print("🤝 تعادل!")
            break

        # تبديل الدور
        turn = 'O' if turn == 'X' else 'X'

    print("\n👋 شكراً للعبك!")


if __name__ == "__main__":
    main()