#!/usr/bin/env python3
"""
RunPod Deployment Script - Production Version with RunPod SDK
Scans for best value GPU in EUR-IS region and deploys a pod on SECURE or COMMUNITY cloud.

ARCHITECTURE:
  - Uses RunPod Python SDK (not CLI subprocess calls)
  - Automatic venv setup (.venv) with runpod SDK
  - REST API for availability-aware deployment
  - Automatic binary upload with SHA256 validation
  - Volume mount architecture (zero downloads at runtime)
  - CI/CD friendly (exit codes, no interactive prompts)

BINARY UPLOAD ARCHITECTURE:
  - Calculates SHA256 hash of local binaries
  - Compares with S3 metadata to detect changes
  - Only uploads if local binary is newer or different
  - Supports multiple binaries (train_*, hyperopt_*)
"""

import os
import sys
import argparse
import hashlib
import subprocess
import logging
from datetime import datetime
from pathlib import Path

# Configure logging with timestamps
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger(__name__)

# Get script directory
SCRIPT_DIR = Path(__file__).parent.absolute()
PROJECT_ROOT = SCRIPT_DIR.parent
VENV_DIR = SCRIPT_DIR / '.venv'
VENV_PYTHON = VENV_DIR / 'bin' / 'python3'
VENV_PIP = VENV_DIR / 'bin' / 'pip'

def setup_venv():
    """
    Auto-setup Python venv if it doesn't exist.
    This ensures the script is self-contained and works in CI/CD.
    """
    if VENV_DIR.exists() and VENV_PYTHON.exists():
        logger.info("✓ Virtual environment already exists")
        return True

    logger.info("Setting up Python virtual environment...")
    try:
        # Create venv
        subprocess.run(
            [sys.executable, '-m', 'venv', str(VENV_DIR)],
            check=True,
            capture_output=True,
            cwd=SCRIPT_DIR
        )
        logger.info("✓ Created .venv")

        # Upgrade pip
        subprocess.run(
            [str(VENV_PIP), 'install', '--upgrade', 'pip'],
            check=True,
            capture_output=True
        )
        logger.info("✓ Upgraded pip")

        # Install requirements
        requirements_file = SCRIPT_DIR / 'requirements.txt'
        if requirements_file.exists():
            subprocess.run(
                [str(VENV_PIP), 'install', '-r', str(requirements_file)],
                check=True,
                capture_output=True
            )
            logger.info("✓ Installed dependencies from requirements.txt")
        else:
            # Fallback: install core packages
            subprocess.run(
                [str(VENV_PIP), 'install', 'runpod', 'boto3', 'requests', 'python-dotenv'],
                check=True,
                capture_output=True
            )
            logger.info("✓ Installed core dependencies")

        logger.info("✓ Virtual environment setup complete")
        return True

    except subprocess.CalledProcessError as e:
        logger.error(f"✗ Failed to setup virtual environment: {e}")
        return False
    except Exception as e:
        logger.error(f"✗ Unexpected error during venv setup: {e}")
        return False

def restart_with_venv():
    """Restart the script using the venv Python interpreter."""
    if sys.executable == str(VENV_PYTHON):
        # Already running in venv
        return

    logger.info("Restarting script with venv Python...")
    os.execv(str(VENV_PYTHON), [str(VENV_PYTHON)] + sys.argv)

# Auto-setup venv on import
if __name__ == '__main__':
    # Check if we're running in venv
    if sys.executable != str(VENV_PYTHON):
        # Not in venv - setup and restart
        if not VENV_DIR.exists():
            if not setup_venv():
                logger.error("✗ Failed to setup venv - exiting")
                sys.exit(1)
        restart_with_venv()

# Now import runpod SDK (after venv is active)
try:
    import runpod
    from dotenv import load_dotenv
    import requests
except ImportError as e:
    logger.error(f"✗ Missing required package: {e}")
    logger.error("  Run: pip install runpod boto3 requests python-dotenv")
    sys.exit(1)

# Load environment variables from .env.runpod
env_path = PROJECT_ROOT / '.env.runpod'
if env_path.exists():
    load_dotenv(env_path)
    logger.info(f"✓ Loaded environment from {env_path}")
