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

NEW: Integrates with foxhunt_runpod module for enhanced features:
- RunPodClient for pod management
- S3LogMonitor for real-time log streaming
- Auto-termination support
"""

import os
import sys
import argparse
from pathlib import Path
from dotenv import load_dotenv

# Check for .venv activation
is_venv = hasattr(sys, 'real_prefix') or (hasattr(sys, 'base_prefix') and sys.base_prefix != sys.prefix)
if not is_venv:
    print("=" * 70)
    print("ERROR: Not running in a virtual environment (.venv)")
    print("=" * 70)
    print("This script requires .venv activation to ensure correct dependencies.")
    print()
    print("To fix this:")
    print("  1. Activate .venv:   source .venv/bin/activate")
    print("  2. Install deps:     pip install -r runpod/requirements.txt")
    print("  3. Re-run script:    python3 scripts/runpod_deploy.py --help")
    print("=" * 70)
    sys.exit(1)

try:
    # Import from runpod module (clean architecture - no sys.path hacks)
    from runpod import RunPodClient, PodMonitor, S3Client
    from runpod.s3_monitor import S3LogMonitor
    import requests  # Required for legacy GraphQL queries
    USE_NEW_MODULE = True
except ImportError as e:
    print(f"ERROR: Could not import runpod module: {e}")
    print()
    print("To fix this:")
    print("  1. Ensure .venv is activated: source .venv/bin/activate")
    print("  2. Install dependencies: pip install -r runpod/requirements.txt")
    print("  3. Check module location: ls runpod/")
    print()
    sys.exit(1)

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

# S3 credentials (optional - for log monitoring)
RUNPOD_S3_BUCKET = os.getenv('RUNPOD_S3_BUCKET')
RUNPOD_S3_ACCESS_KEY = os.getenv('RUNPOD_S3_ACCESS_KEY')
RUNPOD_S3_SECRET_KEY = os.getenv('RUNPOD_S3_SECRET_KEY')

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

    client = RunPodClient(
        api_key=RUNPOD_API_KEY,
        volume_id=RUNPOD_VOLUME_ID,
        registry_auth_id=RUNPOD_CONTAINER_REGISTRY_AUTH_ID
    )

    gpus = client.get_available_gpus(min_vram=16)

    if gpus:
        print(f"    ✅ Found {len(gpus)} GPU type(s) with global secure cloud availability")
    else:
        print(f"    ⚠️  No GPU types found with ≥16GB VRAM in secure cloud")

    return gpus


def get_available_gpu_types_legacy_unused():
    """
    Query available GPU types with ≥16GB VRAM that have SOME availability in SECURE cloud (LEGACY).

    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 (legacy mode)...")

    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(gpu, image, command, container_disk, dry_run=False):
    """Deploy a pod using RunPodClient."""
    client = RunPodClient(
        api_key=RUNPOD_API_KEY,
        volume_id=RUNPOD_VOLUME_ID,
        registry_auth_id=RUNPOD_CONTAINER_REGISTRY_AUTH_ID
    )

    print("\n" + "="*70)
    print("DEPLOYMENT PLAN")
    print("="*70)
    print(f"Pod Name:        foxhunt-training")
    print(f"GPU:             {gpu['name']} ({gpu['vram']}GB VRAM)")
    print(f"Datacenters:     EUR-IS-1 (volume location)")
    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")
        return None

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

    pod_data = client.deploy_pod(
        gpu_id=gpu['id'],
        image=image,
        command=command,
        container_disk=container_disk,
        dry_run=False
    )

    if pod_data:
        print(f"  ✅ Pod created successfully! ID: {pod_data['id']}")

    return pod_data


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

    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 (legacy mode)")
    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
  python3 scripts/runpod_deploy.py

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

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

  # Enable S3 log monitoring (streams logs in real-time)
  python3 scripts/runpod_deploy.py --monitor --timeout 120m

  # Full automation: deploy, monitor, auto-terminate
  python3 scripts/runpod_deploy.py --gpu-type "RTX A4000" --monitor --auto-stop --timeout 2h

  # Custom image and command
  python3 scripts/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)
  python3 scripts/runpod_deploy.py --dry-run

NEW FEATURES:
  - RunPodClient integration for cleaner pod management
  - S3LogMonitor for real-time training log streaming (--monitor)
  - Auto-termination when training completes (--auto-stop)
  - Configurable monitoring timeout and interval
  - Backward compatible with legacy implementation

