import os
import sys
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

import argparse
import time
import math
import json
from typing import Tuple, Dict, Any, List
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, random_split
from nlp_core.tokenizer import BilingualTokenizer
from nlp_core.vocabulary import Vocabulary
from models.abstractive.seq2seq_model import Seq2SeqSummarizer
from training.dataset import SummarizationDataset, pad_collate_fn
from training.sample_data_gen import generate_synthetic_dataset
from evaluation.metrics_manager import MetricsManager

def build_vocabularies(data_path: str, lang: str = "en") -> Tuple[Vocabulary, Vocabulary]:
    """Builds source and target vocabularies from dataset file."""
    tokenizer = BilingualTokenizer()
    src_vocab = Vocabulary(name=f"{lang}_src_vocab")
    trg_vocab = Vocabulary(name=f"{lang}_trg_vocab")

    with open(data_path, "r", encoding="utf-8") as f:
        data = json.load(f)

    for item in data:
        src_tokens = tokenizer.tokenize_words(item["article"], lang=lang, remove_stopwords=False)
        trg_tokens = tokenizer.tokenize_words(item["summary"], lang=lang, remove_stopwords=False)
        src_vocab.add_sentence(src_tokens)
        trg_vocab.add_sentence(trg_tokens)

    src_vocab.build_vocab(min_freq=1, max_size=15000)
    trg_vocab.build_vocab(min_freq=1, max_size=15000)

    return src_vocab, trg_vocab

def train_epoch(model, dataloader, optimizer, criterion, clip: float = 1.0, teacher_forcing_ratio: float = 0.5):
    model.train()
    epoch_loss = 0.0

    for src, trg in dataloader:
        src = src.to(model.device)
        trg = trg.to(model.device)

        optimizer.zero_grad()
        output = model(src, trg, teacher_forcing_ratio=teacher_forcing_ratio)

        # trg shape: [batch_size, trg_len]
        # output shape: [batch_size, trg_len, trg_vocab_size]
        trg_vocab_size = output.shape[-1]
        
        # Flatten excluding the <sos> token at index 0
        output_dim = output[:, 1:].reshape(-1, trg_vocab_size)
        trg_dim = trg[:, 1:].reshape(-1)

        loss = criterion(output_dim, trg_dim)
        loss.backward()

        torch.nn.utils.clip_grad_norm_(model.parameters(), clip)
        optimizer.step()

        epoch_loss += loss.item()

    return epoch_loss / len(dataloader)

def evaluate_epoch(model, dataloader, criterion):
    model.eval()
    epoch_loss = 0.0

    with torch.no_grad():
        for src, trg in dataloader:
            src = src.to(model.device)
            trg = trg.to(model.device)

            output = model(src, trg, teacher_forcing_ratio=0.0) # No teacher forcing during eval
            trg_vocab_size = output.shape[-1]
            output_dim = output[:, 1:].reshape(-1, trg_vocab_size)
            trg_dim = trg[:, 1:].reshape(-1)

            loss = criterion(output_dim, trg_dim)
            epoch_loss += loss.item()

    return epoch_loss / len(dataloader)

