#!/usr/bin/env python3
"""Cloud GPU deployment launcher for FSI_Edge training.

Supports:
  - RunPod (Serverless + Secure Cloud)
  - Vast.ai
  - Lambda Labs
  - Manual mode (generates setup script only)

Usage:
  python cloud_launch.py --provider runpod --preset production --model-size 800M
  python cloud_launch.py --provider vast --preset dev --gpus 8
  python cloud_launch.py --provider lambda --preset quick
  python cloud_launch.py --provider manual --output-dir ./deploy  # generates setup script
"""
import os
import sys
import json
import argparse
import subprocess
import tempfile
import shutil
from pathlib import Path
from datetime import datetime

REPO_ROOT = Path(__file__).resolve().parent.parent

PRESETS = {
    "quick": {
        "model_size": "360M",
        "stages": "stage1,stage1b",
        "batch_size": 4,
        "max_steps": 5000,
        "fp16": True,
        "gpu_type": "RTX 4090",
        "min_vram_gb": 24,
        "num_gpus": 1,
        "data_samples": 10000,
    },
    "dev": {
        "model_size": "800M",
        "stages": "stage1,stage1b,stage2,stage2b",
        "batch_size": 8,
        "max_steps": 50000,
        "fp16": True,
        "gpu_type": "A100-80GB",
        "min_vram_gb": 80,
        "num_gpus": 4,
        "data_samples": 100000,
    },
    "production": {
        "model_size": "800M",
        "stages": "stage0,stage1,stage1b,stage2,stage2b,stage3,stage3b,stage4,stage5",
        "batch_size": 32,
        "max_steps": 500000,
        "fp16": True,
        "gpu_type": "H100-80GB",
        "min_vram_gb": 80,
        "num_gpus": 8,
        "data_samples": 5000000,
    },
    "production_fast": {
        "model_size": "800M",
        "stages": "stage0,stage1,stage1b,stage2,stage2b,stage3,stage3b,stage4",
        "batch_size": 64,
        "max_steps": 200000,
        "fp16": True,
        "gpu_type": "H100-80GB",
        "min_vram_gb": 80,
        "num_gpus": 8,
        "data_samples": 2000000,
    },
}

PROVIDER_CONFIGS = {
    "runpod": {
        "template_id": "runpod-pytorch:2.2.0-cuda12.1.0",
        "container_disk_gb": 50,
        "supported_gpus": ["RTX 4090", "A100-80GB", "A100-SXM-80GB", "H100-80GB", "H100-SXM-80GB"],
    },
    "vast": {
        "supported_gpus": ["RTX 4090", "A100-80GB", "A100-SXM-80GB", "H100-80GB", "H100-SXM-80GB"],
        "disk_gb": 100,
    },
    "lambda": {
        "supported_gpu_types": ["1x A100", "4x A100", "8x A100", "8x H100"],
        "base_image": "nvidia/cuda:12.1.0-runtime-ubuntu22.04",
    },
}

SETUP_SCRIPT = """#!/bin/bash
set -e
echo "=== FSI_Edge Cloud Instance Setup ==="

export DEBIAN_FRONTEND=noninteractive
export CUDA_HOME=/usr/local/cuda
export PATH=$CUDA_HOME/bin:$PATH
export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH
export NCCL_DEBUG=INFO
export NCCL_IB_DISABLE=0
export NCCL_IB_HCA=$(ibstatus 2>/dev/null | grep -oP 'mlx5_\\d+' | head -1 || echo "")
export TORCH_NCCL_ASYNC_ERROR_HANDLING=1

apt-get update && apt-get install -y --no-install-recommends \\
    git curl wget build-essential python3-dev python3-pip \\
    libopenmpi-dev openmpi-bin \\
    libnccl-dev libnccl2 nccl-tools \\
    ibverbs-providers libibverbs-dev \\
    ninja-build

pip install --upgrade pip setuptools wheel

if ! nvidia-smi &>/dev/null; then
    echo "ERROR: No NVIDIA GPU detected!"
    exit 1
fi
echo "GPU: $(nvidia-smi --query-gpu=name --format=csv,noheader | head -1)"
echo "VRAM: $(nvidia-smi --query-gpu=memory.total --format=csv,noheader | head -1)"
echo "CUDA: $(nvcc --version 2>/dev/null | tail -1 || nvidia-smi | grep 'CUDA Version' || echo 'unknown')"

echo "=== Installing FSI_Edge dependencies ==="
cd /workspace/FSI_Edge

pip install torch==2.2.0 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
pip install -r requirements.txt
pip install flash-attn --no-build-isolation
pip install deepspeed mpi4py

echo "=== Generating training data ==="
python data/prepare_data.py --samples {data_samples}

echo "=== Setup complete ==="
echo "To launch training:"
echo "  python run_pipeline.py --preset {preset} --model-size {model_size} --stages {stages} [additional args]"
echo ""
echo "For multi-node (if applicable):"
echo '  torchrun --nnodes=$NNODES --nproc-per-node=$NPROC_PER_NODE --rdzv-endpoint=$MASTER_ADDR run_pipeline.py ...'
"""

