#!/usr/bin/env python3
"""
RunPod Deployment Script - FIXED VERSION
Scans for best value GPU in EUR-IS region and deploys a pod on SECURE cloud.

KEY FIX: Uses REST API for availability-aware deployment instead of GraphQL global counts.
"""

import os
import sys
import argparse
import requests
from dotenv import load_dotenv

# Load environment variables from .env.runpod
env_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), '.env.runpod')
load_dotenv(env_path)

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')

if not RUNPOD_API_KEY:
    print("ERROR: RUNPOD_API_KEY not found in .env.runpod")
    sys.exit(1)

if not RUNPOD_VOLUME_ID:
    print("ERROR: RUNPOD_VOLUME_ID not found in .env.runpod")
    sys.exit(1)

# EUR-IS datacenters to scan
# CRITICAL FIX: Volume se3zdnb5o4 is in EUR-IS-1 ONLY!
# If pod deploys to EUR-IS-2 or EUR-IS-3, volume will NOT be accessible
EUR_IS_DATACENTERS = ['EUR-IS-1']  # ONLY EUR-IS-1 for volume mounting!

# GPU Compatibility Reference (Documentation Only)
# CRITICAL CHANGE (2025-10-28): Removed CUDA 13+ GPU filtering
#
# Previously filtered H100, L40S, RTX 6000 Ada as "CUDA 13.0+ incompatible"
# because we believed driver 580+ could not run CUDA 12.9 binaries.
#
# CONFIRMED SAFE: NVIDIA driver 580+ is backward compatible with CUDA 12.x
# - Driver 580 natively supports CUDA 12.9 binaries (no forward compat package needed)
# - PTX JIT compilation works correctly (verified via NVIDIA docs)
# - Our CUDA 12.9 binaries work on both driver 550 (CUDA 12.x) and 580+ (CUDA 13.x)
# - /usr/local/cuda/compat path in Dockerfile provides additional forward compat
#
# Sources:
# - NVIDIA CUDA Compatibility Docs (Minor Version Compatibility)
# - Perplexity AI verification (2025-10-28)
# - Medium article: "CUDA Hell" (driver backward compatibility)
#
# Result: ALL GPUs with driver 525+ can run our CUDA 12.9 binaries

# Known compatible GPUs (reference only - not used for filtering)
KNOWN_COMPATIBLE_GPU_TYPES = [
    'RTX A4000',
    'RTX A5000',
    'RTX A6000',
    'Tesla V100',
    'RTX 4090',
    'A100',
    'H100',           # CUDA 13 GPU - backward compatible with CUDA 12.9
    'L40S',           # CUDA 13 GPU - backward compatible with CUDA 12.9
    'RTX 6000 Ada',   # CUDA 13 GPU - backward compatible with CUDA 12.9
]

# Deprecated: No longer used (kept for git history)
INCOMPATIBLE_GPU_TYPES = []

# REST API endpoint (NEW - supports datacenter filtering)
REST_API_URL = "https://rest.runpod.io/v1/pods"

# GraphQL endpoint (still needed for listing GPU types and pricing)
GRAPHQL_ENDPOINT = "https://api.runpod.io/graphql"

# GraphQL query to fetch GPU types with global availability and pricing
GPU_QUERY = """
{
  gpuTypes {
    id
    displayName
    memoryInGb
    secureCloud
    communityCloud
    lowestPrice(input: {gpuCount: 1}) {
      uninterruptablePrice
    }
  }
}
"""

def query_graphql(query, variables=None):
    """Execute GraphQL query against RunPod API."""
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {RUNPOD_API_KEY}"
    }

    payload = {"query": query}
    if variables:
        payload["variables"] = variables

    try:
        response = requests.post(GRAPHQL_ENDPOINT, json=payload, headers=headers, timeout=30)
        response.raise_for_status()
        result = response.json()

        # Check for GraphQL errors
        if 'errors' in result:
            error_msg = result['errors'][0].get('message', 'Unknown error')
            print(f"ERROR: GraphQL errors: {result['errors']}")
            return None

        return result
    except requests.exceptions.RequestException as e:
        print(f"ERROR: Failed to query RunPod GraphQL API: {e}")
        return None


