#!/usr/bin/env python
"""
WaveDL Training Launcher.

This module provides a universal training launcher that wraps accelerate
for distributed training. It works seamlessly on both:
- Local machines (uses standard cache locations)
- HPC clusters (uses local caching, offline WandB)

The environment is auto-detected based on scheduler variables (SLURM, PBS, etc.)
and home directory writability.

Usage:
    # Local machine or HPC - same command!
    wavedl-train --model cnn --data_path train.npz --output_dir results

    # Multi-GPU is automatic (uses all available GPUs)
    wavedl-train --model resnet18 --data_path train.npz --num_gpus 4

Example SLURM script:
    #!/bin/bash
    #SBATCH --nodes=1
    #SBATCH --gpus-per-node=4
    #SBATCH --time=12:00:00

    wavedl-train --model cnn --data_path /scratch/data.npz --compile

Author: Ductho Le (ductho.le@outlook.com)
"""

from __future__ import annotations

import argparse
import os
import shutil
import subprocess
import sys
from pathlib import Path


def detect_gpus() -> int:
    """Auto-detect available GPUs using nvidia-smi."""
    if shutil.which("nvidia-smi") is None:
        print("Warning: nvidia-smi not found, defaulting to NUM_GPUS=1")
        return 1

    try:
        result = subprocess.run(
            ["nvidia-smi", "--list-gpus"],
            capture_output=True,
            text=True,
            check=True,
        )
        gpu_output = result.stdout.strip()
        gpu_count = len(gpu_output.split("\n")) if gpu_output else 0
        if gpu_count > 0:
            print(f"Auto-detected {gpu_count} GPU(s)")
            return gpu_count
    except (subprocess.CalledProcessError, FileNotFoundError):
        pass

    print("Warning: No GPUs detected, defaulting to NUM_GPUS=1")
    return 1


def is_hpc_environment() -> bool:
    """Detect if running on an HPC cluster.

    Checks for:
    1. Common HPC scheduler environment variables (SLURM, PBS, LSF, SGE, Cobalt)
    2. Non-writable home directory (common on HPC systems)

    Returns:
        True if HPC environment detected, False otherwise.
    """
    # Check for common HPC scheduler environment variables
    hpc_indicators = [
        "SLURM_JOB_ID",  # SLURM
        "PBS_JOBID",  # PBS/Torque
        "LSB_JOBID",  # LSF
        "SGE_TASK_ID",  # Sun Grid Engine
        "COBALT_JOBID",  # Cobalt
    ]
    if any(var in os.environ for var in hpc_indicators):
        return True

    # Check if home directory is not writable (common on HPC)
    home = os.path.expanduser("~")
    return not os.access(home, os.W_OK)


