"""
train.py - البرنامج الرئيسي لتدريب نموذج XO.
الاستخدام: python train.py
يعتمد على: tokenizer.py و model.py
"""

import os
import json
import random
import numpy as np
from datasets import load_dataset
from transformers import DataCollatorWithPadding, Trainer, TrainingArguments
from sklearn.metrics import accuracy_score

# استيراد الدوال من الملفين المجاورين
from tokenizer import build_tokenizer, load_tokenizer
from model import build_model

# ------------------------------------------------
# توليد البيانات إن لم تكن موجودة
# ------------------------------------------------
def generate_dataset(output_file="xo_dataset.json", n_games=10000):
    """
    تولد مباريات XO عشوائية وتحفظها كـ JSON Lines.
    كل سطر: {"board": "X..O.....", "move": 4}
    """
    def empty_board():
        return ['.'] * 9

    def legal_moves(board):
        return [i for i, cell in enumerate(board) if cell == '.']

    def make_move(board, pos, player):
        new_board = board[:]
        new_board[pos] = player
        return new_board

    def check_win(board, player):
        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 random_game():
        board = empty_board()
        history = []
        players = ['X', 'O']
        turn = 0
        while True:
            player = players[turn % 2]
            moves = legal_moves(board)
            if not moves:
                break
            move = random.choice(moves)
            if player == 'X':
                history.append((''.join(board), move))
            board = make_move(board, move, player)
            if check_win(board, player):
                break
            turn += 1
        return history

    print(f"🔄 جاري توليد {n_games} مباراة عشوائية...")
    all_samples = []
    for _ in range(n_games):
        all_samples.extend(random_game())

    with open(output_file, 'w', encoding='utf-8') as f:
        for board_str, move in all_samples:
            f.write(json.dumps({"board": board_str, "move": move}) + '\n')

    print(f"✅ تم إنشاء {len(all_samples)} عينة في {output_file}")
    return output_file

# ------------------------------------------------
# تجهيز البيانات للتدريب
# ------------------------------------------------
def prepare_datasets(data_file, tokenizer):
    """
    تحميل بيانات JSON، تقسيمها، تطبيق الترميز، وإرجاع مجمّع البيانات.
    """
    dataset = load_dataset('json', data_files=data_file, split='train')
    dataset = dataset.train_test_split(test_size=0.1)
    train_dataset = dataset['train']
    eval_dataset = dataset['test']

    def preprocess_function(examples):
        tokenized = tokenizer(
            examples['board'],
            truncation=True,
            padding=False,
            max_length=12,
        )
        tokenized['labels'] = examples['move']
        return tokenized

    train_dataset = train_dataset.map(
        preprocess_function, batched=True, remove_columns=['board', 'move']
    )
    eval_dataset = eval_dataset.map(
        preprocess_function, batched=True, remove_columns=['board', 'move']
    )

    train_dataset.set_format(type='torch', columns=['input_ids', 'attention_mask', 'labels'])
    eval_dataset.set_format(type='torch', columns=['input_ids', 'attention_mask', 'labels'])

    data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

    print(f"✅ بيانات التدريب: {len(train_dataset)} عينة | الاختبار: {len(eval_dataset)} عينة")
    return train_dataset, eval_dataset, data_collator

# ------------------------------------------------
# التدريب
# ------------------------------------------------
def run_training(model, tokenizer, train_dataset, eval_dataset, data_collator):
    """
    إعداد Trainer وتشغيل التدريب، ثم حفظ النموذج.
    """
    # إصلاحات صغيرة للبيئة
    os.environ["TENSORBOARD_LOGGING_DIR"] = "./logs"
    import datasets.config
    datasets.config.TORCHVISION_AVAILABLE = False

    training_args = TrainingArguments(
        output_dir="./xo_model",
        num_train_epochs=5,
        per_device_train_batch_size=64,
        per_device_eval_batch_size=64,
        eval_strategy="epoch",
        save_strategy="epoch",
        logging_dir="./logs",
        logging_steps=100,
        load_best_model_at_end=True,
        metric_for_best_model="eval_loss",
    )

    def compute_metrics(eval_pred):
        logits, labels = eval_pred
        predictions = np.argmax(logits, axis=-1)
        return {"accuracy": accuracy_score(labels, predictions)}

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        data_collator=data_collator,
        compute_metrics=compute_metrics,
    )

    print("\n🚀 بدء التدريب...")
    trainer.train()

    # حفظ النموذج والـ Tokenizer
    model.save_pretrained("./xo_model")
    tokenizer.save_pretrained("./xo_tokenizer")
    print("✅ تم حفظ النموذج والـ Tokenizer في ./xo_model و ./xo_tokenizer")

# ------------------------------------------------
# البرنامج الرئيسي
# ------------------------------------------------
def main():
    # 1. ملف البيانات
    data_file = "xo_dataset.json"
    if not os.path.exists(data_file):
        generate_dataset(data_file, n_games=10000)
    else:
        print(f"📂 تم العثور على ملف البيانات: {data_file}")

    # 2. Tokenizer (بناء جديد أو تحميل موجود)
    if not os.path.exists("./xo_tokenizer"):
        tokenizer = build_tokenizer(save_path="./xo_tokenizer")
    else:
        tokenizer = load_tokenizer("./xo_tokenizer")

    # 3. بناء النموذج
    model = build_model(tokenizer)

    # 4. تجهيز البيانات
    train_dataset, eval_dataset, data_collator = prepare_datasets(data_file, tokenizer)

    # 5. تدريب
    run_training(model, tokenizer, train_dataset, eval_dataset, data_collator)
    print("\n✨ تم الانتهاء من التدريب بنجاح.")


if __name__ == "__main__":
    main()