#!/usr/bin/env python3
"""
Runpod GraphQL API Deployment Script for Foxhunt FP32 Training

This script replaces runpodctl with direct GraphQL/REST API calls for more reliable deployments.
It supports both smoke tests (1 epoch) and full training (50 epochs).

Requirements:
    pip install requests

Usage:
    # Smoke test (1 epoch, small dataset)
    ./scripts/deploy_runpod_graphql.py --smoke-test
    
    # Full training (50 epochs, 180 days)
    ./scripts/deploy_runpod_graphql.py --full-training
    
    # Custom configuration
    ./scripts/deploy_runpod_graphql.py --epochs 10 --parquet-file /runpod-volume/test_data/ES_FUT_180d.parquet
    
    # List available GPUs
    ./scripts/deploy_runpod_graphql.py --list-gpus
    
    # Get Docker credential ID
    ./scripts/deploy_runpod_graphql.py --get-credentials

Environment Variables:
    RUNPOD_API_KEY: Your Runpod API key (required)

Author: Claude Code
Date: 2025-10-24
"""

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


class RunpodAPI:
    """Runpod API client using GraphQL and REST endpoints."""
    
    GRAPHQL_ENDPOINT = "https://api.runpod.io/graphql"
    REST_ENDPOINT = "https://api.runpod.io/v2"
    
    def __init__(self, api_key: str):
        """Initialize Runpod API client.
        
        Args:
            api_key: Runpod API key
        """
        if not api_key:
            raise ValueError("RUNPOD_API_KEY is required")
        
        self.api_key = api_key
        self.session = requests.Session()
        self.session.headers.update({
            "Content-Type": "application/json"
        })
    
    def _graphql_request(self, query: str, variables: Optional[Dict] = None) -> Dict:
        """Execute a GraphQL query.
        
        Args:
            query: GraphQL query string
            variables: Optional query variables
            
        Returns:
            Response data dictionary
            
        Raises:
            RuntimeError: If the GraphQL request fails
        """
        payload = {"query": query}
        if variables:
            payload["variables"] = variables
        
        response = self.session.post(
            self.GRAPHQL_ENDPOINT,
            headers={"api-key": self.api_key},
            json=payload,
            timeout=30
        )
        
        if response.status_code != 200:
            raise RuntimeError(f"GraphQL request failed: {response.status_code} - {response.text}")
        
        data = response.json()
        
        if "errors" in data:
            errors = data["errors"]
            error_messages = [err.get("message", str(err)) for err in errors]
            raise RuntimeError(f"GraphQL errors: {', '.join(error_messages)}")
        
        return data.get("data", {})
    
    def _rest_request(self, method: str, endpoint: str, data: Optional[Dict] = None) -> Dict:
        """Execute a REST API request.
        
        Args:
            method: HTTP method (GET, POST, DELETE, etc.)
            endpoint: API endpoint path
            data: Optional request body data
            
        Returns:
            Response data dictionary
            
        Raises:
            RuntimeError: If the REST request fails
        """
        url = f"{self.REST_ENDPOINT}/{endpoint}"
        
        response = self.session.request(
            method,
            url,
            headers={"Authorization": f"Bearer {self.api_key}"},
            json=data,
            timeout=30
        )
        
        if response.status_code not in (200, 201):
            raise RuntimeError(f"REST request failed: {response.status_code} - {response.text}")
        
        return response.json()
    
    def get_docker_credential_id(self, name: str = "Docker") -> Optional[str]:
        """Get Docker registry credential ID by name.
        
        Args:
            name: Credential name (default: "Docker")
            
        Returns:
            Credential ID or None if not found
        """
        query = """
        query {
            myself {
                containerRegistryAuths {
                    id
                    name
                }
            }
        }
        """
        
        try:
            data = self._graphql_request(query)
            auths = data.get("myself", {}).get("containerRegistryAuths", [])
            
            for auth in auths:
                if auth.get("name") == name:
                    return auth.get("id")
            
            return None
        except Exception as e:
            print(f"Warning: Failed to get Docker credentials: {e}")
            return None
    
    def list_available_gpus(self) -> List[Dict]:
        """List available GPU types.
        
        Returns:
            List of GPU type dictionaries with id, name, and pricing
        """
        query = """
        query {
            gpuTypes {
                id
                displayName
                memoryInGb
                secureCloud
                communityCloud
                lowestPrice(input: { gpuCount: 1 }) {
                    minimumBidPrice
                    uninterruptablePrice
                }
            }
        }
        """
        
        data = self._graphql_request(query)
        return data.get("gpuTypes", [])
    
    def create_pod_rest(
        self,
        name: str,
        image_name: str,
        gpu_type_id: str,
        docker_start_cmd: List[str],
        network_volume_id: Optional[str] = None,
        container_registry_auth_id: Optional[str] = None,
        env: Optional[Dict[str, str]] = None,
        cloud_type: str = "SECURE",
        gpu_count: int = 1,
        container_disk_in_gb: int = 50,
        volume_mount_path: str = "/workspace",
        ports: Optional[List[str]] = None
    ) -> Dict:
        """Create a pod using the REST API (better support for dockerStartCmd).
        
        Args:
            name: Pod name
            image_name: Docker image name
            gpu_type_id: GPU type ID
            docker_start_cmd: Docker CMD override (list of strings)
            network_volume_id: Network volume ID (optional)
            container_registry_auth_id: Docker credential ID (optional)
            env: Environment variables (optional)
            cloud_type: "SECURE" or "COMMUNITY" (default: SECURE)
            gpu_count: Number of GPUs (default: 1)
            container_disk_in_gb: Container disk size (default: 50GB)
            volume_mount_path: Volume mount path (default: /workspace)
            ports: Exposed ports (optional)
            
        Returns:
            Pod creation response with pod ID
        """
        payload = {
            "name": name,
            "imageName": image_name,
            "gpuTypeIds": [gpu_type_id],
            "cloudType": cloud_type,
            "gpuCount": gpu_count,
            "containerDiskInGb": container_disk_in_gb,
            "volumeMountPath": volume_mount_path,
            "dockerStartCmd": docker_start_cmd,
        }
        
        # Add optional parameters
        if network_volume_id:
            payload["networkVolumeId"] = network_volume_id
        
        if container_registry_auth_id:
            payload["containerRegistryAuthId"] = container_registry_auth_id
        
        if env:
            payload["env"] = env
        
        if ports:
            payload["ports"] = ports
        
        return self._rest_request("POST", "pods", payload)
    
    def get_pod_status(self, pod_id: str) -> Dict:
        """Get pod status.
        
        Args:
            pod_id: Pod ID
            
        Returns:
            Pod status dictionary
        """
        query = """
        query GetPod($podId: String!) {
            pod(input: { podId: $podId }) {
                id
                name
                desiredStatus
                imageName
                costPerHr
                gpuCount
                machine {
                    podHostId
                }
                runtime {
                    uptimeInSeconds
                    ports {
                        ip
                        isIpPublic
                        privatePort
                        publicPort
                        type
                    }
                }
            }
        }
        """
        
        data = self._graphql_request(query, {"podId": pod_id})
        return data.get("pod", {})
    
    def terminate_pod(self, pod_id: str) -> Dict:
        """Terminate a pod.
        
        Args:
            pod_id: Pod ID
            
        Returns:
            Termination response
        """
        return self._rest_request("DELETE", f"pods/{pod_id}", None)
    
    def list_pods(self) -> List[Dict]:
        """List all user pods.
        
        Returns:
            List of pod dictionaries
        """
        query = """
        query {
            myself {
                pods {
                    id
                    name
                    desiredStatus
                    imageName
                    costPerHr
                    gpuCount
                }
            }
        }
        """
        
        data = self._graphql_request(query)
        return data.get("myself", {}).get("pods", [])


