#!/usr/bin/env python3
"""
RunPod Pod Creation Test
Tests pod creation with incremental feature additions.

IMPORTANT: Run test_runpod_auth.py first to verify authentication.
"""

import os
import sys
import requests
import json
import time
from typing import Optional, Dict, Any


def get_api_key() -> Optional[str]:
    """Get RunPod API key from environment"""
    api_key = os.environ.get('RUNPOD_API_KEY')
    if not api_key:
        print("❌ ERROR: RUNPOD_API_KEY environment variable not set")
        return None
    return api_key


def create_minimal_pod(api_key: str) -> Optional[Dict[str, Any]]:
    """
    Test 1: Create pod with absolute minimum required fields

    This tests the basic pod creation without any optional features.
    """
    print("\n" + "="*60)
    print("TEST 1: Minimal Pod Creation")
    print("="*60)

    url = f'https://api.runpod.io/graphql?api_key={api_key}'

    headers = {
        'Content-Type': 'application/json'
    }

    # Absolute minimum fields based on documentation
    mutation = """
    mutation {
      podFindAndDeployOnDemand(
        input: {
          cloudType: SECURE
          gpuTypeId: "NVIDIA RTX A4000"
          gpuCount: 1
          name: "foxhunt-test-minimal"
          imageName: "runpod/pytorch:2.0.1-py3.10-cuda11.8.0-devel"
        }
      ) {
        id
        name
        imageName
        desiredStatus
        machineId
      }
    }
    """

    payload = {'query': mutation}

    print("\n📋 Configuration:")
    print("  Cloud Type: SECURE")
    print("  GPU: NVIDIA RTX A4000 (1x)")
    print("  Image: runpod/pytorch:2.0.1-py3.10-cuda11.8.0-devel")
    print("  Name: foxhunt-test-minimal")
    print("\n📡 Creating minimal pod...")

    try:
        response = requests.post(url, json=payload, headers=headers, timeout=30)

        print(f"📥 Response Status: {response.status_code}")

        if response.status_code != 200:
            print(f"❌ HTTP Error: {response.status_code}")
            print(f"Response: {response.text}")
            return None

        data = response.json()

        if 'errors' in data:
            print(f"❌ GraphQL Errors: {json.dumps(data['errors'], indent=2)}")
            return None

        if 'data' in data and data['data'].get('podFindAndDeployOnDemand'):
            pod = data['data']['podFindAndDeployOnDemand']
            print(f"\n✅ Pod created successfully!")
            print(f"🆔 Pod ID: {pod['id']}")
            print(f"📛 Name: {pod['name']}")
            print(f"🖼️  Image: {pod['imageName']}")
            print(f"📊 Status: {pod['desiredStatus']}")
            print(f"🖥️  Machine ID: {pod.get('machineId', 'N/A')}")
            return pod
        else:
            print(f"❌ Pod creation returned null (likely no capacity)")
            print(f"Response: {json.dumps(data, indent=2)}")
            return None

    except requests.exceptions.Timeout:
        print("❌ Request timed out after 30 seconds")
        return None
    except requests.exceptions.RequestException as e:
        print(f"❌ Request failed: {e}")
        return None
    except Exception as e:
        print(f"❌ Unexpected error: {e}")
        return None


def create_pod_with_storage(api_key: str) -> Optional[Dict[str, Any]]:
    """
    Test 2: Add storage configuration

    Adds network volume mounting to the minimal pod.
    """
    print("\n" + "="*60)
    print("TEST 2: Pod Creation with Network Volume")
    print("="*60)

    url = f'https://api.runpod.io/graphql?api_key={api_key}'

    headers = {
        'Content-Type': 'application/json'
    }

    mutation = """
    mutation {
      podFindAndDeployOnDemand(
        input: {
          cloudType: SECURE
          gpuTypeId: "NVIDIA RTX A4000"
          gpuCount: 1
          name: "foxhunt-test-storage"
          imageName: "runpod/pytorch:2.0.1-py3.10-cuda11.8.0-devel"

          networkVolumeId: "se3zdnb5o4"
          volumeMountPath: "/workspace"
          containerDiskInGb: 50
        }
      ) {
        id
        name
        imageName
        desiredStatus
        machineId
      }
    }
    """

    payload = {'query': mutation}

    print("\n📋 Configuration:")
    print("  Previous + Network Volume: se3zdnb5o4")
    print("  Mount Path: /workspace")
    print("  Container Disk: 50GB")
    print("\n📡 Creating pod with storage...")

    try:
        response = requests.post(url, json=payload, headers=headers, timeout=30)

        print(f"📥 Response Status: {response.status_code}")

        if response.status_code != 200:
            print(f"❌ HTTP Error: {response.status_code}")
            print(f"Response: {response.text}")
            return None

        data = response.json()

        if 'errors' in data:
            print(f"❌ GraphQL Errors: {json.dumps(data['errors'], indent=2)}")
            return None

        if 'data' in data and data['data'].get('podFindAndDeployOnDemand'):
            pod = data['data']['podFindAndDeployOnDemand']
            print(f"\n✅ Pod with storage created successfully!")
            print(f"🆔 Pod ID: {pod['id']}")
            print(f"📛 Name: {pod['name']}")
            return pod
        else:
            print(f"❌ Pod creation returned null")
            return None

    except Exception as e:
        print(f"❌ Error: {e}")
        return None