else:
    logger.warning(f"⚠ Environment file not found: {env_path}")

# Configuration from environment
RUNPOD_API_KEY = os.getenv('RUNPOD_API_KEY')
RUNPOD_VOLUME_ID = os.getenv('RUNPOD_VOLUME_ID')
RUNPOD_CONTAINER_REGISTRY_AUTH_ID = os.getenv('RUNPOD_CONTAINER_REGISTRY_AUTH_ID')

# S3 configuration for binary uploads (Runpod S3 endpoint)
S3_BUCKET = 'se3zdnb5o4'
S3_ENDPOINT = 'https://s3api-eur-is-1.runpod.io'
S3_PROFILE = 'runpod'

# EUR-IS datacenters to scan
# CRITICAL: Volume se3zdnb5o4 is in EUR-IS-1 ONLY!
EUR_IS_DATACENTERS = ['EUR-IS-1']

# REST API endpoint
REST_API_URL = "https://rest.runpod.io/v1/pods"

def validate_environment():
    """Validate required environment variables."""
    if not RUNPOD_API_KEY:
        logger.error("✗ RUNPOD_API_KEY not found in environment")
        logger.error("  Set it in .env.runpod or export RUNPOD_API_KEY=...")
        return False

    if not RUNPOD_VOLUME_ID:
        logger.error("✗ RUNPOD_VOLUME_ID not found in environment")
        logger.error("  Set it in .env.runpod or export RUNPOD_VOLUME_ID=...")
        return False

    logger.info("✓ Environment variables validated")
    return True

def get_binary_hash(binary_path):
    """
    Calculate SHA256 hash of local binary.

    Args:
        binary_path: Path to binary file

    Returns:
        SHA256 hash as hex string
    """
    sha256 = hashlib.sha256()
    try:
        with open(binary_path, 'rb') as f:
            for chunk in iter(lambda: f.read(65536), b''):
                sha256.update(chunk)
        return sha256.hexdigest()
    except FileNotFoundError:
        logger.warning(f"⚠ Binary not found: {binary_path}")
        return None
    except Exception as e:
        logger.error(f"✗ Failed to hash binary: {e}")
        return None

def check_s3_binary_exists(binary_name):
    """
    Check if binary exists in S3 and get its metadata.

    Args:
        binary_name: Name of binary in S3 (e.g., 'hyperopt_mamba2_demo')

    Returns:
        Dict with 'exists', 'etag', 'last_modified' or None on error
    """
    s3_key = f"binaries/current/{binary_name}"
    cmd = [
        'aws', 's3api', 'head-object',
        '--bucket', S3_BUCKET,
        '--key', s3_key,
        '--profile', S3_PROFILE,
        '--endpoint-url', S3_ENDPOINT
    ]

    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=10)
        if result.returncode == 0:
            import json
            metadata = json.loads(result.stdout)
            return {
                'exists': True,
                'etag': metadata.get('ETag', '').strip('"'),
                'last_modified': metadata.get('LastModified', ''),
                'size': metadata.get('ContentLength', 0)
            }
        elif result.returncode == 254 or 'Not Found' in result.stderr:
            return {'exists': False}
        else:
            logger.warning(f"⚠ S3 head-object failed: {result.stderr[:100]}")
            return {'exists': False}
    except subprocess.TimeoutExpired:
        logger.warning("⚠ S3 head-object timeout")
        return {'exists': False}
    except Exception as e:
        logger.error(f"✗ S3 check failed: {e}")
        return {'exists': False}

