#!/usr/bin/env python3
"""
Runpod GraphQL Schema Introspection - Datacenter Field Check

This script introspects the Runpod GraphQL schema to:
1. Find all available fields in PodFindAndDeployOnDemandInput
2. Identify datacenter/region selection fields
3. Document the exact field name and type
4. Query available datacenters (if supported)

Prerequisites:
  - Set RUNPOD_API_KEY environment variable
  - Install requests: pip install requests

Usage:
  export RUNPOD_API_KEY='your-key-here'
  python3 scripts/check_runpod_datacenter_field.py
"""

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


API_URL = 'https://api.runpod.io/graphql'


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("\nSet it with:")
        print("  export RUNPOD_API_KEY='your-key-here'")
        print("\nGet your API key from: https://www.runpod.io/console/user/settings")
        return None
    return api_key


def make_graphql_request(api_key: str, query: str) -> Optional[Dict[str, Any]]:
    """Make GraphQL request to RunPod API"""
    url = f'{API_URL}?api_key={api_key}'
    headers = {'Content-Type': 'application/json'}
    payload = {'query': query}

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

        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:")
            for error in data['errors']:
                print(f"  - {error.get('message', 'Unknown error')}")
            return None

        return data

    except requests.exceptions.RequestException as e:
        print(f"❌ Request failed: {e}")
        return None


def introspect_pod_input_type(api_key: str):
    """Introspect PodFindAndDeployOnDemandInput type"""
    print("\n" + "="*70)
    print("🔍 INTROSPECTING: PodFindAndDeployOnDemandInput")
    print("="*70)

    query = '''
    {
      __type(name: "PodFindAndDeployOnDemandInput") {
        name
        inputFields {
          name
          description
          type {
            name
            kind
            ofType {
              name
              kind
            }
          }
        }
      }
    }
    '''

    data = make_graphql_request(api_key, query)
    if not data or 'data' not in data:
        return None

    type_info = data['data'].get('__type')
    if not type_info:
        print("❌ Type not found in schema")
        return None

    fields = type_info['inputFields']
    print(f"\n✅ Found {len(fields)} fields\n")

    # Print all fields
    print("📋 All Available Fields:")
    print("-" * 70)
    for field in sorted(fields, key=lambda f: f['name']):
        field_name = field['name']
        field_type = field['type']
        type_name = field_type.get('name') or field_type.get('ofType', {}).get('name', 'Unknown')
        required = '!' in str(field_type)
        req_marker = '(required)' if required else '(optional)'

        print(f"  {field_name:<30} {type_name:<20} {req_marker}")

        if field.get('description'):
            print(f"    → {field['description']}")

    # Search for datacenter/region fields
    print("\n" + "="*70)
    print("🎯 SEARCHING FOR DATACENTER/REGION FIELDS")
    print("="*70)

    keywords = ['datacenter', 'region', 'location', 'zone', 'area', 'center']
    datacenter_fields = [
        f for f in fields
        if any(keyword in f['name'].lower() for keyword in keywords)
    ]

    if datacenter_fields:
        print("\n✅ FOUND DATACENTER/REGION FIELDS:\n")
        for field in datacenter_fields:
            field_name = field['name']
            field_type = field['type']
            type_name = field_type.get('name') or field_type.get('ofType', {}).get('name', 'Unknown')

            print(f"  🎯 {field_name}")
            print(f"     Type: {type_name}")
            if field.get('description'):
                print(f"     Description: {field['description']}")
            print()

        return datacenter_fields
    else:
        print("\n❌ NO datacenter/region fields found")
        print("\nThis could mean:")
        print("  1. Region is implicitly selected based on networkVolumeId")
        print("  2. Region selection not supported in current API version")
        print("  3. Different field naming convention")
        return None