def create_pod_with_env(api_key: str) -> Optional[Dict[str, Any]]:
    """
    Test 3: Add environment variables

    Adds environment variable configuration.
    """
    print("\n" + "="*60)
    print("TEST 3: Pod Creation with Environment Variables")
    print("="*60)

    url = f'https://api.runpod.io/graphql?api_key={api_key}'

    headers = {
        'Content-Type': 'application/json'
    }

    mutation = """
    mutation {
      podFindAndDeployOnDemand(
        input: {
          cloudType: SECURE
          gpuTypeId: "NVIDIA RTX A4000"
          gpuCount: 1
          name: "foxhunt-test-env"
          imageName: "runpod/pytorch:2.0.1-py3.10-cuda11.8.0-devel"

          networkVolumeId: "se3zdnb5o4"
          volumeMountPath: "/workspace"
          containerDiskInGb: 50

          env: [
            { key: "TEST_VAR", value: "test_value" },
            { key: "RUST_LOG", value: "info" }
          ]
        }
      ) {
        id
        name
        imageName
        env
        desiredStatus
      }
    }
    """

    payload = {'query': mutation}

    print("\n📋 Configuration:")
    print("  Previous + Environment Variables:")
    print("    TEST_VAR=test_value")
    print("    RUST_LOG=info")
    print("\n📡 Creating pod with environment...")

    try:
        response = requests.post(url, json=payload, headers=headers, timeout=30)

        print(f"📥 Response Status: {response.status_code}")

        if response.status_code != 200:
            print(f"❌ HTTP Error: {response.status_code}")
            return None

        data = response.json()

        if 'errors' in data:
            print(f"❌ GraphQL Errors: {json.dumps(data['errors'], indent=2)}")
            return None

        if 'data' in data and data['data'].get('podFindAndDeployOnDemand'):
            pod = data['data']['podFindAndDeployOnDemand']
            print(f"\n✅ Pod with environment created successfully!")
            print(f"🆔 Pod ID: {pod['id']}")
            print(f"📛 Name: {pod['name']}")
            if pod.get('env'):
                print(f"🔧 Environment: {pod['env']}")
            return pod
        else:
            print(f"❌ Pod creation returned null")
            return None

    except Exception as e:
        print(f"❌ Error: {e}")
        return None


def create_pod_with_custom_image(api_key: str) -> Optional[Dict[str, Any]]:
    """
    Test 4: Use custom Docker image

    Tests with the actual Foxhunt image.
    """
    print("\n" + "="*60)
    print("TEST 4: Pod Creation with Custom Image")
    print("="*60)

    url = f'https://api.runpod.io/graphql?api_key={api_key}'

    headers = {
        'Content-Type': 'application/json'
    }

    mutation = """
    mutation {
      podFindAndDeployOnDemand(
        input: {
          cloudType: SECURE
          gpuTypeId: "NVIDIA RTX A4000"
          gpuCount: 1
          name: "foxhunt-test-custom"
          imageName: "jgrusewski/foxhunt:latest"

          networkVolumeId: "se3zdnb5o4"
          volumeMountPath: "/workspace"
          containerDiskInGb: 50

          minVcpuCount: 4
          minMemoryInGb: 16

          env: [
            { key: "RUST_LOG", value: "info" }
          ]
        }
      ) {
        id
        name
        imageName
        desiredStatus
      }
    }
    """

    payload = {'query': mutation}

    print("\n📋 Configuration:")
    print("  Custom Image: jgrusewski/foxhunt:latest")
    print("  Resources: 4 vCPU, 16GB RAM")
    print("  Storage: Network volume + 50GB container disk")
    print("\n📡 Creating pod with custom image...")

    try:
        response = requests.post(url, json=payload, headers=headers, timeout=30)

        print(f"📥 Response Status: {response.status_code}")

        if response.status_code != 200:
            print(f"❌ HTTP Error: {response.status_code}")
            return None

        data = response.json()

        if 'errors' in data:
            print(f"❌ GraphQL Errors: {json.dumps(data['errors'], indent=2)}")
            return None

        if 'data' in data and data['data'].get('podFindAndDeployOnDemand'):
            pod = data['data']['podFindAndDeployOnDemand']
            print(f"\n✅ Pod with custom image created!")
            print(f"🆔 Pod ID: {pod['id']}")
            print(f"📛 Name: {pod['name']}")
            print(f"🖼️  Image: {pod['imageName']}")
            return pod
        else:
            print(f"❌ Pod creation returned null")
            return None

    except Exception as e:
        print(f"❌ Error: {e}")
        return None