TRAIN_LAUNCH_TEMPLATE = """#!/bin/bash
set -e
echo "=== FSI_Edge Training Launch ==="

cd /workspace/FSI_Edge

NNODES=${{NNODES:-1}}
NPROC_PER_NODE=${{NPROC_PER_NODE:-{num_gpus}}}
MASTER_ADDR=${{MASTER_ADDR:-localhost}}
MASTER_PORT=${{MASTER_PORT:-29500}}
WORLD_SIZE=$((NNODES * NPROC_PER_NODE))

EXTRA_ARGS="${{@}}"

if [ $NNODES -eq 1 ] && [ $NPROC_PER_NODE -eq 1 ]; then
    CMD="python run_pipeline.py --preset {preset} --model-size {model_size} --stages \\"{stages}\\" --batch-size {batch_size} --max-steps {max_steps} --output-dir /workspace/FSI_Edge/output/{run_id}"
else
    CMD="torchrun --nnodes=$NNODES --nproc-per-node=$NPROC_PER_NODE --rdzv-endpoint=$MASTER_ADDR:$MASTER_PORT --rdzv-backend=c10d --max-restarts=3 run_pipeline.py --preset {preset} --model-size {model_size} --stages \\"{stages}\\" --batch-size {batch_size} --max-steps {max_steps} --output-dir /workspace/FSI_Edge/output/{run_id}"
fi

if [ -n "$WANDB_API_KEY" ]; then
    CMD="$CMD --wandb-project fsi_edge"
fi
if [ -n "$HF_TOKEN" ]; then
    CMD="$CMD --repo-id fsi_edge/fsi_edge-{model_size} --hf-token $HF_TOKEN"
fi

echo "Running: $CMD $EXTRA_ARGS"
eval "$CMD $EXTRA_ARGS"
"""


def generate_setup_script(preset_name, preset_config):
    data_samples = preset_config.get("data_samples", 500000)
    return SETUP_SCRIPT.format(
        data_samples=data_samples,
        preset=preset_name,
        model_size=preset_config["model_size"],
        stages=preset_config["stages"],
    )


def generate_launch_script(preset_name, preset_config, run_id=None):
    run_id = run_id or datetime.now().strftime("%Y%m%d_%H%M%S")
    return TRAIN_LAUNCH_TEMPLATE.format(
        preset=preset_name,
        model_size=preset_config["model_size"],
        stages=preset_config["stages"],
        batch_size=preset_config["batch_size"],
        max_steps=preset_config["max_steps"],
        num_gpus=preset_config.get("num_gpus", 1),
        run_id=run_id,
    )


def deploy_manual(preset_name, preset_config, output_dir):
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    setup_script = output_dir / "setup.sh"
    setup_script.write_text(generate_setup_script(preset_name, preset_config))
    setup_script.chmod(0o755)

    launch_script = output_dir / "launch.sh"
    launch_script.write_text(generate_launch_script(preset_name, preset_config))
    launch_script.chmod(0o755)

    readme = output_dir / "README.md"
    readme.write_text(f"""# FSI_Edge Cloud Deployment — {preset_name}

## Upload to GPU Instance

```bash
# Copy FSI_Edge to instance
rsync -avz --exclude='__pycache__' --exclude='.git' /FSI_Edge/ user@instance:/workspace/FSI_Edge/

# SSH in and run setup
ssh user@instance
cd /workspace/FSI_Edge
bash {output_dir.name}/setup.sh

# Launch training
bash {output_dir.name}/launch.sh
```

## Multi-Node (if applicable)

```bash
# On master:
NNODES=2 NPROC_PER_NODE={preset_config['num_gpus']} bash {output_dir.name}/launch.sh

# On worker (same command, auto-discovers via env):
NNODES=2 NPROC_PER_NODE={preset_config['num_gpus']} MASTER_ADDR=<master_ip> bash {output_dir.name}/launch.sh
```

## Config

| Setting | Value |
|---------|-------|
| Model | {preset_config['model_size']} |
| Stages | {preset_config['stages']} |
| Batch Size | {preset_config['batch_size']} |
| Max Steps | {preset_config['max_steps']} |
| GPUs | {preset_config['num_gpus']} |
| Data Samples | {preset_config['data_samples']} |
""")

    # Create tarball
    shutil.make_archive(
        str(output_dir / f"fsi_edge_{preset_name}_deploy"),
        "gztar",
        REPO_ROOT,
    )

    print(f"Deployment package created in {output_dir}/")
    print(f"  setup.sh        — Instance setup script")
    print(f"  launch.sh       — Training launch script")
    print(f"  README.md       — Deployment instructions")
    print(f"  fsi_edge_{preset_name}_deploy.tar.gz — Full source tarball")
    return output_dir


