#!/usr/bin/env python3
"""
Verify RunPod pod deployment - check datacenter and logs
"""

import os
import sys
import time
import requests
from dotenv import load_dotenv

# Load environment variables
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)

REST_API_URL = "https://rest.runpod.io/v1/pods"
GRAPHQL_ENDPOINT = "https://api.runpod.io/graphql"

def get_pod_status(pod_id):
    """Get current status of a pod."""
    print(f"\n🔍 Checking pod {pod_id} status...")

    headers = {
        "Authorization": f"Bearer {RUNPOD_API_KEY}"
    }

    try:
        response = requests.get(
            f"{REST_API_URL}/{pod_id}",
            headers=headers,
            timeout=30
        )

        if response.status_code == 200:
            pod_data = response.json()

            print("\n" + "="*70)
            print("POD STATUS")
            print("="*70)

            print(f"Pod ID:          {pod_data.get('id', 'N/A')}")
            print(f"Name:            {pod_data.get('name', 'N/A')}")
            print(f"Status:          {pod_data.get('desiredStatus', 'UNKNOWN')}")

            runtime = pod_data.get('runtime', {})
            print(f"Runtime Status:  {runtime.get('status', 'UNKNOWN')}")

            machine = pod_data.get('machine', {})
            datacenter = machine.get('dataCenterId', 'N/A')
            print(f"Datacenter:      {datacenter}")

            gpu_info = machine.get('gpuType', {})
            print(f"GPU:             {gpu_info.get('displayName', 'N/A')}")
            print(f"GPU Count:       {machine.get('gpuCount', 'N/A')}")

            print(f"Image:           {pod_data.get('imageName', 'N/A')}")
            print(f"Cost:            ${pod_data.get('costPerHr', 'N/A')}/hr")

            print("="*70)

            # Check if datacenter is EUR-IS-1
            if datacenter == 'EUR-IS-1':
                print(f"\n✅ DATACENTER VERIFIED: Pod is in EUR-IS-1 (volume accessible)")
            else:
                print(f"\n❌ DATACENTER MISMATCH: Pod is in {datacenter}, volume is in EUR-IS-1!")

            return pod_data
        elif response.status_code == 404:
            print(f"  ⚠️  Pod not found")
            return None
        else:
            print(f"  ⚠️  Failed to get pod status (status {response.status_code}): {response.text[:200]}")
            return None

    except requests.exceptions.RequestException as e:
        print(f"  ❌ Failed to check pod status: {e}")
        return None


def get_pod_logs_graphql(pod_id):
    """Get logs from a pod using GraphQL (REST API logs endpoint doesn't work)."""
    print(f"\n📝 Fetching logs for pod {pod_id} via GraphQL...")

    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {RUNPOD_API_KEY}"
    }

    query = """
    query GetPodLogs($podId: String!) {
      pod(input: {podId: $podId}) {
        logs
      }
    }
    """

    variables = {
        "podId": pod_id
    }

    try:
        response = requests.post(
            GRAPHQL_ENDPOINT,
            json={"query": query, "variables": variables},
            headers=headers,
            timeout=30
        )

        if response.status_code == 200:
            result = response.json()

            if 'errors' in result:
                print(f"  ⚠️  GraphQL errors: {result['errors']}")
                return None

            logs = result.get('data', {}).get('pod', {}).get('logs', '')

            if logs:
                print("\n" + "="*70)
                print("CONTAINER LOGS")
                print("="*70)
                # Print last 50 lines
                log_lines = logs.split('\n')
                for line in log_lines[-50:]:
                    print(line)
                print("="*70)

                # Check for volume mounting indicators
                if '/runpod-volume' in logs:
                    print("\n✅ Volume path /runpod-volume found in logs")
                if 'No such file or directory' in logs and '/runpod-volume' in logs:
                    print("\n❌ Volume NOT mounted - file not found errors detected")
                if 'test_data/ES_FUT_180d.parquet' in logs:
                    print("\n✅ Parquet file path found in logs")

                return logs
            else:
                print("  ⚠️  No logs available yet (pod may still be initializing)")
                return None
        else:
            print(f"  ⚠️  Failed to get logs (status {response.status_code}): {response.text[:200]}")
            return None

    except requests.exceptions.RequestException as e:
        print(f"  ❌ Failed to fetch logs: {e}")
        return None


def main():
    print("="*70)
    print("VERIFY RUNPOD DEPLOYMENT")
    print("="*70)

    if len(sys.argv) < 2:
        print("\nUsage: python3 verify_pod_deployment.py <pod_id>")
        print("\nExample:")
        print("  python3 verify_pod_deployment.py b1m7v451nexg5r")
        sys.exit(1)

    pod_id = sys.argv[1]

    # Check pod status
    pod_data = get_pod_status(pod_id)

    if not pod_data:
        print(f"\n❌ Failed to get pod status")
        sys.exit(1)

    # Wait a bit for container to start
    runtime_status = pod_data.get('runtime', {}).get('status', 'UNKNOWN')
    if runtime_status != 'running':
        print(f"\n⏳ Container not running yet (status: {runtime_status})")
        print("   Waiting 30 seconds for container to start...")
        time.sleep(30)

    # Get logs
    get_pod_logs_graphql(pod_id)

    print("\n" + "="*70)
    print("VERIFICATION COMPLETE")
    print("="*70)


if __name__ == "__main__":
    main()