#!/usr/bin/env python3
"""
Fetch RunPod Pod Logs via API
"""
import os
import sys
import requests
from dotenv import load_dotenv
from pathlib import Path

# Load environment variables
env_path = Path(__file__).parent.parent / '.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)

def get_pod_logs(pod_id, lines=50):
    """Fetch pod logs via GraphQL API."""

    # GraphQL query to get pod info including logs
    query = """
    query GetPod($input: PodFindInput!) {
      pod(input: $input) {
        id
        name
        desiredStatus
        runtime {
          uptimeInSeconds
          ports {
            ip
            isIpPublic
            privatePort
            publicPort
            type
          }
        }
        machine {
          gpuDisplayName
        }
      }
    }
    """

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

    payload = {
        "query": query,
        "variables": {
            "input": {"podId": pod_id}
        }
    }

    try:
        response = requests.post(
            "https://api.runpod.io/graphql",
            json=payload,
            headers=headers,
            timeout=30
        )
        response.raise_for_status()
        result = response.json()

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

        return result.get('data', {}).get('pod')

    except requests.exceptions.RequestException as e:
        print(f"ERROR: Failed to query RunPod API: {e}")
        return None


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: ./get_runpod_logs.py <pod_id> [lines]")
        sys.exit(1)

    pod_id = sys.argv[1]
    lines = int(sys.argv[2]) if len(sys.argv) > 2 else 50

    print(f"Fetching pod info for: {pod_id}")
    print("=" * 70)

    pod = get_pod_logs(pod_id, lines)

    if not pod:
        print("Failed to fetch pod information")
        sys.exit(1)

    # Display pod info
    print(f"\nPod ID: {pod.get('id')}")
    print(f"Name: {pod.get('name')}")
    print(f"Status: {pod.get('desiredStatus')}")

    runtime = pod.get('runtime')
    if runtime:
        uptime = runtime.get('uptimeInSeconds', 0)
        print(f"Uptime: {uptime}s ({uptime // 60}m)")

        # Show ports
        ports = runtime.get('ports', [])
        if ports:
            print("\nPorts:")
            for port in ports:
                print(f"  {port.get('privatePort')} -> {port.get('publicPort')} ({port.get('type')})")

    machine = pod.get('machine')
    if machine:
        print(f"\nGPU: {machine.get('gpuDisplayName', 'N/A')}")

    print("\n" + "=" * 70)
    print("\nNOTE: RunPod GraphQL API does not expose container logs directly.")
    print("To view logs, use one of these methods:")
    print(f"1. Web UI: https://www.runpod.io/console/pods/{pod_id}")
    print(f"2. SSH: ssh root@{pod_id}.ssh.runpod.io")
    print("3. RunPod CLI: runpodctl logs <pod_id>")
    print("\nOnce connected via SSH, check logs with:")
    print("  - docker logs <container_id>")
    print("  - tail -f /var/log/entrypoint.log (if logging to file)")
    print("  - ps aux | grep train_tft_parquet (check if process running)")