def stop_pod(api_key: str, pod_id: str) -> bool:
    """Stop a running pod"""
    url = f'https://api.runpod.io/graphql?api_key={api_key}'

    headers = {
        'Content-Type': 'application/json'
    }

    mutation = f"""
    mutation {{
      podStop(input: {{podId: "{pod_id}"}}) {{
        id
        desiredStatus
      }}
    }}
    """

    payload = {'query': mutation}

    print(f"\n🛑 Stopping pod {pod_id}...")

    try:
        response = requests.post(url, json=payload, headers=headers, timeout=10)

        if response.status_code != 200:
            print(f"❌ Failed to stop pod: HTTP {response.status_code}")
            return False

        data = response.json()

        if 'errors' in data:
            print(f"❌ Errors stopping pod: {data['errors']}")
            return False

        print(f"✅ Pod stopped successfully")
        return True

    except Exception as e:
        print(f"❌ Error stopping pod: {e}")
        return False


def main():
    """Run pod creation tests incrementally"""
    print("="*60)
    print("🧪 RunPod Pod Creation Test Suite")
    print("="*60)
    print("\nThis script tests pod creation with incremental features.")
    print("Each test builds on the previous one.")
    print("\n⚠️  WARNING: These tests will CREATE REAL PODS that incur costs!")
    print("Make sure to stop/terminate pods after testing.")

    # Get API key
    api_key = get_api_key()
    if not api_key:
        sys.exit(1)

    # Ask for confirmation
    print("\n" + "="*60)
    response = input("Continue with pod creation tests? (yes/no): ")
    if response.lower() != 'yes':
        print("Tests cancelled.")
        return 0

    # Track created pods for cleanup
    created_pods = []

    # Run tests
    print("\n🚀 Starting pod creation tests...\n")

    # Test 1: Minimal pod
    pod = create_minimal_pod(api_key)
    if pod:
        created_pods.append(pod['id'])
        time.sleep(2)  # Brief pause between tests
    else:
        print("\n❌ Minimal pod creation failed. Stopping tests.")
        return 1

    # Test 2: With storage
    pod = create_pod_with_storage(api_key)
    if pod:
        created_pods.append(pod['id'])
        time.sleep(2)
    else:
        print("\n⚠️  Storage test failed, but continuing...")

    # Test 3: With environment
    pod = create_pod_with_env(api_key)
    if pod:
        created_pods.append(pod['id'])
        time.sleep(2)
    else:
        print("\n⚠️  Environment test failed, but continuing...")

    # Test 4: With custom image
    pod = create_pod_with_custom_image(api_key)
    if pod:
        created_pods.append(pod['id'])
    else:
        print("\n⚠️  Custom image test failed")

    # Summary
    print("\n" + "="*60)
    print("📊 TEST SUMMARY")
    print("="*60)
    print(f"\n✅ Created {len(created_pods)} pods:")
    for pod_id in created_pods:
        print(f"  - {pod_id}")

    # Cleanup option
    print("\n" + "="*60)
    cleanup = input("\nStop all created pods? (yes/no): ")
    if cleanup.lower() == 'yes':
        print("\n🧹 Cleaning up...")
        for pod_id in created_pods:
            stop_pod(api_key, pod_id)

    print("\n✅ Tests complete!")
    print("\n💡 Next step: Run the full deployment script")
    print("   ./scripts/deploy_runpod_training.py")

    return 0


if __name__ == '__main__':
    sys.exit(main())