def format_docker_cmd(
    binary: str,
    parquet_file: str,
    epochs: int,
    use_int8: bool = False,
    use_qat: bool = False,
    additional_args: Optional[List[str]] = None
) -> List[str]:
    """Format Docker CMD for training.
    
    Args:
        binary: Binary name (e.g., "train_tft_parquet")
        parquet_file: Path to parquet file
        epochs: Number of epochs
        use_int8: Enable INT8 quantization
        use_qat: Enable QAT (requires INT8)
        additional_args: Additional arguments
        
    Returns:
        List of command arguments
    """
    cmd = [
        "--parquet-file", parquet_file,
        "--epochs", str(epochs)
    ]
    
    if use_int8:
        cmd.append("--use-int8")
    
    if use_qat:
        cmd.append("--use-qat")
    
    if additional_args:
        cmd.extend(additional_args)
    
    return cmd


def wait_for_pod_ready(api: RunpodAPI, pod_id: str, timeout: int = 300) -> bool:
    """Wait for pod to be ready.
    
    Args:
        api: Runpod API client
        pod_id: Pod ID
        timeout: Timeout in seconds (default: 300)
        
    Returns:
        True if pod is ready, False if timeout
    """
    start_time = time.time()
    
    while time.time() - start_time < timeout:
        try:
            status = api.get_pod_status(pod_id)
            desired_status = status.get("desiredStatus", "")
            
            if desired_status == "RUNNING":
                print(f"✅ Pod {pod_id} is RUNNING")
                return True
            
            print(f"⏳ Pod status: {desired_status} (waiting...)")
            time.sleep(10)
        except Exception as e:
            print(f"Warning: Failed to get pod status: {e}")
            time.sleep(10)
    
    print(f"❌ Timeout waiting for pod to be ready")
    return False


