import os
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
import time
import glob
from tqdm import tqdm

from config import *
from model import PegeModel
from tokenizer import TurkishTokenizer
from dataset import TextDataset

def train():
    print("=" * 50)
    print("PEGE TRAINING STARTED")
    print("=" * 50)

    # Tokenizer
    tokenizer = TurkishTokenizer(vocab_size=VOCAB_SIZE)
    tokenizer_path = os.path.expanduser(f"{MODEL_DIR}/tokenizer.pkl")

    data_dir = os.path.expanduser(DATA_DIR)
    txt_files = glob.glob(f"{data_dir}/**/*.txt", recursive=True)
    if not os.path.exists(data_dir) or not txt_files:
        print(f"\nUYARI: {data_dir} klasöründe .txt dosyası bulunamadı!")
        print("Lütfen makalelerini TXT olarak buraya at:")
        print(f"  {data_dir}/")
        return

    tokenizer_json = tokenizer_path.replace(".pkl", ".json")
    if os.path.exists(tokenizer_json):
        tokenizer.load(tokenizer_path)
    else:
        tokenizer.train(txt_files)
        tokenizer.save(tokenizer_path)

    dataset = TextDataset(data_dir, tokenizer, max_length=MAX_SEQ_LEN)
    dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=0)

    vocab_size = len(tokenizer.stoi)
    print(f"Vocabulary size: {vocab_size}")

    # Model
    model = PegeModel(vocab_size).to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE)

    # Checkpoint yükle
    start_step = 0
    checkpoint_files = glob.glob(os.path.expanduser(f"{CHECKPOINT_DIR}/checkpoint_*.pt"))
    if checkpoint_files:
        latest = max(checkpoint_files, key=os.path.getctime)
        checkpoint = torch.load(latest, map_location=device)
        model.load_state_dict(checkpoint['model_state'])
        optimizer.load_state_dict(checkpoint['optimizer_state'])
        start_step = checkpoint['step']
        print(f"Resumed from checkpoint: {latest} (step {start_step})")

    model.train()

    # Training loop
    step = start_step
    losses = []

    print(f"\nTraining on {len(dataset)} samples")
    print(f"Batch size: {BATCH_SIZE}")
    print(f"Device: {device}\n")

    progress = tqdm(total=MAX_STEPS, initial=start_step, desc="Training")

    while step < MAX_STEPS:
        for batch_x, batch_y in dataloader:
            batch_x = batch_x.to(device)
            batch_y = batch_y.to(device)

            # Forward
            logits, loss = model(batch_x, batch_y)

            # Backward
            optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()

            losses.append(loss.item())

            # Logging
            if step % 100 == 0:
                avg_loss = sum(losses[-100:]) / min(len(losses), 100)
                progress.set_postfix({'loss': f'{avg_loss:.4f}'})

            # Checkpoint
            if step % CHECKPOINT_INTERVAL == 0 and step > 0:
                os.makedirs(os.path.expanduser(CHECKPOINT_DIR), exist_ok=True)
                checkpoint_path = os.path.expanduser(f"{CHECKPOINT_DIR}/checkpoint_{step}.pt")
                torch.save({
                    'step': step,
                    'model_state': model.state_dict(),
                    'optimizer_state': optimizer.state_dict(),
                    'vocab_size': vocab_size,
                }, checkpoint_path)

                # Eski checkpoint'leri sil
                old_checkpoints = sorted(glob.glob(os.path.expanduser(f"{CHECKPOINT_DIR}/checkpoint_*.pt")))
                for old in old_checkpoints[:-3]:  # Son 3'ü tut
                    os.remove(old)

                print(f"\nCheckpoint saved: {checkpoint_path}")

            step += 1
            progress.update(1)

            if step >= MAX_STEPS:
                break

    progress.close()

    # Final save
    os.makedirs(os.path.expanduser(MODEL_DIR), exist_ok=True)
    final_path = os.path.expanduser(f"{MODEL_DIR}/pege_final.pt")
    torch.save({
        'model_state': model.state_dict(),
        'vocab_size': vocab_size,
        'config': {
            'embed_dim': EMBED_DIM,
            'num_heads': NUM_HEADS,
            'num_layers': NUM_LAYERS,
            'hidden_dim': HIDDEN_DIM,
        }
    }, final_path)

    print(f"\n{'=' * 50}")
    print(f"Training complete! Model saved to: {final_path}")
    print(f"{'=' * 50}")

if __name__ == "__main__":
    train()