def deploy_runpod(preset_name, preset_config, api_key=None, pod_id=None):
    """Generate RunPod deployment config or launch via API."""
    config = {
        "name": f"fsi_edge_{preset_name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}",
        "imageName": "nvidia/cuda:12.1.0-runtime-ubuntu22.04",
        "containerDiskSizeGb": 50,
        "volumeInGb": 0,
        "ports": "22/tcp,8888/tcp,29500/tcp",
        "gpuTypeName": preset_config["gpu_type"],
        "gpuCount": preset_config["num_gpus"],
        "env": [
            {"key": "WANDB_API_KEY", "value": os.environ.get("WANDB_API_KEY", "")},
            {"key": "HF_TOKEN", "value": os.environ.get("HF_TOKEN", "")},
        ],
    }

    config_path = Path.cwd() / f"runpod_{preset_name}.json"
    config_path.write_text(json.dumps(config, indent=2))
    print(f"RunPod config written: {config_path}")
    print()
    print("To deploy via RunPod CLI:")
    print(f"  # Install: pip install runpod")
    print(f"  # Deploy: runpodctl pod create --config {config_path.name}")
    print()
    print("Or use the RunPod web UI at https://www.runpod.io")
    print("1. Click 'Deploy' → 'Secure Cloud'")
    print("2. Select template: RunPod PyTorch 2.2")
    print(f"3. GPU: {preset_config['gpu_type']} x {preset_config['num_gpus']}")
    print(f"4. Container Disk: 50GB")
    print("5. Launch and SSH in, then:")
    print(f"   - git clone https://github.com/YOUR_ORG/FSI_Edge.git /workspace/FSI_Edge")
    print(f"   - cd /workspace/FSI_Edge && bash scripts/setup.sh")
    print(f"   - bash scripts/launch.sh")

    if api_key:
        try:
            import runpod
            runpod.api_key = api_key
            print("API-based deployment coming soon — use config file for now.")
        except ImportError:
            print("Install 'runpod' for API-based deployment: pip install runpod")

    return config_path


def deploy_vast(preset_name, preset_config, api_key=None):
    """Generate Vast.ai deployment config."""
    config = {
        "image": "nvidia/cuda:12.1.0-cudnn8-runtime-ubuntu22.04",
        "env": {
            "WANDB_API_KEY": os.environ.get("WANDB_API_KEY", ""),
            "HF_TOKEN": os.environ.get("HF_TOKEN", ""),
        },
        "disk": 100,
        "gpu_ids": [preset_config["gpu_type"]],
        "num_gpus": preset_config["num_gpus"],
        "ssh_host": "0.0.0.0",
        "ssh_port": 42001,
        "label": f"fsi_edge_{preset_name}",
    }

    config_path = Path.cwd() / f"vast_{preset_name}.json"
    config_path.write_text(json.dumps(config, indent=2))
    print(f"Vast.ai config written: {config_path}")
    print()
    print("To deploy on Vast.ai:")
    print(f"  pip install vastai")
    print(f"  vastai create instance {config_path.name}")
    print()
    print(f"Or browse: https://cloud.vast.ai/?gpu_ids={preset_config['gpu_type'].replace(' ', '+')}")
    print(f"Then SSH in and run setup + launch scripts.")

    if api_key:
        try:
            # Attempt API-based instance creation
            result = subprocess.run(
                ["vastai", "create", "instance", str(config_path)],
                capture_output=True, text=True, timeout=30,
            )
            print(f"Vast.ai API result: {result.stdout}")
            if result.stderr:
                print(f"  stderr: {result.stderr}")
        except Exception as e:
            print(f"  Auto-deploy attempted but failed: {e}")
            print("  Use the config file above to deploy manually.")

    return config_path


