#!/usr/bin/env python3
"""Launch Gemini embedding sweep on multiple GPUs.

This script:
1. Checks available GPUs
2. Assigns each embedding variant (256, 512, 768) to a separate GPU
3. Launches training in parallel with proper logging
4. Monitors progress and reports results

Usage:
    python scripts/launch_gemini_sweep.py [--dry-run] [--gpus 0,1,2]
"""

import argparse
import logging
import os
import subprocess
import sys
import time
from pathlib import Path
from typing import List, Optional

_SCRIPT_DIR = Path(__file__).resolve().parent
_MODULE_DIR = _SCRIPT_DIR.parent

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
logger = logging.getLogger(__name__)

# Embedding variants to train
VARIANTS = [
    {"name": "gemini_256", "config": "v8_gemini_256.yaml"},
    {"name": "gemini_512", "config": "v8_gemini_512.yaml"},
    {"name": "gemini_768", "config": "v8_gemini_768.yaml"},
]


def get_available_gpus() -> List[int]:
    """Detect available CUDA GPUs."""
    try:
        import torch
        if torch.cuda.is_available():
            return list(range(torch.cuda.device_count()))
    except ImportError:
        pass

    # Fallback: check nvidia-smi
    try:
        result = subprocess.run(
            ["nvidia-smi", "--query-gpu=index", "--format=csv,noheader"],
            capture_output=True,
            text=True,
            check=True,
        )
        return [int(idx.strip()) for idx in result.stdout.strip().split("\n") if idx.strip()]
    except (subprocess.CalledProcessError, FileNotFoundError):
        return []


def verify_prerequisites() -> bool:
    """Verify all required files exist."""
    all_ok = True

    for variant in VARIANTS:
        config_path = _MODULE_DIR / "registry" / variant["config"]
        if not config_path.exists():
            logger.error(f"Missing config: {config_path}")
            all_ok = False

        # Check data file path from config
        data_path = _MODULE_DIR / "data" / f"router_data_{variant['name']}.parquet"
        if not data_path.exists():
            logger.error(f"Missing data: {data_path}")
            all_ok = False

        # Check LLM embeddings
        llm_path = _MODULE_DIR / "configs" / f"llm_embedding_{variant['name']}.pkl"
        if not llm_path.exists():
            logger.error(f"Missing LLM embeddings: {llm_path}")
            all_ok = False

    if not all_ok:
        logger.error("\nRun 'python scripts/prepare_gemini_sweep.py' first to create required files.")

    return all_ok


def launch_training(
    variant_name: str,
    config_file: str,
    gpu_id: int,
    log_dir: Path,
    dry_run: bool = False,
) -> Optional[subprocess.Popen]:
    """Launch a single training job on a specific GPU.

    Args:
        variant_name: Name of the variant for logging
        config_file: Config filename in registry/
        gpu_id: CUDA device ID
        log_dir: Directory for log files
        dry_run: If True, print command but don't execute

    Returns:
        Popen object for the launched process, or None if dry_run
    """
    config_path = _MODULE_DIR / "registry" / config_file
    log_file = log_dir / f"{variant_name}.log"

    # Build command
    cmd = [
        sys.executable,
        str(_MODULE_DIR / "run_exp.py"),
        "--config_file", str(config_path),
    ]

    env = os.environ.copy()
    env["CUDA_VISIBLE_DEVICES"] = str(gpu_id)

    logger.info(f"Launching {variant_name} on GPU {gpu_id}")
    logger.info(f"  Config: {config_file}")
    logger.info(f"  Log: {log_file}")

    if dry_run:
        logger.info(f"  Command: CUDA_VISIBLE_DEVICES={gpu_id} {' '.join(cmd)}")
        return None

    # Launch process with output redirected to log file
    with open(log_file, "w") as f:
        process = subprocess.Popen(
            cmd,
            env=env,
            stdout=f,
            stderr=subprocess.STDOUT,
            cwd=str(_MODULE_DIR),
        )

    logger.info(f"  Started PID {process.pid}")
    return process


