#!/usr/bin/env python3
"""
RunPod Pod Monitoring Script
Continuously monitors pod metrics and displays GPU utilization, memory usage, etc.
"""

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

# Load RunPod API key
env_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), '.env.runpod')
load_dotenv(env_path)

api_key = os.getenv('RUNPOD_API_KEY')
if not api_key:
    print("ERROR: RUNPOD_API_KEY not found in .env.runpod")
    sys.exit(1)

# Default pod ID (can be overridden via command line)
pod_id = 'k18xwnvja2mk1s'
if len(sys.argv) > 1:
    pod_id = sys.argv[1]

query = """
query GetPodMetrics($podId: String!) {
  pod(input: {podId: $podId}) {
    id
    name
    desiredStatus
    runtime {
      uptimeInSeconds
      container {
        cpuPercent
        memoryPercent
      }
      gpus {
        gpuUtilPercent
        memoryUtilPercent
      }
    }
  }
}
"""

print(f"📊 Monitoring Pod: {pod_id}")
print(f"Press Ctrl+C to stop\n")

iteration = 0
while True:
    try:
        response = requests.post(
            "https://api.runpod.io/graphql",
            json={"query": query, "variables": {"podId": pod_id}},
            headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
            timeout=30
        )

        if response.status_code != 200:
            print(f"⚠️  API Error: HTTP {response.status_code}")
            time.sleep(60)
            continue

        data = response.json()

        if 'errors' in data:
            print(f"⚠️  GraphQL Error: {data['errors']}")
            time.sleep(60)
            continue

        pod = data.get('data', {}).get('pod', {})
        if not pod:
            print(f"⚠️  Pod {pod_id} not found")
            time.sleep(60)
            continue

        status = pod.get('desiredStatus', 'UNKNOWN')
        runtime = pod.get('runtime', {})

        timestamp = datetime.now().strftime('%H:%M:%S')

        if runtime:
            uptime = runtime.get('uptimeInSeconds', 0)
            container = runtime.get('container', {})
            gpus = runtime.get('gpus', [{}])

            cpu_pct = container.get('cpuPercent', 0)
            mem_pct = container.get('memoryPercent', 0)
            gpu_util = gpus[0].get('gpuUtilPercent', 0) if gpus else 0
            gpu_mem = gpus[0].get('memoryUtilPercent', 0) if gpus else 0

            # Format uptime
            hours = uptime // 3600
            minutes = (uptime % 3600) // 60
            seconds = uptime % 60
            uptime_str = f"{hours}h {minutes}m {seconds}s"

            # Color coding for GPU utilization
            gpu_indicator = "🟢" if gpu_util > 80 else ("🟡" if gpu_util > 50 else "🔴")

            print(f"[{timestamp}] {gpu_indicator} Uptime: {uptime_str:>12} | "
                  f"CPU: {cpu_pct:5.1f}% | Mem: {mem_pct:5.1f}% | "
                  f"GPU: {gpu_util:5.1f}% | VRAM: {gpu_mem:5.1f}% | "
                  f"Status: {status}")

            # Alert on low GPU utilization after 10 minutes
            if uptime > 600 and gpu_util < 50:
                print(f"    ⚠️  WARNING: Low GPU utilization after {uptime}s - check if training started")

            # Alert on high memory usage
            if mem_pct > 90:
                print(f"    ⚠️  WARNING: High memory usage ({mem_pct:.1f}%) - possible OOM risk")

        else:
            print(f"[{timestamp}] ⏳ Pod initializing... (Status: {status})")

        iteration += 1
        time.sleep(60)  # Check every 60 seconds

    except KeyboardInterrupt:
        print("\n\n✅ Monitoring stopped")
        sys.exit(0)
    except Exception as e:
        print(f"⚠️  Error: {e}")
        time.sleep(60)