#!/usr/bin/env python3
"""
Train BlogLLM directly from JSONL without database dependency.
Loads blog-training-alpaca.jsonl and runs LoRA fine-tuning.

Usage:
    python3 scripts/train_blog_direct.py
    python3 scripts/train_blog_direct.py --model Qwen/Qwen2.5-7B-Instruct --output models/fo-blog-v4
"""

import json
import sys
from pathlib import Path
from typing import Optional

# Add package to path
REPO_ROOT = Path(__file__).parent.parent
sys.path.insert(0, str(REPO_ROOT.parent.parent))
sys.path.insert(0, str(REPO_ROOT))

from src.trainer import run_lora_training
import yaml


def load_jsonl(path: Path) -> list[dict]:
    """Load Alpaca-format JSONL samples and convert to trainer format."""
    samples = []
    with open(path) as f:
        for line in f:
            if line.strip():
                sample = json.loads(line)
                # Convert Alpaca format to trainer's expected format
                # trainer.py expects: system_prompt, input_text, output_text
                samples.append({
                    "system_prompt": "You are an expert blog writer specializing in optical transceiver technology and network infrastructure. Write clear, practical blog posts targeted at network engineers and architects.",
                    "input_text": sample.get("instruction", ""),
                    "output_text": sample.get("output", ""),
                })
    return samples


def main():
    import argparse

    parser = argparse.ArgumentParser(
        description="Train BlogLLM from JSONL directly",
    )
    parser.add_argument(
        "--model",
        default="Qwen/Qwen2.5-14B-Instruct",
        help="Base model ID (default: Qwen/Qwen2.5-14B-Instruct)",
    )
    parser.add_argument(
        "--output",
        default="models/fo-blog-v4",
        help="Output directory for adapters",
    )
    parser.add_argument(
        "--config",
        default="config/fine_tuner.yaml",
        help="Path to fine_tuner.yaml",
    )
    parser.add_argument(
        "--jsonl",
        default="data/blog-training-alpaca.jsonl",
        help="Path to training JSONL",
    )

    args = parser.parse_args()

    # Load config
    config_path = REPO_ROOT / args.config
    with open(config_path) as f:
        cfg = yaml.safe_load(f)

    training_cfg = cfg["training"]
    eval_cfg = cfg["evaluation"]
    llama_cfg = cfg["llama_cpp"]
    models_cfg = cfg["models"]

    # Load training data
    jsonl_path = REPO_ROOT / args.jsonl
    print(f"📂 Loading {jsonl_path}...")

    if not jsonl_path.exists():
        print(f"❌ File not found: {jsonl_path}")
        sys.exit(1)

    samples = load_jsonl(jsonl_path)
    print(f"✅ Loaded {len(samples)} training samples")

    # Split into train/val (80/20)
    split_idx = int(len(samples) * 0.8)
    train_samples = samples[:split_idx]
    val_samples = samples[split_idx:]

    print(f"   Train: {len(train_samples)} | Val: {len(val_samples)}")

    # Run training
    output_dir = REPO_ROOT / args.output
    output_dir.mkdir(parents=True, exist_ok=True)

    print(f"\n🚀 Starting LoRA fine-tuning...")
    print(f"   Model: {args.model}")
    print(f"   Output: {output_dir}")
    print(f"   Epochs: {training_cfg['sft']['num_epochs']}")
    print(f"   LoRA r={training_cfg['lora_r']}, alpha={training_cfg['lora_alpha']}")
    print(f"   Max seq length: {training_cfg['max_seq_length']}")
    print()

    try:
        metrics = run_lora_training(
            base_model_path=args.model,
            train_examples=train_samples,
            val_examples=val_samples,
            output_dir=str(output_dir),
            task_type="tip_blog",
            lora_r=training_cfg["lora_r"],
            lora_alpha=training_cfg["lora_alpha"],
            lora_dropout=training_cfg["lora_dropout"],
            max_seq_length=training_cfg["max_seq_length"],
            num_epochs=training_cfg["sft"]["num_epochs"],
            batch_size=training_cfg["sft"]["batch_size"],
            gradient_accumulation_steps=training_cfg["sft"]["gradient_accumulation"],
            learning_rate=training_cfg["sft"]["learning_rate"],
            warmup_ratio=training_cfg["sft"]["warmup_ratio"],
        )

        print(f"\n✅ Training complete!")
        print(f"   Metrics: {metrics}")
        print(f"\n📦 Converting to GGUF for Ollama...")

        # Optional: Convert to GGUF
        gguf_path = output_dir / "model.gguf"
        print(f"   Output: {gguf_path}")
        print(f"\n🎯 Next: Load into Ollama with: ollama create fo-blog-v4 -f Modelfile")

    except Exception as e:
        print(f"\n❌ Training failed: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()