#!/usr/bin/env python3
"""
RunPod API Authentication Test
Tests that API key authentication works correctly before attempting pod creation.
"""

import os
import sys
import requests
import json
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")
        print("Set it with: export RUNPOD_API_KEY='your_api_key'")
        return None
    return api_key


def test_authentication(api_key: str) -> bool:
    """Test API authentication with simple GPU types query"""
    print("\n" + "="*60)
    print("TEST 1: Basic Authentication")
    print("="*60)

    # Correct URL format per documentation
    url = f'https://api.runpod.io/graphql?api_key={api_key}'

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

    # Simple query that doesn't create resources
    query = """
    query {
      gpuTypes {
        id
        displayName
        memoryInGb
      }
    }
    """

    payload = {'query': query}

    print(f"\n📡 Sending request to: {url[:50]}...")
    print(f"📝 Query: Get all GPU types")

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

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

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

        data = response.json()

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

        # Check for successful data
        if 'data' in data and 'gpuTypes' in data['data']:
            gpu_types = data['data']['gpuTypes']
            print(f"\n✅ Authentication successful!")
            print(f"📊 Found {len(gpu_types)} GPU types")

            # Show first 5 GPU types
            print("\n🎮 Sample GPU Types:")
            for gpu in gpu_types[:5]:
                print(f"  - {gpu['displayName']} ({gpu['memoryInGb']}GB)")

            if len(gpu_types) > 5:
                print(f"  ... and {len(gpu_types) - 5} more")

            return True
        else:
            print(f"❌ Unexpected response structure: {json.dumps(data, indent=2)}")
            return False

    except requests.exceptions.Timeout:
        print("❌ Request timed out after 10 seconds")
        return False
    except requests.exceptions.RequestException as e:
        print(f"❌ Request failed: {e}")
        return False
    except json.JSONDecodeError as e:
        print(f"❌ Failed to parse JSON response: {e}")
        print(f"Raw response: {response.text[:500]}")
        return False


def test_user_info(api_key: str) -> bool:
    """Test getting user information (requires authentication)"""
    print("\n" + "="*60)
    print("TEST 2: User Information (Authenticated Query)")
    print("="*60)

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

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

    # Query that requires authentication
    query = """
    query {
      myself {
        id
        email
      }
    }
    """

    payload = {'query': query}

    print(f"\n📡 Sending authenticated request...")
    print(f"📝 Query: Get user information")

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

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

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

        data = response.json()

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

        # Check for successful data
        if 'data' in data and 'myself' in data['data']:
            user = data['data']['myself']
            print(f"\n✅ User information retrieved!")
            print(f"👤 User ID: {user['id']}")
            print(f"📧 Email: {user['email']}")
            return True
        else:
            print(f"❌ Unexpected response structure: {json.dumps(data, indent=2)}")
            return False

    except requests.exceptions.Timeout:
        print("❌ Request timed out after 10 seconds")
        return False
    except requests.exceptions.RequestException as e:
        print(f"❌ Request failed: {e}")
        return False
    except json.JSONDecodeError as e:
        print(f"❌ Failed to parse JSON response: {e}")
        print(f"Raw response: {response.text[:500]}")
        return False


def test_gpu_availability(api_key: str, gpu_type: str = "NVIDIA RTX A4000") -> bool:
    """Test checking GPU availability and pricing"""
    print("\n" + "="*60)
    print(f"TEST 3: GPU Availability Check ({gpu_type})")
    print("="*60)

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

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

    # Query to check specific GPU availability
    query = f"""
    query {{
      gpuTypes(input: {{ id: "{gpu_type}" }}) {{
        id
        displayName
        memoryInGb
        secureCloud
        communityCloud
        lowestPrice(input: {{
          gpuCount: 1
          secureCloud: true
        }}) {{
          minimumBidPrice
          uninterruptablePrice
          stockStatus
        }}
      }}
    }}
    """

    payload = {'query': query}

    print(f"\n📡 Checking availability for {gpu_type}...")

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

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

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

        data = response.json()

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

        if 'data' in data and 'gpuTypes' in data['data']:
            gpu_types = data['data']['gpuTypes']

            if not gpu_types:
                print(f"❌ GPU type '{gpu_type}' not found")
                return False

            gpu = gpu_types[0]
            print(f"\n✅ GPU information retrieved!")
            print(f"🎮 Display Name: {gpu['displayName']}")
            print(f"💾 Memory: {gpu['memoryInGb']}GB")
            print(f"🔒 Secure Cloud: {gpu['secureCloud']}")
            print(f"🌍 Community Cloud: {gpu['communityCloud']}")

            if gpu.get('lowestPrice'):
                price = gpu['lowestPrice']
                print(f"\n💰 Pricing:")
                print(f"  On-Demand: ${price.get('uninterruptablePrice', 'N/A')}/hour")
                print(f"  Spot Min Bid: ${price.get('minimumBidPrice', 'N/A')}/hour")
                print(f"  Stock Status: {price.get('stockStatus', 'Unknown')}")

            return True
        else:
            print(f"❌ Unexpected response structure")
            return False

    except requests.exceptions.Timeout:
        print("❌ Request timed out after 10 seconds")
        return False
    except requests.exceptions.RequestException as e:
        print(f"❌ Request failed: {e}")
        return False
    except json.JSONDecodeError as e:
        print(f"❌ Failed to parse JSON response: {e}")
        return False


def main():
    """Run all authentication tests"""
    print("="*60)
    print("🧪 RunPod API Authentication Test Suite")
    print("="*60)
    print("\nThis script tests authentication with the RunPod API")
    print("before attempting pod creation.")

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

    print(f"\n🔑 API Key: {api_key[:10]}...{api_key[-10:]}")

    # Run tests
    results = []

    results.append(("Authentication Test", test_authentication(api_key)))
    results.append(("User Info Test", test_user_info(api_key)))
    results.append(("GPU Availability Test", test_gpu_availability(api_key)))

    # Summary
    print("\n" + "="*60)
    print("📊 TEST SUMMARY")
    print("="*60)

    for test_name, result in results:
        status = "✅ PASS" if result else "❌ FAIL"
        print(f"{status} - {test_name}")

    passed = sum(1 for _, result in results if result)
    total = len(results)

    print(f"\n🎯 Overall: {passed}/{total} tests passed")

    if passed == total:
        print("\n✅ All tests passed! Authentication is working correctly.")
        print("✅ Ready to proceed with pod creation tests.")
        return 0
    else:
        print("\n❌ Some tests failed. Fix authentication before proceeding.")
        return 1


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