def upload_binary_to_s3(binary_path, binary_name):
    """
    Upload binary to S3 with timestamped filename and standard symlink.

    Args:
        binary_path: Local path to binary
        binary_name: Name for S3 object

    Returns:
        True on success, False on failure
    """
    local_hash = get_binary_hash(binary_path)

    if not local_hash:
        return False

    # Get binary modification time as timestamp
    try:
        mtime = os.path.getmtime(binary_path)
        build_timestamp = datetime.fromtimestamp(mtime).isoformat()
    except Exception as e:
        logger.warning(f"⚠ Failed to get binary timestamp: {e}")
        build_timestamp = datetime.now().isoformat()

    # Generate timestamp for filename (YYYYMMDD_HHMMSS)
    timestamp_suffix = datetime.now().strftime("%Y%m%d_%H%M%S")
    timestamped_name = f"{binary_name}_{timestamp_suffix}"

    # S3 keys
    timestamped_key = f"binaries/timestamped/{timestamped_name}"
    standard_key = f"binaries/current/{binary_name}"
    flat_key = f"binaries/{binary_name}"  # Backward compatibility

    # Metadata for all objects
    metadata = f"sha256={local_hash},uploaded={datetime.now().isoformat()},build_timestamp={build_timestamp}"

    logger.info(f"📤 Uploading {binary_name} to S3...")
    logger.info(f"   Timestamped version: {timestamped_name}")
    logger.info(f"   Build timestamp: {build_timestamp}")

    # Step 1: Upload to timestamped path
    timestamped_cmd = [
        'aws', 's3', 'cp',
        binary_path,
        f"s3://{S3_BUCKET}/{timestamped_key}",
        '--profile', S3_PROFILE,
        '--endpoint-url', S3_ENDPOINT,
        '--metadata', metadata
    ]

    try:
        result = subprocess.run(timestamped_cmd, capture_output=True, text=True, timeout=120)

        if result.returncode != 0:
            logger.error(f"✗ Timestamped upload failed: {result.stderr[:200]}")
            return False

        logger.info(f"✓ Timestamped binary uploaded: {timestamped_name}")

    except subprocess.TimeoutExpired:
        logger.error("✗ Upload timeout (binary too large?)")
        return False
    except Exception as e:
        logger.error(f"✗ Upload error: {e}")
        return False

    # Step 2: Copy to standard path (binaries/current/{name})
    logger.info(f"🔗 Creating standard path link: binaries/current/{binary_name}")

    standard_cmd = [
        'aws', 's3', 'cp',
        f"s3://{S3_BUCKET}/{timestamped_key}",
        f"s3://{S3_BUCKET}/{standard_key}",
        '--profile', S3_PROFILE,
        '--endpoint-url', S3_ENDPOINT,
        '--metadata', metadata,
        '--metadata-directive', 'REPLACE'
    ]

    try:
        result = subprocess.run(standard_cmd, capture_output=True, text=True, timeout=120)

        if result.returncode != 0:
            logger.warning(f"⚠ Standard path copy failed: {result.stderr[:200]}")
            logger.info(f"   Timestamped version still available: {timestamped_name}")
            return True

        logger.info(f"✓ Standard path updated: binaries/current/{binary_name} → {timestamped_name}")

    except Exception as e:
        logger.warning(f"⚠ Standard path copy error: {e}")
        logger.info(f"   Timestamped version still available: {timestamped_name}")
        return True

    # Step 3: Copy to flat path (binaries/{name}) for backward compatibility
    logger.info(f"🔗 Creating flat path: binaries/{binary_name}")

    flat_cmd = [
        'aws', 's3', 'cp',
        f"s3://{S3_BUCKET}/{timestamped_key}",
        f"s3://{S3_BUCKET}/{flat_key}",
        '--profile', S3_PROFILE,
        '--endpoint-url', S3_ENDPOINT,
        '--metadata', metadata,
        '--metadata-directive', 'REPLACE'
    ]

    try:
        result = subprocess.run(flat_cmd, capture_output=True, text=True, timeout=120)

        if result.returncode != 0:
            logger.warning(f"⚠ Flat path copy failed: {result.stderr[:200]}")
            return True

        logger.info(f"✓ Flat path updated: binaries/{binary_name}")
        logger.info(f"   SHA256: {local_hash[:16]}...")
        return True

    except Exception as e:
        logger.warning(f"⚠ Flat path copy error: {e}")
        return True