def get_available_gpu_types(allow_cuda13=False):
    """
    Query available GPU types with ≥16GB VRAM that have SOME availability in SECURE cloud.

    NOTE: This returns GLOBAL secure cloud availability, not EUR-IS specific.
    The actual datacenter filtering happens during deployment via REST API.

    Args:
        allow_cuda13: [DEPRECATED] No longer used - all GPUs supported via backward compatibility
    """
    print("  Querying GPU types and pricing...")

    # Deprecation warning if allow_cuda13 was explicitly used
    if allow_cuda13:
        print("    ⚠️  Note: --allow-cuda13 flag is deprecated (all GPUs now supported)")

    data = query_graphql(GPU_QUERY)
    if not data:
        return []

    gpu_types = data.get('data', {}).get('gpuTypes', [])

    # Filter criteria (SIMPLIFIED - no CUDA version filtering):
    # 1. memoryInGb >= 16
    # 2. secureCloud > 0 (available SOMEWHERE in secure cloud - not necessarily EUR-IS)
    # 3. Has pricing information
    #
    # CRITICAL: No longer filtering by CUDA version!
    # Driver 580+ is backward compatible with CUDA 12.9 binaries (verified 2025-10-28)
    available_gpus = []

    for gpu in gpu_types:
        memory = gpu.get('memoryInGb', 0)
        secure_count = gpu.get('secureCloud', 0)
        lowest_price = gpu.get('lowestPrice', {})
        price = lowest_price.get('uninterruptablePrice') if lowest_price else None
        gpu_name = gpu.get('displayName', 'Unknown')

        if memory >= 16 and secure_count > 0 and price is not None:
            available_gpus.append({
                'id': gpu.get('id', ''),
                'name': gpu_name,
                'vram': memory,
                'price': float(price),
                'global_available': secure_count  # This is GLOBAL, not EUR-IS specific
            })

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

    return available_gpus


def sanitize_command(command):
    """
    Sanitize deployment command to ensure it works with Docker ENTRYPOINT chain.
    
    CRITICAL: RunPod's dockerStartCmd sets Docker CMD, NOT ENTRYPOINT.
    The ENTRYPOINT chain (entrypoint-self-terminate.sh → entrypoint-generic.sh) 
    MUST execute first for pod auto-termination to work.
    
    This function:
    1. Detects and strips /bin/bash -c wrappers (which bypass entrypoint)
    2. Removes chmod +x commands (entrypoint-generic.sh handles this)
    3. Removes tee redirection (container logs capture everything)
    4. Returns clean binary path + arguments
    
    Args:
        command: Raw command string from user
        
    Returns:
        Sanitized command string suitable for dockerStartCmd
    """
    import re
    
    if not command:
        return command
        
    # Strip leading/trailing whitespace
    cmd = command.strip()
    
    # Detect /bin/bash -c wrapper
    if cmd.startswith('/bin/bash -c'):
        print("  ⚠️  WARNING: Detected /bin/bash -c wrapper - stripping to preserve entrypoint chain")
        # Extract command from quotes
        match = re.search(r'/bin/bash -c ["\'](.+)["\']', cmd)
        if match:
            cmd = match.group(1)
            print(f"      Extracted: {cmd[:80]}...")
    
    # Remove chmod +x prefix (entrypoint-generic.sh handles this)
    cmd = re.sub(r'chmod \+x [^\s]+ && ', '', cmd)
    
    # Remove tee redirection suffix (container logs capture everything)
    cmd = re.sub(r' 2>&1 \| tee [^\s]+$', '', cmd)
    
    return cmd.strip()