REQUIREMENTS:
  - Python 3.8+, preferably in .venv (source .venv/bin/activate)
  - Dependencies: dotenv, requests, boto3 (for S3 monitoring)
  - .env.runpod with RUNPOD_API_KEY, RUNPOD_VOLUME_ID
  - Optional: RUNPOD_S3_* credentials for log monitoring
        """
    )

    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'
    )
    parser.add_argument(
        '--monitor',
        action='store_true',
        help='Enable S3 log monitoring (streams training logs in real-time)'
    )
    parser.add_argument(
        '--auto-stop',
        action='store_true',
        help='Auto-terminate pod when training completes (requires --monitor)'
    )
    parser.add_argument(
        '--timeout',
        default='120m',
        help='Maximum monitoring time (e.g., 30m, 2h). Default: 120m'
    )
    parser.add_argument(
        '--monitor-interval',
        type=int,
        default=10,
        help='S3 log polling interval in seconds (default: 10)'
    )

    args = parser.parse_args()

    # Validate --auto-stop requires --monitor
    if args.auto_stop and not args.monitor:
        print("ERROR: --auto-stop requires --monitor to be enabled")
        sys.exit(1)

    # Parse timeout to seconds
    timeout_seconds = None
    if args.timeout:
        import re
        match = re.match(r'(\d+)([mh]?)$', args.timeout.lower())
        if match:
            value, unit = match.groups()
            value = int(value)
            if unit == 'h':
                timeout_seconds = value * 3600
            else:  # Default to minutes
                timeout_seconds = value * 60
        else:
            print(f"ERROR: Invalid timeout format '{args.timeout}'. Use format like '30m' or '2h'")
            sys.exit(1)

    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(
            gpu,
            args.image,
            args.command,
            args.container_disk,
            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)

        # Enable monitoring if requested
        if args.monitor and USE_NEW_MODULE:
            # Check S3 credentials
            if not all([RUNPOD_S3_BUCKET, RUNPOD_S3_ACCESS_KEY, RUNPOD_S3_SECRET_KEY]):
                print("\n⚠️  WARNING: S3 credentials not configured in .env.runpod")
                print("           Cannot enable log monitoring")
                print("           Required: RUNPOD_S3_BUCKET, RUNPOD_S3_ACCESS_KEY, RUNPOD_S3_SECRET_KEY")
            else:
                print("\n" + "="*70)
                print("S3 LOG MONITORING ENABLED")
                print("="*70)
                print(f"Bucket:          {RUNPOD_S3_BUCKET}")
                print(f"Pod ID:          {pod_data['id']}")
                print(f"Interval:        {args.monitor_interval}s")
                print(f"Timeout:         {args.timeout}")
                if args.auto_stop:
                    print(f"Auto-stop:       Enabled (auto-terminate on completion)")
                print("="*70)
                print("\nStarting log monitoring... (Press Ctrl+C to stop)\n")

                try:
                    # Use PodMonitor for integrated monitoring and auto-termination
                    from runpod.config import RunPodConfig

                    # Create config with credentials
                    config = RunPodConfig(
                        api_key=RUNPOD_API_KEY,
                        volume_id=RUNPOD_VOLUME_ID,
                        s3_bucket=RUNPOD_S3_BUCKET,
                        s3_access_key=RUNPOD_S3_ACCESS_KEY,
                        s3_secret_key=RUNPOD_S3_SECRET_KEY,
                        log_poll_interval=args.monitor_interval
                    )

                    monitor = PodMonitor(
                        pod_id=pod_data['id'],
                        config=config
                    )

                    # Stream logs with auto-termination if requested
                    if args.auto_stop:
                        # Use auto_terminate which streams logs and terminates on completion
                        success = monitor.auto_terminate(wait_for_completion=True)

                        if success:
                            print("\n" + "="*70)
                            print("AUTO-TERMINATION COMPLETE")
                            print("="*70)
                            print(f"  ✅ Pod {pod_data['id']} terminated successfully")
                            print(f"  💰 Final cost estimate: ~${pod_data.get('costPerHr', 0) * (timeout_seconds or 7200) / 3600:.2f}")
                            print("="*70)
                        else:
                            print("\n⚠️  Auto-termination failed or training incomplete")
                            print(f"    Please terminate manually: https://www.runpod.io/console/pods")
                    else:
                        # Just stream logs without auto-termination
                        monitor.stream_s3_logs(
                            follow=True,
                            poll_interval=args.monitor_interval
                        )

                except Exception as e:
                    print(f"\n⚠️  Monitoring error: {e}")
                    print("     Pod is still running. Remember to stop it manually!")

        elif args.monitor and not USE_NEW_MODULE:
            print("\n⚠️  WARNING: Log monitoring requires runpod module")
            print("           Install dependencies: pip install -r runpod/requirements.txt")

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