def check_and_upload_binary(binary_path, binary_name, force_upload=False):
    """
    Check if binary needs upload and upload if necessary.

    Args:
        binary_path: Local path to binary
        binary_name: Name for S3 object
        force_upload: Skip version check and force upload

    Returns:
        True if binary is available in S3, False on error
    """
    if not os.path.exists(binary_path):
        logger.warning(f"⚠ Binary not found locally: {binary_path}")
        logger.info("   Checking if it exists in S3...")
        s3_info = check_s3_binary_exists(binary_name)
        if s3_info.get('exists'):
            logger.info("✓ Binary exists in S3 (uploaded previously)")
            return True
        else:
            logger.error("✗ Binary not found in S3 either - build it first!")
            return False

    # Calculate local hash
    local_hash = get_binary_hash(binary_path)
    if not local_hash:
        return False

    # Get local file info
    local_stat = os.stat(binary_path)
    local_size = local_stat.st_size
    local_mtime = datetime.fromtimestamp(local_stat.st_mtime)

    logger.info(f"📦 Local binary: {binary_name}")
    logger.info(f"   Size: {local_size / (1024*1024):.1f}MB")
    logger.info(f"   Modified: {local_mtime.strftime('%Y-%m-%d %H:%M:%S')}")
    logger.info(f"   SHA256: {local_hash[:16]}...")

    if force_upload:
        logger.info("🔄 Force upload requested")
        return upload_binary_to_s3(binary_path, binary_name)

    # Check S3 version
    logger.info("🔍 Checking S3 version...")
    s3_info = check_s3_binary_exists(binary_name)

    if not s3_info.get('exists'):
        logger.info("📤 Binary not in S3 - uploading...")
        return upload_binary_to_s3(binary_path, binary_name)

    # Compare metadata
    s3_size = s3_info.get('size', 0)

    logger.info(f"📦 S3 binary: {binary_name}")
    logger.info(f"   Size: {s3_size / (1024*1024):.1f}MB")

    # If sizes differ, upload new version
    if local_size != s3_size:
        logger.info(f"🔄 Size mismatch - uploading new version...")
        return upload_binary_to_s3(binary_path, binary_name)

    logger.info("✓ Binary up-to-date in S3")
    return True

def sanitize_command(command):
    """
    Sanitize deployment command to ensure it works with Docker ENTRYPOINT chain.

    Args:
        command: Raw command string from user

    Returns:
        Sanitized command string
    """
    import re

    if not command:
        return command

    cmd = command.strip()

    # Detect /bin/bash -c wrapper
    if cmd.startswith('/bin/bash -c'):
        logger.warning("⚠ Detected /bin/bash -c wrapper - stripping to preserve entrypoint")
        match = re.search(r'/bin/bash -c ["\'](.+)["\']', cmd)
        if match:
            cmd = match.group(1)
            logger.info(f"   Extracted: {cmd[:80]}...")

    # Remove chmod +x prefix
    cmd = re.sub(r'chmod \+x [^\s]+ && ', '', cmd)

    # Remove tee redirection suffix
    cmd = re.sub(r' 2>&1 \| tee [^\s]+$', '', cmd)

    return cmd.strip()

def get_available_gpus():
    """
    Query available GPU types using RunPod SDK.

    Returns:
        List of GPU dicts with id, name, vram, price
    """
    logger.info("🔍 Querying available GPU types...")

    try:
        # Initialize RunPod with API key
        runpod.api_key = RUNPOD_API_KEY

        # Get GPU types
        gpu_types = runpod.get_gpus()

        available_gpus = []
        for gpu in gpu_types:
            # Filter: >= 16GB VRAM, has secure OR community cloud availability
            memory_gb = gpu.get('memoryInGb', 0)
            secure_count = gpu.get('secureCloud', 0)
            community_count = gpu.get('communityCloud', 0)
            total_count = secure_count + community_count
            lowest_price = gpu.get('lowestPrice', {})
            price = lowest_price.get('uninterruptablePrice') if lowest_price else None
            gpu_name = gpu.get('displayName', 'Unknown')

            if memory_gb < 16:
                logger.info(f"   ⏭ {gpu_name}: Skipped (VRAM {memory_gb}GB < 16GB)")
                continue
            if total_count == 0:
                logger.info(f"   ⏭ {gpu_name}: Skipped (0 instances available)")
                continue
            if price is None:
                logger.info(f"   ⏭ {gpu_name}: Skipped (no pricing info)")
                continue

            cloud_type = []
            if secure_count > 0:
                cloud_type.append(f"secure:{secure_count}")
            if community_count > 0:
                cloud_type.append(f"community:{community_count}")
            cloud_info = ", ".join(cloud_type)

            available_gpus.append({
                'id': gpu.get('id', ''),
                'name': gpu_name,
                'vram': memory_gb,
                'price': float(price),
                'global_available': total_count
            })
            logger.info(f"   ✓ {gpu_name}: Available ({memory_gb}GB VRAM, ${price:.3f}/hr, {cloud_info})")

        if available_gpus:
            # Sort by price (cheapest first)
            available_gpus.sort(key=lambda x: x['price'])
            logger.info(f"✓ Found {len(available_gpus)} GPU type(s) with ≥16GB VRAM")
        else:
            logger.warning("⚠ No GPU types found with ≥16GB VRAM")

        return available_gpus

    except Exception as e:
        logger.error(f"✗ Failed to query GPU types: {e}")
        return []

