"""
ViuAI-1M Training Script.

End-to-end training pipeline:
1. Load or prepare data
2. Train BPE tokenizer on the corpus
3. Create datasets and dataloaders
4. Initialize model
5. Train with the ViuAI trainer
6. Save the final model and tokenizer

Usage:
    python scripts/train.py
    python scripts/train.py --epochs 20 --batch_size 16 --lr 1e-4
"""

import argparse
import os
import sys
import json

# Add project root to path
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, project_root)

import torch
from torch.utils.data import DataLoader

from viuai.model.config import ViuAIConfig
from viuai.model.viuai_model import ViuAIForCausalLM
from viuai.tokenizer.bpe_tokenizer import ViuAITokenizer
from viuai.training.dataset import ViuAIDataset, ChatDataCollator
from viuai.training.trainer import ViuAITrainer


def parse_args():
    parser = argparse.ArgumentParser(description="Train ViuAI-1M")
    parser.add_argument("--config", type=str, default="configs/viuai_1m.json",
                        help="Path to model config JSON")
    parser.add_argument("--train_data", type=str, default="data/train.jsonl",
                        help="Path to training data JSONL")
    parser.add_argument("--eval_data", type=str, default="data/eval.jsonl",
                        help="Path to evaluation data JSONL")
    parser.add_argument("--tokenizer_corpus", type=str, default="data/tokenizer_corpus.txt",
                        help="Path to tokenizer training corpus")
    parser.add_argument("--output_dir", type=str, default="output",
                        help="Output directory for checkpoints and tokenizer")
    parser.add_argument("--epochs", type=int, default=10,
                        help="Number of training epochs")
    parser.add_argument("--batch_size", type=int, default=32,
                        help="Training batch size")
    parser.add_argument("--lr", type=float, default=3e-4,
                        help="Learning rate")
    parser.add_argument("--warmup_steps", type=int, default=100,
                        help="Number of warmup steps")
    parser.add_argument("--device", type=str, default=None,
                        help="Device (cuda/cpu/mps)")
    return parser.parse_args()


def main():
    args = parse_args()

    print(f"\n{'='*60}")
    print(f"  ViuAI-1M Training Pipeline")
    print(f"{'='*60}\n")

    # ── Step 1: Load Config ──────────────────────────────────────
    print("[1/6] Loading configuration...")
    config_path = os.path.join(project_root, args.config)
    with open(config_path, "r") as f:
        config_data = json.load(f)

    config = ViuAIConfig(**config_data["architecture"])
    print(f"  Config: {config_path}")
    print(f"  Estimated parameters: {config.num_parameters:,}")

    # ── Step 2: Train Tokenizer ──────────────────────────────────
    print(f"\n[2/6] Training BPE tokenizer...")
    tokenizer = ViuAITokenizer(vocab_size=config.vocab_size)

    corpus_path = os.path.join(project_root, args.tokenizer_corpus)
    if os.path.exists(corpus_path):
        with open(corpus_path, "r", encoding="utf-8") as f:
            corpus_texts = [line.strip() for line in f if line.strip()]
        tokenizer.train(corpus_texts, verbose=True)
    else:
        print(f"  WARNING: Corpus not found at {corpus_path}")
        print(f"  Run 'python scripts/prepare_data.py' first!")
        return

    # Update config vocab size to match actual tokenizer
    config.vocab_size = tokenizer.vocab_size

    # Save tokenizer
    tokenizer_dir = os.path.join(project_root, args.output_dir, "tokenizer")
    tokenizer.save(tokenizer_dir)

    # ── Step 3: Create Datasets ──────────────────────────────────
    print(f"\n[3/6] Creating datasets...")
    train_path = os.path.join(project_root, args.train_data)
    eval_path = os.path.join(project_root, args.eval_data)

    train_dataset = ViuAIDataset(train_path, tokenizer, max_seq_len=config.max_seq_len)
    eval_dataset = ViuAIDataset(eval_path, tokenizer, max_seq_len=config.max_seq_len)

    collator = ChatDataCollator(pad_token_id=config.pad_token_id)

    train_loader = DataLoader(
        train_dataset,
        batch_size=args.batch_size,
        shuffle=True,
        collate_fn=collator,
        num_workers=0,
    )
    eval_loader = DataLoader(
        eval_dataset,
        batch_size=args.batch_size,
        shuffle=False,
        collate_fn=collator,
        num_workers=0,
    )

    print(f"  Train batches: {len(train_loader)}")
    print(f"  Eval batches:  {len(eval_loader)}")

    # ── Step 4: Initialize Model ─────────────────────────────────
    print(f"\n[4/6] Initializing ViuAI-1M model...")
    model = ViuAIForCausalLM(config)
    print(model.get_parameter_summary())

    # ── Step 5: Train ────────────────────────────────────────────
    print(f"\n[5/6] Starting training...")
    checkpoint_dir = os.path.join(project_root, args.output_dir, "checkpoints")

    trainer = ViuAITrainer(
        model=model,
        train_dataloader=train_loader,
        eval_dataloader=eval_loader,
        learning_rate=args.lr,
        weight_decay=0.01,
        warmup_steps=args.warmup_steps,
        num_epochs=args.epochs,
        save_dir=checkpoint_dir,
        log_interval=10,
        eval_interval=50,
        device=args.device,
    )

    metrics = trainer.train()

    # ── Step 6: Save Final Model ─────────────────────────────────
    print(f"\n[6/6] Saving final model...")
    final_dir = os.path.join(project_root, args.output_dir, "final")
    os.makedirs(final_dir, exist_ok=True)

    torch.save(model.state_dict(), os.path.join(final_dir, "model.pt"))
    config.save(os.path.join(final_dir, "config.json"))

    print(f"  Model saved to: {final_dir}")
    print(f"\n  To chat with your model:")
    print(f"    python scripts/chat.py --checkpoint {final_dir} --tokenizer {tokenizer_dir}")


if __name__ == "__main__":
    main()