def setup_environment() -> None:
    """Configure environment for HPC or local machine.

    Automatically detects the environment and configures accordingly:
    - HPC: Uses CWD-based caching, offline WandB (compute nodes lack internet)
    - Local: Uses standard cache locations (~/.cache), doesn't override WandB
    """
    is_hpc = is_hpc_environment()

    if is_hpc:
        # HPC: use CWD-based caching (compute nodes lack internet)
        cache_base = os.getcwd()

        # TORCH_HOME set to CWD - compute nodes need pre-cached weights
        os.environ.setdefault("TORCH_HOME", f"{cache_base}/.torch_cache")
        Path(os.environ["TORCH_HOME"]).mkdir(parents=True, exist_ok=True)

        # Triton/Inductor caches - prevents permission errors with --compile
        os.environ.setdefault("TRITON_CACHE_DIR", f"{cache_base}/.triton_cache")
        os.environ.setdefault(
            "TORCHINDUCTOR_CACHE_DIR", f"{cache_base}/.inductor_cache"
        )
        Path(os.environ["TRITON_CACHE_DIR"]).mkdir(parents=True, exist_ok=True)
        Path(os.environ["TORCHINDUCTOR_CACHE_DIR"]).mkdir(parents=True, exist_ok=True)

        # Check if home is writable for other caches
        home = os.path.expanduser("~")
        home_writable = os.access(home, os.W_OK)

        # Other caches only if home is not writable
        if not home_writable:
            os.environ.setdefault("MPLCONFIGDIR", f"{cache_base}/.matplotlib")
            os.environ.setdefault("FONTCONFIG_CACHE", f"{cache_base}/.fontconfig")
            os.environ.setdefault("XDG_CACHE_HOME", f"{cache_base}/.cache")

            for env_var in [
                "MPLCONFIGDIR",
                "FONTCONFIG_CACHE",
                "XDG_CACHE_HOME",
            ]:
                Path(os.environ[env_var]).mkdir(parents=True, exist_ok=True)

        # WandB configuration (offline by default for HPC)
        os.environ.setdefault("WANDB_MODE", "offline")
        os.environ.setdefault("WANDB_DIR", f"{cache_base}/.wandb")
        os.environ.setdefault("WANDB_CACHE_DIR", f"{cache_base}/.wandb_cache")
        os.environ.setdefault("WANDB_CONFIG_DIR", f"{cache_base}/.wandb_config")

        print("🖥️  HPC environment detected - using local caching")
    else:
        # Local machine: use standard locations, don't override user settings
        # TORCH_HOME defaults to ~/.cache/torch (PyTorch default)
        # WANDB_MODE defaults to online (WandB default)
        print("💻 Local environment detected - using standard cache locations")

    # Suppress non-critical warnings (both environments)
    os.environ.setdefault(
        "PYTHONWARNINGS",
        "ignore::UserWarning,ignore::FutureWarning,ignore::DeprecationWarning",
    )


def handle_fast_path_args() -> int | None:
    """Handle utility flags that don't need accelerate launch.

    Returns:
        Exit code if handled (0 for success), None if should continue to full launch.
    """
    # --list_models: print models and exit immediately
    if "--list_models" in sys.argv:
        from wavedl.models import list_models

        print("Available models:")
        for name in list_models():
            print(f"  {name}")
        return 0

    return None  # Continue to full launch