def deploy_pod(gpu, image, command, container_disk, dry_run=False):
    """
    Deploy a pod using RunPod SDK.
    Tries SECURE cloud first, falls back to COMMUNITY cloud if needed.

    Args:
        gpu: GPU dict from get_available_gpus()
        image: Docker image name
        command: Command to run
        container_disk: Container disk size in GB
        dry_run: If True, don't actually deploy

    Returns:
        Pod dict on success, None on failure
    """
    pod_name = "foxhunt-training"

    logger.info("="*70)
    logger.info("DEPLOYMENT PLAN")
    logger.info("="*70)
    logger.info(f"Pod Name:        {pod_name}")
    logger.info(f"GPU:             {gpu['name']} ({gpu['vram']}GB VRAM)")
    logger.info(f"Datacenters:     {', '.join(EUR_IS_DATACENTERS)}")
    logger.info(f"Price:           ${gpu['price']:.3f}/hr (estimate)")
    logger.info(f"Docker Image:    {image}")
    logger.info(f"Container Disk:  {container_disk}GB")
    logger.info(f"Network Volume:  {RUNPOD_VOLUME_ID} → /runpod-volume")
    logger.info(f"Ports:           8888/http (Jupyter), 22/tcp (SSH)")
    logger.info(f"Auto-Terminate:  entrypoint-self-terminate.sh")
    if command:
        logger.info(f"Command:         {command}")
    logger.info("="*70)

    if dry_run:
        logger.info("\n✓ DRY RUN: Skipping actual deployment")
        return None

    # Initialize RunPod
    runpod.api_key = RUNPOD_API_KEY

    # Prepare base deployment config
    base_config = {
        "name": pod_name,
        "gpu_type_id": gpu['id'],
        "gpu_count": 1,
        "image_name": image,
        "container_disk_in_gb": container_disk,
        "volume_in_gb": 0,
        "network_volume_id": RUNPOD_VOLUME_ID,
        "volume_mount_path": "/runpod-volume",
        "ports": "8888/http,22/tcp",
        "min_vcpu_count": 2,
        "min_memory_in_gb": 8,
        "data_center_id": EUR_IS_DATACENTERS[0],  # EUR-IS-1
    }

    # Add Docker command if specified
    if command:
        import shlex
        base_config["docker_args"] = shlex.split(command)

    # Add registry auth if configured
    if RUNPOD_CONTAINER_REGISTRY_AUTH_ID:
        base_config["container_registry_auth_id"] = RUNPOD_CONTAINER_REGISTRY_AUTH_ID

    # Try SECURE cloud first, then COMMUNITY cloud
    for cloud_type in ["SECURE", "COMMUNITY"]:
        logger.info(f"\n🚀 Trying {cloud_type} cloud...")

        deploy_config = base_config.copy()
        deploy_config["cloud_type"] = cloud_type

        try:
            pod = runpod.create_pod(**deploy_config)

            if pod and 'id' in pod:
                logger.info(f"✓ Pod created successfully on {cloud_type} cloud! ID: {pod['id']}")
                return pod
            else:
                logger.warning(f"⚠ {cloud_type} cloud returned no pod ID - trying next")

        except Exception as e:
            logger.warning(f"⚠ {cloud_type} cloud deployment failed: {e}")
            if cloud_type == "COMMUNITY":
                # Last attempt failed
                logger.error("✗ All cloud types exhausted")
                return None

    return None