def evaluate_model_rouge_bleu(
    model: Seq2SeqSummarizer,
    val_dataset,
    raw_data_file: str,
    lang: str = "en",
    num_samples: int = 15
) -> Dict[str, float]:
    """Computes empirical ROUGE and BLEU metrics over validation samples."""
    tokenizer = BilingualTokenizer()
    metrics_mgr = MetricsManager()

    with open(raw_data_file, "r", encoding="utf-8") as f:
        all_samples = json.load(f)

    # Pick validation sample subset
    eval_samples = all_samples[-num_samples:] if len(all_samples) >= num_samples else all_samples

    rouge1_f1_list, rouge2_f1_list, rougel_f1_list = [], [], []
    bleu1_list, bleu2_list, bleu_cum_list = [], [], []

    for item in eval_samples:
        src_tokens = tokenizer.tokenize_words(item["article"], lang=lang)
        pred_tokens = model.summarize_beam(src_tokens, beam_width=3, max_len=40)
        pred_text = " ".join(pred_tokens)
        ref_text = item["summary"]

        eval_res = metrics_mgr.evaluate_summary(
            original_text=item["article"],
            generated_summary=pred_text,
            reference_summary=ref_text,
            lang=lang
        )

        if "rouge" in eval_res and "bleu" in eval_res:
            rouge1_f1_list.append(eval_res["rouge"]["rouge-1"]["f1"])
            rouge2_f1_list.append(eval_res["rouge"]["rouge-2"]["f1"])
            rougel_f1_list.append(eval_res["rouge"]["rouge-l"]["f1"])

            bleu1_list.append(eval_res["bleu"]["bleu-1"])
            bleu2_list.append(eval_res["bleu"]["bleu-2"])
            bleu_cum_list.append(eval_res["bleu"]["bleu_cumulative"])

    avg_metrics = {
        "rouge_1_f1": sum(rouge1_f1_list) / len(rouge1_f1_list) if rouge1_f1_list else 0.0,
        "rouge_2_f1": sum(rouge2_f1_list) / len(rouge2_f1_list) if rouge2_f1_list else 0.0,
        "rouge_l_f1": sum(rougel_f1_list) / len(rougel_f1_list) if rougel_f1_list else 0.0,
        "bleu_1": sum(bleu1_list) / len(bleu1_list) if bleu1_list else 0.0,
        "bleu_2": sum(bleu2_list) / len(bleu2_list) if bleu2_list else 0.0,
        "cumulative_bleu": sum(bleu_cum_list) / len(bleu_cum_list) if bleu_cum_list else 0.0,
    }
    return avg_metrics