def deploy_pod_rest_api(gpu, image, command, container_disk, datacenters, dry_run=False):
    """
    Deploy a pod using the REST API with datacenter-specific availability filtering.

    This is the KEY FIX: REST API checks EUR-IS datacenter availability at deployment time,
    not relying on global GraphQL counts.
    """
    from datetime import datetime, timedelta

    pod_name = "foxhunt-training"

    # NOTE: Auto-termination is now handled by entrypoint-self-terminate.sh wrapper script
    # inside the Docker container, not via the RunPod API "terminateAfter" field (which causes HTTP 400).
    # The wrapper script terminates the pod after training completes to prevent restart loops.

    # Build REST API request payload
    deployment_payload = {
        "cloudType": "SECURE",  # CRITICAL: Only secure cloud
        "dataCenterIds": datacenters,  # CRITICAL: EUR-IS specific datacenters
        "dataCenterPriority": "availability",  # Try datacenters in order, prefer available
        "gpuTypeIds": [gpu['id']],  # GPU type to deploy
        "gpuTypePriority": "availability",  # Use available GPU
        "gpuCount": 1,
        "name": pod_name,
        "imageName": image,
        "containerDiskInGb": container_disk,
        "volumeInGb": 0,  # Using network volume instead
        "networkVolumeId": RUNPOD_VOLUME_ID,
        "volumeMountPath": "/runpod-volume",  # CRITICAL: WHERE to mount the volume
        "ports": ["8888/http", "22/tcp"],
        "env": {},
        "interruptible": False,  # On-demand (non-spot)
        "minRAMPerGPU": 8,
        "minVCPUPerGPU": 2
    }

    # Sanitize and add Docker command if specified
    # CRITICAL ARCHITECTURE:
    # 1. dockerStartCmd sets Docker CMD (NOT ENTRYPOINT)
    # 2. Docker execution order: ENTRYPOINT args... + CMD args...
    # 3. Our ENTRYPOINT: /entrypoint.sh (→ entrypoint-self-terminate.sh → entrypoint-generic.sh)
    # 4. entrypoint-generic.sh calls: exec "$@" (passes CMD to binary)
    # 5. entrypoint-self-terminate.sh captures exit code and terminates pod on success
    # 
    # MUST AVOID (sanitized automatically):
    # - /bin/bash -c "..." wrappers (override ENTRYPOINT, bypass auto-termination)
    # - chmod commands (entrypoint-generic.sh handles this)
    # - Shell redirections like tee (container logs capture everything)
    #
    # Command is pre-sanitized in main() before reaching here
    if command:
        # Split command into arguments (respects quoted strings)
        import shlex
        deployment_payload["dockerStartCmd"] = shlex.split(command)

    # Add Docker registry authentication if configured
    if RUNPOD_CONTAINER_REGISTRY_AUTH_ID:
        deployment_payload["containerRegistryAuthId"] = RUNPOD_CONTAINER_REGISTRY_AUTH_ID

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

    if dry_run:
        print("\nDRY RUN: Skipping actual deployment")
        print("\nPayload would be:")
        import json
        print(json.dumps(deployment_payload, indent=2))
        return None

    print("\nDeploying pod via REST API (checks EUR-IS availability)...")

    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {RUNPOD_API_KEY}"
    }

    try:
        response = requests.post(REST_API_URL, json=deployment_payload, headers=headers, timeout=60)

        # Debug output
        print(f"  HTTP Status: {response.status_code}")

        # SUCCESS: HTTP 200 (OK) or HTTP 201 (Created)
        if response.status_code in [200, 201]:
            try:
                pod_data = response.json()

                # Validate pod_data structure
                if not isinstance(pod_data, dict):
                    print(f"  ⚠️  Unexpected response format: {type(pod_data)}")
                    print(f"  Raw response: {response.text[:500]}")
                    return None

                # Check if pod was actually created
                if 'id' not in pod_data:
                    print(f"  ⚠️  Pod data missing 'id' field")
                    print(f"  Raw response: {response.text[:500]}")
                    return None

                print(f"  ✅ Pod created successfully! ID: {pod_data['id']}")
                return pod_data
            except ValueError as e:
                print(f"  ⚠️  Failed to parse JSON response: {e}")
                print(f"  Raw response: {response.text[:500]}")
                return None
        elif response.status_code == 400:
            try:
                error_data = response.json()
                error_msg = error_data.get('error', 'Unknown error')
                print(f"  ⚠️  Deployment failed: {error_msg}")

                # Check if it's an availability issue
                if 'not available' in error_msg.lower() or 'no machines' in error_msg.lower():
                    print(f"  💡 GPU not available in EUR-IS datacenters at this time")
            except ValueError:
                print(f"  ⚠️  Deployment failed: {response.text[:200]}")

            return None
        else:
            print(f"  ❌ Unexpected response (status {response.status_code}): {response.text[:200]}")
            return None

    except requests.exceptions.RequestException as e:
        print(f"  ⚠️  Deployment failed: {str(e)[:100]}")
        return None


def display_deployment_result(pod_data):
    """Display deployment results and connection info."""
    print("\n" + "="*70)
    print("✅ POD DEPLOYED SUCCESSFULLY")
    print("="*70)
    print(f"Pod ID:          {pod_data['id']}")

    # Extract GPU info safely
    machine = pod_data.get('machine', {})
    gpu_info = machine.get('gpuType', {}) if machine else {}
    print(f"GPU:             {gpu_info.get('displayName', 'N/A')}")

    print(f"GPU Count:       {pod_data.get('gpu', {}).get('count', 'N/A')}")
    print(f"Cost:            ${pod_data.get('costPerHr', 'N/A')}/hr")

    # Extract datacenter
    datacenter = machine.get('dataCenterId', 'N/A')
    print(f"Datacenter:      {datacenter}")

    print(f"Image:           {pod_data.get('imageName', pod_data.get('image', 'N/A'))}")
    print(f"Container Disk:  {pod_data['containerDiskInGb']}GB")
    print(f"Status:          {pod_data.get('desiredStatus', 'RUNNING')}")
    print("="*70)
    print("\n📝 NEXT STEPS:")
    print("1. Wait 2-3 minutes for pod to initialize")
    print(f"2. Access Jupyter at: https://{pod_data['id']}-8888.proxy.runpod.net")
    print(f"3. SSH access: ssh root@{pod_data['id']}.ssh.runpod.io")
    print("4. Monitor pod: https://www.runpod.io/console/pods")
    print(f"\n⚠️  IMPORTANT: Pod will cost ${pod_data.get('costPerHr', 'N/A')}/hr - remember to stop when done!")
    print()