def query_datacenters(api_key: str):
    """Try to query available datacenters"""
    print("\n" + "="*70)
    print("🌍 QUERYING AVAILABLE DATACENTERS")
    print("="*70)

    # Try different possible query names
    queries = [
        ('dataCenters', '{ dataCenters { id name location } }'),
        ('datacenters', '{ datacenters { id name location } }'),
        ('regions', '{ regions { id name location } }'),
        ('locations', '{ locations { id name location } }'),
    ]

    for query_name, query in queries:
        print(f"\nTrying query: {query_name}...")
        data = make_graphql_request(api_key, query)

        if data and 'data' in data:
            result = data['data'].get(query_name)
            if result is not None:
                print(f"\n✅ Found datacenters via '{query_name}' query!")
                print(json.dumps(result, indent=2))
                return result

    print("\n⚠️  No datacenter query endpoint found")
    print("Datacenters may not be queryable via GraphQL")


def generate_recommendations(datacenter_fields: Optional[List[Dict]]):
    """Generate recommendations based on findings"""
    print("\n" + "="*70)
    print("📝 RECOMMENDATIONS")
    print("="*70)

    if datacenter_fields:
        field = datacenter_fields[0]
        field_name = field['name']

        print(f"\n✅ SOLUTION FOUND: Use '{field_name}' field")
        print("\nCode change for scripts/runpod_deploy_production.py:")
        print("-" * 70)
        print(f"""
  mutation = f\"\"\"
  mutation {{{{
    podFindAndDeployOnDemand(
      input: {{{{
        cloudType: {{cloud_type}}
        {field_name}: "eur-is-1"  # <-- ADD THIS LINE
        gpuTypeId: "{{gpu_id}}"
        name: "{{name}}"
        imageName: "{{config['image_name']}}"
        networkVolumeId: "{{config['network_volume_id']}}"
        ...
      }}}}
    ) {{{{ ... }}}}
  }}}}
  \"\"\"
        """)
        print("-" * 70)

    else:
        print("\n⚠️  NO DATACENTER FIELD FOUND")
        print("\nPossible solutions:")
        print("  1. TEST CURRENT BEHAVIOR:")
        print("     - Deploy a test pod with current script")
        print("     - Check if volume mounts correctly")
        print("     - Runpod may auto-select region based on networkVolumeId")
        print()
        print("  2. CONTACT RUNPOD SUPPORT:")
        print("     - Ask about region selection for network volumes")
        print("     - Request documentation for datacenter parameter")
        print()
        print("  3. USE SPOT INSTANCES:")
        print("     - Try podRentInterruptable mutation")
        print("     - May have different region selection options")
        print()
        print("  4. MIGRATE VOLUME:")
        print("     - Create new volume in high-availability region (us-ca-1)")
        print("     - Copy data to new volume")


def main():
    """Main execution"""
    print("="*70)
    print("RUNPOD DATACENTER FIELD INTROSPECTION")
    print("="*70)

    # Get API key
    api_key = get_api_key()
    if not api_key:
        return 1

    print("\n✅ API key found")
    print(f"🌐 API endpoint: {API_URL}")

    # Test authentication
    print("\n🔐 Testing authentication...")
    test_query = '{ gpuTypes { id displayName } }'
    test_data = make_graphql_request(api_key, test_query)

    if not test_data:
        print("❌ Authentication failed")
        return 1

    print("✅ Authentication successful")

    # Introspect PodFindAndDeployOnDemandInput
    datacenter_fields = introspect_pod_input_type(api_key)

    # Try to query datacenters
    query_datacenters(api_key)

    # Generate recommendations
    generate_recommendations(datacenter_fields)

    print("\n" + "="*70)
    print("✅ INTROSPECTION COMPLETE")
    print("="*70)

    # Save results
    output_file = 'RUNPOD_DATACENTER_INTROSPECTION_RESULTS.json'
    results = {
        'timestamp': '2025-10-24',
        'datacenter_fields_found': bool(datacenter_fields),
        'datacenter_fields': datacenter_fields if datacenter_fields else [],
    }

    with open(output_file, 'w') as f:
        json.dump(results, f, indent=2)

    print(f"\n💾 Results saved to: {output_file}")
    print("\nNext steps:")
    print("  1. Review recommendations above")
    print("  2. Update deployment script if datacenter field found")
    print("  3. Test with smoke test deployment")
    print("  4. Update CLAUDE.md with findings")

    return 0


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