def run_training(
    lang: str = "ar",
    epochs: int = 20,
    batch_size: int = 8,
    lr: float = 0.002,
    emb_dim: int = 128,
    enc_hid_dim: int = 256,
    dec_hid_dim: int = 256,
    device: str = "cpu"
):
    if sys.stdout.encoding != 'utf-8':
        try:
            sys.stdout.reconfigure(encoding='utf-8')
        except Exception:
            pass

    print(f"\n========================================================")
    print(f"[*] Starting Seq2Seq Training for [{lang.upper()}] (Kaggle Benchmark)")
    print(f"========================================================")

    # Ensure dataset exists
    data_dir = "data/datasets"
    data_file = os.path.join(data_dir, f"{'arabic' if lang == 'ar' else 'english'}_corpus.json")
    if not os.path.exists(data_file):
        print(f"Dataset not found at {data_file}. Generating synthetic dataset...")
        generate_synthetic_dataset(data_dir)

    # Build vocabularies
    src_vocab, trg_vocab = build_vocabularies(data_file, lang=lang)
    print(f"Source Vocabulary Size: {len(src_vocab)} | Target Vocabulary Size: {len(trg_vocab)}")

    # Prepare Dataset & DataLoader
    full_dataset = SummarizationDataset(data_file, src_vocab, trg_vocab, lang=lang)
    train_size = int(0.85 * len(full_dataset))
    val_size = len(full_dataset) - train_size
    train_dataset, val_dataset = random_split(full_dataset, [train_size, val_size])

    train_loader = DataLoader(
        train_dataset,
        batch_size=batch_size,
        shuffle=True,
        collate_fn=lambda b: pad_collate_fn(b, pad_idx=src_vocab.PAD_IDX)
    )
    val_loader = DataLoader(
        val_dataset,
        batch_size=batch_size,
        shuffle=False,
        collate_fn=lambda b: pad_collate_fn(b, pad_idx=src_vocab.PAD_IDX)
    )

    # Initialize Model
    model = Seq2SeqSummarizer(
        src_vocab=src_vocab,
        trg_vocab=trg_vocab,
        emb_dim=emb_dim,
        enc_hid_dim=enc_hid_dim,
        dec_hid_dim=dec_hid_dim,
        dropout=0.2,
        device=device
    )

    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs, eta_min=1e-5)
    criterion = nn.CrossEntropyLoss(ignore_index=trg_vocab.PAD_IDX)

    best_val_loss = float("inf")
    checkpoint_path = f"checkpoints/seq2seq_{'arabic' if lang == 'ar' else 'english'}.pt"

    for epoch in range(1, epochs + 1):
        start_time = time.time()
        # Scheduled teacher forcing decay
        tf_ratio = max(0.2, 0.85 - (epoch / epochs) * 0.5)

        train_loss = train_epoch(model, train_loader, optimizer, criterion, teacher_forcing_ratio=tf_ratio)
        val_loss = evaluate_epoch(model, val_loader, criterion)
        scheduler.step()
        elapsed = time.time() - start_time

        print(f"Epoch {epoch:02d}/{epochs:02d} | Train Loss: {train_loss:.4f} | Val Loss: {val_loss:.4f} | LR: {scheduler.get_last_lr()[0]:.6f} | Time: {elapsed:.2f}s")

        if val_loss < best_val_loss:
            best_val_loss = val_loss
            model.save_checkpoint(checkpoint_path)
            print(f"  [+] Saved Best Checkpoint -> {checkpoint_path}")

    print(f"\n[+] Training finished successfully for {lang.upper()}! Best Val Loss: {best_val_loss:.4f}")

    # Load best checkpoint for evaluation
    best_model = Seq2SeqSummarizer.load_checkpoint(checkpoint_path, device=device)
    print(f"\n[*] Evaluating Quantitative Metrics (ROUGE & BLEU) on Validation Benchmark...")
    eval_metrics = evaluate_model_rouge_bleu(best_model, val_dataset, data_file, lang=lang)
    print(f"--------------------------------------------------------")
    print(f"  • ROUGE-1 F1 Score:  {eval_metrics['rouge_1_f1'] * 100:.2f}%")
    print(f"  • ROUGE-2 F1 Score:  {eval_metrics['rouge_2_f1'] * 100:.2f}%")
    print(f"  • ROUGE-L F1 Score:  {eval_metrics['rouge_l_f1'] * 100:.2f}%")
    print(f"  • BLEU-1 Score:      {eval_metrics['bleu_1'] * 100:.2f}%")
    print(f"  • BLEU-2 Score:      {eval_metrics['bleu_2'] * 100:.2f}%")
    print(f"  • Cumulative BLEU:   {eval_metrics['cumulative_bleu'] * 100:.2f}%")
    print(f"--------------------------------------------------------")

    # Test sample inference
    tokenizer = BilingualTokenizer()
    if lang == "ar":
        sample_text = "تعد معالجة اللغات الطبيعية فرعا جوهريا من فروع الذكاء الاصطناعي يهدف إلى جعل الحواسيب قادرة على فهم النصوص البشرية وتوليدها بفعالية."
        ref_text = "معالجة اللغات الطبيعية تمكن الحواسيب من فهم النصوص البشرية وتوليدها واستخراج المعلومات."
    else:
        sample_text = "Natural language processing is an essential branch of artificial intelligence enabling computing systems to comprehend and synthesize human language."
        ref_text = "Natural language processing enables computers to understand, process, and generate human language seamlessly."

    sample_tokens = tokenizer.tokenize_words(sample_text, lang=lang)
    greedy_summary, _ = best_model.summarize_greedy(sample_tokens, max_len=30)
    beam_summary = best_model.summarize_beam(sample_tokens, beam_width=3, max_len=30)

    print(f"\n[Sample Test Inference]")
    print(f"Input:       {sample_text}")
    print(f"Reference:   {ref_text}")
    print(f"Greedy:      {' '.join(greedy_summary)}")
    print(f"Beam (k=3):  {' '.join(beam_summary)}\n")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Train Bilingual Seq2Seq Text Summarization Model on Kaggle Benchmark")
    parser.add_argument("--lang", type=str, default="all", choices=["ar", "en", "all"], help="Language to train ('ar', 'en', or 'all')")
    parser.add_argument("--epochs", type=int, default=20, help="Number of training epochs")
    parser.add_argument("--batch_size", type=int, default=8, help="Batch size")
    parser.add_argument("--lr", type=float, default=0.002, help="Learning rate")
    parser.add_argument("--device", type=str, default="cpu", help="Device ('cpu' or 'cuda')")
    args = parser.parse_args()

    if args.lang == "all":
        run_training("ar", epochs=args.epochs, batch_size=args.batch_size, lr=args.lr, device=args.device)
        run_training("en", epochs=args.epochs, batch_size=args.batch_size, lr=args.lr, device=args.device)
    else:
        run_training(args.lang, epochs=args.epochs, batch_size=args.batch_size, lr=args.lr, device=args.device)