def print_pod_info(pod_data: Dict):
    """Print pod information.
    
    Args:
        pod_data: Pod data dictionary
    """
    print("\n" + "="*80)
    print("POD CREATED SUCCESSFULLY")
    print("="*80)
    
    pod_id = pod_data.get("id") or pod_data.get("podId")
    print(f"\n📦 Pod ID: {pod_id}")
    print(f"📛 Name: {pod_data.get('name', 'N/A')}")
    print(f"🖼️  Image: {pod_data.get('imageName', 'N/A')}")
    print(f"💰 Cost: ${pod_data.get('costPerHr', 0):.4f}/hr")
    
    # Print SSH/connection info if available
    runtime = pod_data.get("runtime", {})
    if runtime:
        ports = runtime.get("ports", [])
        for port in ports:
            if port.get("type") == "tcp" and port.get("privatePort") == 22:
                ip = port.get("ip")
                public_port = port.get("publicPort")
                print(f"\n🔌 SSH Connection:")
                print(f"   ssh root@{ip} -p {public_port}")
    
    print(f"\n⚠️  IMPORTANT: Terminate pod when done to avoid charges!")
    print(f"   Terminate: curl -X DELETE https://api.runpod.io/v2/pods/{pod_id} \\")
    print(f"              -H 'Authorization: Bearer $RUNPOD_API_KEY'")
    print("\n" + "="*80 + "\n")