def parse_args() -> tuple[argparse.Namespace, list[str]]:
    """Parse launcher-specific arguments, pass remaining to wavedl.train."""
    parser = argparse.ArgumentParser(
        description="WaveDL Training Launcher (works on local machines and HPC clusters)",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Basic training (auto-detects GPUs and environment)
  wavedl-train --model cnn --data_path train.npz --output_dir results

  # Specify GPU count explicitly
  wavedl-train --model cnn --data_path train.npz --num_gpus 4

  # Full configuration
  wavedl-train --model resnet18 --data_path train.npz --batch_size 256 \\
               --lr 1e-3 --epochs 100 --compile --output_dir ./results

  # List available models
  wavedl-train --list_models

Environment Detection:
  The launcher automatically detects your environment:
  - HPC (SLURM, PBS, etc.): Uses local caching, offline WandB
  - Local machine: Uses standard cache locations (~/.cache)

For full training options, see: python -m wavedl.train --help
""",
    )

    # HPC-specific arguments
    parser.add_argument(
        "--num_gpus",
        type=int,
        default=None,
        help="Number of GPUs to use (default: auto-detect)",
    )
    parser.add_argument(
        "--num_machines",
        type=int,
        default=1,
        help="Number of machines for multi-node training (default: 1)",
    )
    parser.add_argument(
        "--machine_rank",
        type=int,
        default=0,
        help="Rank of this machine in multi-node setup (default: 0)",
    )
    parser.add_argument(
        "--main_process_ip",
        type=str,
        default=None,
        help="IP address of the main process for multi-node training",
    )
    parser.add_argument(
        "--main_process_port",
        type=int,
        default=None,
        help="Port for multi-node communication (default: accelerate auto-selects)",
    )
    parser.add_argument(
        "--mixed_precision",
        type=str,
        choices=["bf16", "fp16", "no"],
        default="bf16",
        help="Mixed precision mode (default: bf16)",
    )
    parser.add_argument(
        "--dynamo_backend",
        type=str,
        default="no",
        help="PyTorch dynamo backend (default: no)",
    )

    # Parse known args, pass rest to wavedl.train
    args, remaining = parser.parse_known_args()
    return args, remaining


def print_summary(
    exit_code: int, wandb_enabled: bool, wandb_mode: str, wandb_dir: str
) -> None:
    """Print post-training summary and instructions."""
    print()
    print("=" * 40)

    if exit_code == 0:
        print("✅ Training completed successfully!")
        print("=" * 40)

        # Only show WandB sync instructions if user enabled wandb
        if wandb_enabled and wandb_mode == "offline":
            print()
            print("📊 WandB Sync Instructions:")
            print("   From the login node, run:")
            print(f"   wandb sync {wandb_dir}/wandb/offline-run-*")
            print()
            print("   This will upload your training logs to wandb.ai")
    else:
        print(f"❌ Training failed with exit code: {exit_code}")
        print("=" * 40)
        print()
        print("Common issues:")
        print("  - Missing data file (check --data_path)")
        print("  - Insufficient GPU memory (reduce --batch_size)")
        print("  - Invalid model name (run: wavedl-train --list_models)")
        print()

    print("=" * 40)
    print()


def main() -> int:
    """Main entry point for wavedl-train command."""
    # Fast path for utility flags (avoid accelerate launch overhead)
    exit_code = handle_fast_path_args()
    if exit_code is not None:
        return exit_code

    # Parse arguments
    args, train_args = parse_args()

    # Setup environment (smart detection)
    setup_environment()

    # Check if wavedl package is importable
    try:
        import wavedl  # noqa: F401
    except ImportError:
        print("Error: wavedl package not found. Run: pip install -e .", file=sys.stderr)
        return 1

    # Auto-detect GPUs if not specified
    if args.num_gpus is not None:
        num_gpus = args.num_gpus
        print(f"Using NUM_GPUS={num_gpus} (set via command line)")
    else:
        num_gpus = detect_gpus()

    # Build accelerate launch command
    cmd = [
        "accelerate",
        "launch",
        f"--num_processes={num_gpus}",
        f"--num_machines={args.num_machines}",
        f"--machine_rank={args.machine_rank}",
        f"--mixed_precision={args.mixed_precision}",
        f"--dynamo_backend={args.dynamo_backend}",
    ]

    # Explicitly set multi_gpu to suppress accelerate auto-detection warning
    if num_gpus > 1:
        cmd.append("--multi_gpu")

    # Add multi-node networking args if specified (required for some clusters)
    if args.main_process_ip:
        cmd.append(f"--main_process_ip={args.main_process_ip}")
    if args.main_process_port:
        cmd.append(f"--main_process_port={args.main_process_port}")

    cmd += ["-m", "wavedl.train"] + train_args

    # Create output directory if specified
    for i, arg in enumerate(train_args):
        if arg == "--output_dir" and i + 1 < len(train_args):
            Path(train_args[i + 1]).mkdir(parents=True, exist_ok=True)
            break
        if arg.startswith("--output_dir="):
            Path(arg.split("=", 1)[1]).mkdir(parents=True, exist_ok=True)
            break

    # Launch training
    try:
        result = subprocess.run(cmd, check=False)
        exit_code = result.returncode
    except KeyboardInterrupt:
        print("\n\n⚠️  Training interrupted by user")
        exit_code = 130

    # Print summary
    wandb_enabled = "--wandb" in train_args
    print_summary(
        exit_code,
        wandb_enabled,
        os.environ.get("WANDB_MODE", "online"),
        os.environ.get("WANDB_DIR", "/tmp/wandb"),
    )

    return exit_code


if __name__ == "__main__":
    sys.exit(main())