def terminate_pod(pod_id):
    """
    Terminate a specific pod by ID.

    Args:
        pod_id: Pod ID to terminate

    Returns:
        True on success, False on failure
    """
    logger.info("="*70)
    logger.info("POD TERMINATION")
    logger.info("="*70)
    logger.info(f"Pod ID: {pod_id}")
    logger.info("="*70)

    try:
        runpod.api_key = RUNPOD_API_KEY
        result = runpod.terminate_pod(pod_id)

        if result:
            logger.info("✓ Pod termination request sent successfully")
            logger.info("  It may take 30-60 seconds for the pod to fully stop")
            logger.info("  Check status at: https://www.runpod.io/console/pods")
            return True
        else:
            logger.warning("⚠ Pod not found or already terminated")
            return False

    except Exception as e:
        logger.error(f"✗ Termination failed: {e}")
        return False

def display_pod_info(pod):
    """Display pod deployment results."""
    logger.info("")
    logger.info("="*70)
    logger.info("✓ POD DEPLOYED SUCCESSFULLY")
    logger.info("="*70)
    logger.info(f"Pod ID:          {pod['id']}")
    logger.info(f"Cost:            ${pod.get('costPerHr', 'N/A')}/hr")
    logger.info(f"Status:          {pod.get('desiredStatus', 'RUNNING')}")
    logger.info("="*70)
    logger.info("\n📝 NEXT STEPS:")
    logger.info("1. Wait 2-3 minutes for pod to initialize")
    logger.info(f"2. Access Jupyter at: https://{pod['id']}-8888.proxy.runpod.net")
    logger.info(f"3. SSH access: ssh root@{pod['id']}.ssh.runpod.io")
    logger.info("4. Monitor pod: https://www.runpod.io/console/pods")
    logger.info(f"\n⚠ IMPORTANT: Pod will cost ${pod.get('costPerHr', 'N/A')}/hr")
    logger.info("")