def deploy_lambda(preset_name, preset_config, api_key=None):
    """Generate Lambda Labs deployment config."""
    gpu_map = {
        1: "1x A100",
        4: "4x A100",
        8: "8x A100",
    }
    gpu_type = gpu_map.get(preset_config["num_gpus"], "8x A100")

    config = {
        "name": f"fsi_edge_{preset_name}",
        "instance_type": {
            "gpu_type": gpu_type,
            "region": "us-east-1",
        },
        "ssh_key_names": [],
        "file_system_names": [],
        "quantity": 1,
    }

    config_path = Path.cwd() / f"lambda_{preset_name}.json"
    config_path.write_text(json.dumps(config, indent=2))
    print(f"Lambda Labs config written: {config_path}")
    print()
    print("To deploy on Lambda Labs:")
    print(f"  pip install lambda-cloud")
    print(f"  lambda-cloud launch instance {config_path.name}")
    print()
    print("Or use the web UI: https://cloud.lambdalabs.com/instances")
    print(f"Select: {gpu_type}")
    print("Then SSH in and run setup + launch scripts.")

    return config_path


def main():
    parser = argparse.ArgumentParser(
        description="FSI_Edge Cloud GPU Deployment Launcher",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python scripts/cloud_launch.py --provider runpod --preset production --model-size 800M
  python scripts/cloud_launch.py --provider vast --preset dev --gpus 4
  python scripts/cloud_launch.py --provider manual --output-dir ./deploy
        """,
    )
    parser.add_argument("--provider", type=str, default="manual",
                        choices=["runpod", "vast", "lambda", "manual"],
                        help="Cloud provider to deploy to")
    parser.add_argument("--preset", type=str, default="dev",
                        choices=list(PRESETS.keys()),
                        help="Training preset configuration")
    parser.add_argument("--model-size", type=str, default=None,
                        choices=["360M", "800M", "1.5B"],
                        help="Override model size")
    parser.add_argument("--stages", type=str, default=None,
                        help="Override stages (e.g. 'stage1,stage2,stage3')")
    parser.add_argument("--gpus", type=int, default=None,
                        help="Override number of GPUs")
    parser.add_argument("--batch-size", type=int, default=None,
                        help="Override batch size")
    parser.add_argument("--max-steps", type=int, default=None,
                        help="Override max training steps")
    parser.add_argument("--output-dir", type=str, default="./deploy",
                        help="Output directory for deployment files (manual mode)")
    parser.add_argument("--api-key", type=str, default=None,
                        help="Provider API key (optional, for automated deployment)")
    parser.add_argument("--data-samples", type=int, default=None,
                        help="Override number of synthetic data samples")

    args = parser.parse_args()

    # Build config from preset + overrides
    config = dict(PRESETS[args.preset])
    if args.model_size:
        config["model_size"] = args.model_size
    if args.stages:
        config["stages"] = args.stages
    if args.gpus:
        config["num_gpus"] = args.gpus
    if args.batch_size:
        config["batch_size"] = args.batch_size
    if args.max_steps:
        config["max_steps"] = args.max_steps
    if args.data_samples:
        config["data_samples"] = args.data_samples

    print(f"FSI_Edge Cloud Deployment")
    print(f"  Provider: {args.provider}")
    print(f"  Preset:   {args.preset}")
    print(f"  Model:    {config['model_size']}")
    print(f"  Stages:   {config['stages']}")
    print(f"  GPUs:     {config['num_gpus']}")
    print(f"  Batch:    {config['batch_size']}")
    print(f"  Steps:    {config['max_steps']}")
    print()

    if args.provider == "manual":
        deploy_manual(args.preset, config, args.output_dir)
    elif args.provider == "runpod":
        deploy_runpod(args.preset, config, args.api_key)
    elif args.provider == "vast":
        deploy_vast(args.preset, config, args.api_key)
    elif args.provider == "lambda":
        deploy_lambda(args.preset, config, args.api_key)

    print(f"\nDeployment package generated. Copy to your GPU instance and run setup.sh -> launch.sh")


if __name__ == "__main__":
    main()