def monitor_processes(
    processes: dict,
    log_dir: Path,
    poll_interval: int = 30,
) -> dict:
    """Monitor running processes and report status.

    Args:
        processes: Dict of {variant_name: Popen}
        log_dir: Directory containing log files
        poll_interval: Seconds between status checks

    Returns:
        Dict of {variant_name: return_code}
    """
    results = {}
    running = set(processes.keys())

    while running:
        time.sleep(poll_interval)

        for name in list(running):
            proc = processes[name]
            ret = proc.poll()

            if ret is not None:
                running.remove(name)
                results[name] = ret

                if ret == 0:
                    logger.info(f"[SUCCESS] {name} completed")
                else:
                    logger.error(f"[FAILED] {name} exited with code {ret}")

                # Print last few lines of log
                log_file = log_dir / f"{name}.log"
                if log_file.exists():
                    with open(log_file) as f:
                        lines = f.readlines()
                        tail = lines[-10:] if len(lines) > 10 else lines
                        logger.info(f"  Last log lines:\n{''.join(tail)}")

        if running:
            logger.info(f"Still running: {', '.join(sorted(running))}")

    return results


def tail_logs(log_dir: Path, lines: int = 5):
    """Print tail of all log files."""
    logger.info("\n--- Current log tails ---")
    for log_file in sorted(log_dir.glob("*.log")):
        logger.info(f"\n{log_file.name}:")
        with open(log_file) as f:
            all_lines = f.readlines()
            tail = all_lines[-lines:] if len(all_lines) > lines else all_lines
            for line in tail:
                print(f"  {line.rstrip()}")


def main():
    parser = argparse.ArgumentParser(
        description="Launch Gemini embedding sweep on multiple GPUs"
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Print commands but don't execute",
    )
    parser.add_argument(
        "--gpus",
        type=str,
        default=None,
        help="Comma-separated GPU IDs to use (default: auto-detect)",
    )
    parser.add_argument(
        "--variants",
        type=str,
        default=None,
        help="Comma-separated variant names to run (default: all)",
    )
    parser.add_argument(
        "--log-dir",
        type=str,
        default=None,
        help="Directory for log files (default: runs/gemini_sweep_<timestamp>)",
    )
    args = parser.parse_args()

    logger.info("=" * 70)
    logger.info("GEMINI EMBEDDING SWEEP LAUNCHER")
    logger.info("=" * 70)

    # Filter variants if specified
    variants_to_run = VARIANTS
    if args.variants:
        variant_names = [v.strip() for v in args.variants.split(",")]
        variants_to_run = [v for v in VARIANTS if v["name"] in variant_names]
        if not variants_to_run:
            logger.error(f"No matching variants found for: {args.variants}")
            sys.exit(1)

    # Verify prerequisites
    if not verify_prerequisites():
        sys.exit(1)

    # Get GPUs
    if args.gpus:
        gpus = [int(g.strip()) for g in args.gpus.split(",")]
    else:
        gpus = get_available_gpus()

    if not gpus:
        logger.error("No GPUs available!")
        sys.exit(1)

    logger.info(f"Available GPUs: {gpus}")
    logger.info(f"Variants to run: {[v['name'] for v in variants_to_run]}")

    if len(gpus) < len(variants_to_run):
        logger.warning(
            f"Only {len(gpus)} GPUs available for {len(variants_to_run)} variants. "
            "Some variants will run sequentially."
        )

    # Create log directory
    if args.log_dir:
        log_dir = Path(args.log_dir)
    else:
        timestamp = time.strftime("%Y%m%d_%H%M%S")
        log_dir = _MODULE_DIR / "runs" / f"gemini_sweep_{timestamp}"

    log_dir.mkdir(parents=True, exist_ok=True)
    logger.info(f"Log directory: {log_dir}")

    # Launch processes
    processes = {}
    for i, variant in enumerate(variants_to_run):
        gpu_id = gpus[i % len(gpus)]
        proc = launch_training(
            variant_name=variant["name"],
            config_file=variant["config"],
            gpu_id=gpu_id,
            log_dir=log_dir,
            dry_run=args.dry_run,
        )
        if proc:
            processes[variant["name"]] = proc

    if args.dry_run:
        logger.info("\n[DRY RUN] No processes launched")
        return

    # Monitor progress
    logger.info("\nMonitoring training progress...")
    logger.info("(Press Ctrl+C to stop monitoring - training will continue in background)")

    try:
        results = monitor_processes(processes, log_dir)

        # Summary
        logger.info("\n" + "=" * 70)
        logger.info("SWEEP COMPLETE")
        logger.info("=" * 70)
        for name, ret in results.items():
            status = "SUCCESS" if ret == 0 else f"FAILED (code {ret})"
            logger.info(f"  {name}: {status}")

    except KeyboardInterrupt:
        logger.info("\nMonitoring stopped. Training continues in background.")
        logger.info(f"Check logs in: {log_dir}")
        for name, proc in processes.items():
            if proc.poll() is None:
                logger.info(f"  {name}: PID {proc.pid} still running")


if __name__ == "__main__":
    main()