def main():
    """Main execution function."""
    parser = argparse.ArgumentParser(
        description="Deploy a RunPod GPU pod in EUR-IS region (SECURE or COMMUNITY cloud)",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Deploy with PSO-fixed binary (25 trials validation)
  ./runpod_deploy.py --command "/runpod-volume/binaries/hyperopt_mamba2_demo_20251029_094049 --parquet-file /runpod-volume/test_data/ES_FUT_180d.parquet --base-dir /runpod-volume/ml_training --trials 25 --epochs 1 --batch-size-max 256 --seed 42" --skip-upload

  # Deploy with specific GPU
  ./runpod_deploy.py --gpu-type "RTX 4090"

  # Dry run (show plan without deploying)
  ./runpod_deploy.py --dry-run

  # Terminate a specific pod
  ./runpod_deploy.py --terminate POD_ID
        """
    )

    parser.add_argument(
        '--gpu-type',
        help='Preferred GPU type (e.g., "RTX 4090"). Auto-selects best value if not specified.'
    )
    parser.add_argument(
        '--image',
        default='jgrusewski/foxhunt:latest',
        help='Docker image to use (default: jgrusewski/foxhunt:latest)'
    )
    parser.add_argument(
        '--command',
        default='/runpod-volume/binaries/train_dqn --parquet-file /runpod-volume/test_data/ES_FUT_small.parquet --epochs 1 --output-dir /runpod-volume/models',
        help='Full training command including binary path'
    )
    parser.add_argument(
        '--container-disk',
        type=int,
        default=50,
        help='Container disk size in GB (default: 50)'
    )
    parser.add_argument(
        '--dry-run',
        action='store_true',
        help='Show deployment plan without actually deploying'
    )
    parser.add_argument(
        '--skip-upload',
        action='store_true',
        help='Skip automatic binary upload check'
    )
    parser.add_argument(
        '--force-upload',
        action='store_true',
        help='Force re-upload binaries'
    )
    parser.add_argument(
        '--terminate',
        metavar='POD_ID',
        help='Terminate a specific pod by ID'
    )

    args = parser.parse_args()

    # Handle termination
    if args.terminate:
        if not validate_environment():
            sys.exit(1)
        success = terminate_pod(args.terminate)
        sys.exit(0 if success else 1)

    # Validate environment
    if not validate_environment():
        sys.exit(1)

    # Sanitize command
    if args.command:
        args.command = sanitize_command(args.command)

    # STEP 1: Binary upload check
    if not args.skip_upload:
        logger.info("")
        logger.info("="*70)
        logger.info("STEP 1: BINARY VERSION CHECK")
        logger.info("="*70)

        # Detect binaries from command
        binaries_to_check = []

        if args.command:
            import shlex

            # Check if wrapper script
            is_wrapper = (
                args.command.strip().startswith('aws ') or
                '&&' in args.command or
                's3://' in args.command
            )

            if is_wrapper:
                logger.info("ℹ Detected wrapper script - skipping binary validation")
            else:
                cmd_parts = shlex.split(args.command)
                if cmd_parts:
                    binary_name = os.path.basename(cmd_parts[0])
                    if binary_name != 'jupyter':
                        local_binary_path = PROJECT_ROOT / 'target' / 'release' / 'examples' / binary_name
                        binaries_to_check.append((str(local_binary_path), binary_name))

        if not binaries_to_check:
            logger.info("ℹ No binaries detected - skipping upload check")
        else:
            logger.info(f"🔍 Detected {len(binaries_to_check)} binary(ies) to check")

            all_ok = True
            for local_path, name in binaries_to_check:
                logger.info(f"\n📦 Checking: {name}")
                result = check_and_upload_binary(local_path, name, args.force_upload)
                if not result:
                    logger.error(f"✗ Binary check failed: {name}")
                    all_ok = False

            if not all_ok:
                logger.error("\n✗ Some binaries failed upload check")
                logger.error("   Build with: cargo build --release --examples")
                sys.exit(1)

            logger.info("\n✓ All binaries ready in S3")

        logger.info("="*70)
    else:
        logger.info("\nℹ Skipping binary upload check (--skip-upload)")

    # STEP 2: Query GPUs
    logger.info("")
    logger.info("="*70)
    logger.info("STEP 2: GPU AVAILABILITY SCAN")
    logger.info("="*70)

    gpus = get_available_gpus()

    if not gpus:
        logger.error("\n✗ No GPUs available with ≥16GB VRAM in SECURE or COMMUNITY cloud")
        sys.exit(1)

    logger.info(f"\n✓ Found {len(gpus)} GPU type(s) to try")
    logger.info("="*70)

    # Prioritize user-preferred GPU
    priority_gpus = []
    if args.gpu_type:
        for gpu in gpus:
            if args.gpu_type.lower() in gpu['name'].lower():
                priority_gpus.append(gpu)
                break
        if not priority_gpus:
            logger.warning(f"⚠ Preferred GPU '{args.gpu_type}' not found - trying all GPUs")

    # Add remaining GPUs (already sorted by price)
    for gpu in gpus:
        if gpu not in priority_gpus:
            priority_gpus.append(gpu)

    # STEP 3: Deploy
    logger.info("")
    logger.info("="*70)
    logger.info("STEP 3: POD DEPLOYMENT")
    logger.info("="*70)

    pod = None
    for gpu in priority_gpus:
        logger.info(f"\n🎯 Attempting: {gpu['name']} (${gpu['price']:.3f}/hr)")

        pod = deploy_pod(gpu, args.image, args.command, args.container_disk, args.dry_run)

        if pod:
            break
        else:
            logger.warning(f"✗ {gpu['name']} not available - trying next")

    # Display results
    if pod:
        display_pod_info(pod)
        sys.exit(0)
    else:
        logger.error("\n✗ Failed to deploy pod on any available GPU")
        logger.info("\n💡 TIP: RunPod availability changes frequently. Try again in a few minutes.")
        sys.exit(1)

if __name__ == "__main__":
    main()