def main():
    """Main entry point."""
    parser = argparse.ArgumentParser(
        description="Deploy Foxhunt ML training to Runpod using GraphQL API",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Smoke test (1 epoch)
  %(prog)s --smoke-test
  
  # Full training (50 epochs)
  %(prog)s --full-training
  
  # Custom training
  %(prog)s --epochs 10 --gpu-type "NVIDIA RTX A4000"
  
  # List available GPUs
  %(prog)s --list-gpus
  
  # Get Docker credential ID
  %(prog)s --get-credentials
        """
    )
    
    # Deployment options
    parser.add_argument("--smoke-test", action="store_true",
                       help="Run smoke test (1 epoch, small dataset)")
    parser.add_argument("--full-training", action="store_true",
                       help="Run full training (50 epochs, 180 days)")
    
    # Training parameters
    parser.add_argument("--epochs", type=int, default=10,
                       help="Number of training epochs (default: 10)")
    parser.add_argument("--binary", default="train_tft_parquet",
                       help="Training binary name (default: train_tft_parquet)")
    parser.add_argument("--parquet-file", 
                       default="/runpod-volume/test_data/ES_FUT_180d.parquet",
                       help="Path to parquet file on network volume")
    parser.add_argument("--use-int8", action="store_true",
                       help="Enable INT8 post-training quantization")
    parser.add_argument("--use-qat", action="store_true",
                       help="Enable QAT quantization (experimental)")
    
    # Pod configuration
    parser.add_argument("--pod-name", 
                       help="Pod name (auto-generated if not specified)")
    parser.add_argument("--image", default="jgrusewski/foxhunt:latest",
                       help="Docker image (default: jgrusewski/foxhunt:latest)")
    parser.add_argument("--gpu-type", default="NVIDIA RTX A4000",
                       help="GPU type ID (default: NVIDIA RTX A4000)")
    parser.add_argument("--network-volume-id", default="se3zdnb5o4",
                       help="Network volume ID (default: se3zdnb5o4)")
    parser.add_argument("--cloud-type", default="SECURE",
                       choices=["SECURE", "COMMUNITY"],
                       help="Cloud type (default: SECURE)")
    parser.add_argument("--container-disk", type=int, default=50,
                       help="Container disk size in GB (default: 50)")
    
    # Utility options
    parser.add_argument("--list-gpus", action="store_true",
                       help="List available GPU types and exit")
    parser.add_argument("--get-credentials", action="store_true",
                       help="Get Docker credential ID and exit")
    parser.add_argument("--list-pods", action="store_true",
                       help="List all active pods and exit")
    parser.add_argument("--terminate", metavar="POD_ID",
                       help="Terminate a specific pod and exit")
    
    args = parser.parse_args()
    
    # Get API key
    api_key = os.getenv("RUNPOD_API_KEY")
    if not api_key:
        print("❌ Error: RUNPOD_API_KEY environment variable is not set")
        print("\nTo set your API key:")
        print("  export RUNPOD_API_KEY='your-api-key-here'")
        print("\nGet your API key from: https://www.runpod.io/console/user/settings")
        sys.exit(1)
    
    # Initialize API client
    try:
        api = RunpodAPI(api_key)
    except Exception as e:
        print(f"❌ Error initializing Runpod API: {e}")
        sys.exit(1)
    
    # Handle utility commands
    if args.list_gpus:
        print("\n📊 Available GPU Types:\n")
        gpus = api.list_available_gpus()
        for gpu in gpus:
            name = gpu.get("displayName", "Unknown")
            gpu_id = gpu.get("id", "Unknown")
            memory = gpu.get("memoryInGb", 0)
            lowest_price = gpu.get("lowestPrice", {})
            spot_price = lowest_price.get("minimumBidPrice", 0)
            on_demand = lowest_price.get("uninterruptablePrice", 0)
            
            print(f"  {name} ({gpu_id})")
            print(f"    Memory: {memory}GB")
            print(f"    Spot: ${spot_price:.4f}/hr | On-Demand: ${on_demand:.4f}/hr")
            print()
        sys.exit(0)
    
    if args.get_credentials:
        print("\n🔑 Fetching Docker credentials...\n")
        cred_id = api.get_docker_credential_id("Docker")
        if cred_id:
            print(f"✅ Docker credential ID: {cred_id}")
        else:
            print("❌ No credential named 'Docker' found")
            print("\nCreate credentials at: https://www.runpod.io/console/user/settings")
        sys.exit(0)
    
    if args.list_pods:
        print("\n📦 Active Pods:\n")
        pods = api.list_pods()
        if not pods:
            print("  No active pods")
        else:
            for pod in pods:
                print(f"  {pod.get('id')} - {pod.get('name')}")
                print(f"    Status: {pod.get('desiredStatus')}")
                print(f"    Image: {pod.get('imageName')}")
                print(f"    Cost: ${pod.get('costPerHr', 0):.4f}/hr")
                print()
        sys.exit(0)
    
    if args.terminate:
        print(f"\n🛑 Terminating pod {args.terminate}...\n")
        try:
            api.terminate_pod(args.terminate)
            print(f"✅ Pod {args.terminate} terminated successfully")
        except Exception as e:
            print(f"❌ Error terminating pod: {e}")
            sys.exit(1)
        sys.exit(0)
    
    # Handle deployment presets
    if args.smoke_test:
        args.epochs = 1
        args.parquet_file = "/runpod-volume/test_data/ES_FUT_small.parquet"
        print("\n🧪 SMOKE TEST MODE: 1 epoch, small dataset\n")
    
    if args.full_training:
        args.epochs = 50
        args.parquet_file = "/runpod-volume/test_data/ES_FUT_180d.parquet"
        print("\n🚀 FULL TRAINING MODE: 50 epochs, 180 days\n")
    
    # Generate pod name if not specified
    if not args.pod_name:
        timestamp = int(time.time())
        args.pod_name = f"foxhunt-{args.binary}-{args.epochs}ep-{timestamp}"
    
    # Get Docker credential ID
    print("🔑 Fetching Docker credentials...")
    container_registry_auth_id = api.get_docker_credential_id("Docker")
    if container_registry_auth_id:
        print(f"✅ Using Docker credential: {container_registry_auth_id}")
    else:
        print("⚠️  Warning: No Docker credential found (using public images only)")
    
    # Format Docker CMD
    docker_cmd = format_docker_cmd(
        args.binary,
        args.parquet_file,
        args.epochs,
        args.use_int8,
        args.use_qat
    )
    
    # Environment variables
    env = {
        "BINARY_NAME": args.binary,
        "RUST_LOG": "info",
        "CUDA_VISIBLE_DEVICES": "0"
    }
    
    # Print deployment summary
    print("\n" + "="*80)
    print("DEPLOYMENT SUMMARY")
    print("="*80)
    print(f"Pod Name:          {args.pod_name}")
    print(f"Image:             {args.image}")
    print(f"GPU Type:          {args.gpu_type}")
    print(f"Cloud Type:        {args.cloud_type}")
    print(f"Network Volume:    {args.network_volume_id}")
    print(f"Binary:            {args.binary}")
    print(f"Parquet File:      {args.parquet_file}")
    print(f"Epochs:            {args.epochs}")
    print(f"INT8:              {args.use_int8}")
    print(f"QAT:               {args.use_qat}")
    print(f"Docker CMD:        {' '.join(docker_cmd)}")
    print("="*80 + "\n")
    
    # Confirm deployment
    confirm = input("Proceed with deployment? (yes/no): ")
    if confirm.lower() != "yes":
        print("❌ Deployment cancelled")
        sys.exit(0)
    
    # Create pod
    print("\n🚀 Creating pod...")
    try:
        response = api.create_pod_rest(
            name=args.pod_name,
            image_name=args.image,
            gpu_type_id=args.gpu_type,
            docker_start_cmd=docker_cmd,
            network_volume_id=args.network_volume_id,
            container_registry_auth_id=container_registry_auth_id,
            env=env,
            cloud_type=args.cloud_type,
            container_disk_in_gb=args.container_disk,
            volume_mount_path="/runpod-volume",
            ports=["8888/http", "22/tcp"]
        )
        
        pod_id = response.get("id") or response.get("podId")
        if not pod_id:
            raise RuntimeError(f"No pod ID in response: {response}")
        
        print(f"✅ Pod created: {pod_id}")
        
        # Wait for pod to be ready
        print("\n⏳ Waiting for pod to be ready...")
        if wait_for_pod_ready(api, pod_id, timeout=300):
            # Get full pod info
            pod_data = api.get_pod_status(pod_id)
            print_pod_info(pod_data)
            
            print("\n📝 Next Steps:")
            print("  1. Monitor pod logs in Runpod console")
            print("  2. Training will start automatically")
            print("  3. Checkpoints saved to /runpod-volume/models/")
            print("  4. Terminate pod when done to avoid charges")
            print(f"\n  Terminate command:")
            print(f"    ./scripts/deploy_runpod_graphql.py --terminate {pod_id}")
        else:
            print("⚠️  Pod creation timed out, but pod may still be starting")
            print(f"   Check status: ./scripts/deploy_runpod_graphql.py --list-pods")
            print(f"   Terminate if needed: ./scripts/deploy_runpod_graphql.py --terminate {pod_id}")
    
    except Exception as e:
        print(f"❌ Error creating pod: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()