#!/usr/bin/env python3
"""
RunPod GPU Scanner
Queries RunPod API to show available SECURE cloud GPUs with ≥16GB VRAM.
"""

import os
import sys
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')
if not RUNPOD_API_KEY:
    print("ERROR: RUNPOD_API_KEY not found in .env.runpod")
    sys.exit(1)

GRAPHQL_ENDPOINT = "https://api.runpod.io/graphql"

# GraphQL query to fetch GPU types
QUERY = """
{
  gpuTypes {
    id
    displayName
    memoryInGb
    secureCloud
    communityCloud
    lowestPrice(input: {gpuCount: 1}) {
      uninterruptablePrice
    }
  }
}
"""

def query_runpod_api():
    """Query RunPod GraphQL API for GPU types."""
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {RUNPOD_API_KEY}"
    }

    payload = {
        "query": QUERY
    }

    try:
        response = requests.post(GRAPHQL_ENDPOINT, json=payload, headers=headers, timeout=10)
        response.raise_for_status()
        return response.json()
    except requests.exceptions.RequestException as e:
        print(f"ERROR: Failed to query RunPod API: {e}")
        sys.exit(1)

def filter_and_sort_gpus(data):
    """Filter for secure cloud GPUs with ≥16GB VRAM, sorted by price."""
    gpu_types = data.get('data', {}).get('gpuTypes', [])

    # Filter criteria:
    # 1. memoryInGb >= 16
    # 2. secureCloud > 0 (available in secure cloud)
    # 3. Has pricing information
    filtered_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:
            filtered_gpus.append({
                'name': gpu.get('displayName', 'Unknown'),
                'vram': memory,
                'price': float(price),
                'available': secure_count
            })

    # Sort by price (cheapest first)
    filtered_gpus.sort(key=lambda x: x['price'])

    return filtered_gpus

def display_gpus(gpus):
    """Display GPU information in clean format."""
    if not gpus:
        print("No secure cloud GPUs found with ≥16GB VRAM in EUR-IS region.")
        return

    print("\n" + "="*70)
    print("RUNPOD SECURE CLOUD GPUs (≥16GB VRAM) - EUR-IS REGION")
    print("="*70)
    print(f"{'GPU Name':<30} {'VRAM':<10} {'Price/hr':<12} {'Available':<10}")
    print("-"*70)

    for gpu in gpus:
        name = gpu['name']
        vram = f"{gpu['vram']}GB"
        price = f"${gpu['price']:.3f}"
        available = f"{gpu['available']} pods"
        print(f"{name:<30} {vram:<10} {price:<12} {available:<10}")

    print("="*70)
    print(f"Total GPUs found: {len(gpus)}")
    print()

def main():
    """Main execution function."""
    print("Querying RunPod API...")

    # Query API
    data = query_runpod_api()

    # Filter and sort GPUs
    gpus = filter_and_sort_gpus(data)

    # Display results
    display_gpus(gpus)

if __name__ == "__main__":
    main()