def main():
    """Main execution function."""
    parser = argparse.ArgumentParser(
        description="Deploy a RunPod GPU pod in EUR-IS region (SECURE cloud) - FIXED VERSION",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Auto-select best value GPU with default training command (DQN 1 epoch smoke test)
  ./runpod_deploy.py

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

  # Custom TFT training command
  ./runpod_deploy.py --command "/runpod-volume/binaries/train_tft_parquet --parquet-file /runpod-volume/test_data/ES_FUT_180d.parquet --epochs 50 --use-gpu"

  # Custom DQN training with different dataset
  ./runpod_deploy.py --command "/runpod-volume/binaries/train_dqn --parquet-file /runpod-volume/test_data/NQ_FUT_180d.parquet --epochs 100 --output-dir /runpod-volume/models"

  # Custom image and command
  ./runpod_deploy.py --image runpod/pytorch:2.1.0 --command "jupyter lab --ip=0.0.0.0 --allow-root"

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

KEY FIXES:
  - Uses REST API for deployment (checks EUR-IS datacenter availability)
  - Properly formats dockerStartCmd as array (RunPod API requirement)
  - Default command: DQN smoke test (1 epoch, small dataset, proven working)
  - Command parameter accepts FULL command including binary path
        """
    )

    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 (default: DQN 1-epoch smoke test)'
    )
    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(
        '--allow-cuda13',
        action='store_true',
        help='[DEPRECATED] No longer needed - all GPUs supported via backward compatibility'
    )

    args = parser.parse_args()

    # Sanitize command to ensure entrypoint chain works
    if args.command:
        args.command = sanitize_command(args.command)

    print("🔍 Querying available GPU types (global secure cloud)...")

    # Query available GPUs (global availability) - no CUDA filtering
    gpus = get_available_gpu_types(allow_cuda13=args.allow_cuda13)

    if not gpus:
        print("\nERROR: No GPUs available with ≥16GB VRAM in SECURE cloud")
        print("\n💡 TIP: This checks global availability.")
        print("        EUR-IS specific availability is checked during deployment via REST API.")
        sys.exit(1)

    print(f"\n✅ Found {len(gpus)} GPU type(s) to try")

    # Create priority list: user preference (if specified), then sorted by price
    priority_gpus = []

    # Add user's preferred GPU if specified
    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:
            print(f"WARNING: Preferred GPU '{args.gpu_type}' not found, trying all GPUs...")

    # Add remaining GPUs sorted by price (cheapest first)
    # gpus list is already sorted by price at line 121
    for gpu in gpus:
        if gpu not in priority_gpus:
            priority_gpus.append(gpu)

    # Try each GPU until one succeeds
    attempted_gpus = []
    pod_data = None

    for gpu in priority_gpus:
        if gpu in attempted_gpus:
            continue

        attempted_gpus.append(gpu)

        print(f"\n🎯 Attempting deployment: {gpu['name']} (${gpu['price']:.3f}/hr)...")
        print(f"   (Global availability: {gpu['global_available']} in secure cloud)")
        print(f"   (Will check EUR-IS specific availability during deployment...)")

        pod_data = deploy_pod_rest_api(
            gpu,
            args.image,
            args.command,
            args.container_disk,
            EUR_IS_DATACENTERS,
            args.dry_run
        )

        if pod_data:
            break
        else:
            print(f"❌ {gpu['name']} not available in EUR-IS, trying next option...")

    # Display results
    if pod_data:
        display_deployment_result(pod_data)
    else:
        print("\n❌ ERROR: Failed to deploy pod on any available GPU in EUR-IS")
        print("\nAttempted GPUs:")
        for gpu in attempted_gpus:
            print(f"  - {gpu['name']} ({gpu['vram']}GB VRAM) @ ${gpu['price']:.3f}/hr")
        print("\n💡 TIP: RunPod availability changes frequently. Try again in a few minutes.")
        print("         The script now correctly checks EUR-IS datacenter availability,")
        print("         not just global secure cloud counts.")
        sys.exit(1)

if __name__ == "__main__":
    main()