#!/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!

# 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():
    """
    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.
    """
    print("  Querying GPU types and pricing...")

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

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

    # Filter criteria:
    # 1. memoryInGb >= 16
    # 2. secureCloud > 0 (available SOMEWHERE in secure cloud - not necessarily EUR-IS)
    # 3. Has pricing information
    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

        if memory >= 16 and secure_count > 0 and price is not None:
            available_gpus.append({
                'id': gpu.get('id', ''),
                'name': gpu.get('displayName', 'Unknown'),
                '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 global secure cloud availability")
    else:
        print(f"    ⚠️  No GPU types found with ≥16GB VRAM in secure cloud")

    return available_gpus


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.
    """
    pod_name = "foxhunt-training"

    # Build REST API request payload
    # CRITICAL: Only use fields documented in official RunPod REST API
    # Reference: https://docs.runpod.io/api-reference/pods/POST/pods
    deployment_payload = {
        "cloudType": "SECURE",  # CRITICAL: Only secure cloud
        "computeType": "GPU",  # GPU pod (required field)
        "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
    }

    # Add Docker command if specified
    # RunPod expects dockerStartCmd as an array of strings (shell arguments)
    # Split the command string into proper arguments
    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"Registry Auth:   {RUNPOD_CONTAINER_REGISTRY_AUTH_ID}")
    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
  ./runpod_deploy.py

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

  # Custom training command (overrides Dockerfile CMD)
  ./runpod_deploy.py --command "--parquet-file /runpod-volume/test_data/ES_FUT_180d.parquet --epochs 100"

  # 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: TFT training with Parquet data (overrides Dockerfile CMD)
        """
    )

    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='--parquet-file /runpod-volume/test_data/ES_FUT_180d.parquet --epochs 50 --learning-rate 0.001',
        help='Training command arguments (default: TFT training with Parquet data)'
    )
    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'
    )

    args = parser.parse_args()

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

    # Query available GPUs (global availability)
    gpus = get_available_gpu_types()

    if not gpus:
        print("\nERROR: No GPUs available with ≥16GB VRAM in SECURE cloud")
        print("\n💡 TIP: This checks global availability. EUR-IS specific availability")
        print("        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()