#!/usr/bin/env python3
"""
RunPod Training Pod Deployment Script

Deploys a Foxhunt training pod to RunPod with proper configuration.
Based on thorough documentation research and incremental testing.

Prerequisites:
  1. Run test_runpod_auth.py to verify API authentication
  2. Run test_runpod_pod_creation.py to test basic pod creation
  3. Set RUNPOD_API_KEY environment variable
  4. Ensure jgrusewski/foxhunt:latest image is pushed to Docker Hub

Usage:
  ./scripts/deploy_runpod_training.py [--model MODEL] [--epochs EPOCHS] [--dry-run]

Examples:
  # Deploy TFT training with default settings
  ./scripts/deploy_runpod_training.py

  # Deploy MAMBA-2 training for 100 epochs
  ./scripts/deploy_runpod_training.py --model mamba2 --epochs 100

  # Test without actually creating pod
  ./scripts/deploy_runpod_training.py --dry-run
"""

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


# Configuration defaults
DEFAULT_CONFIG = {
    'gpu_type': 'NVIDIA RTX A4000',
    'gpu_count': 1,
    'cloud_type': 'SECURE',
    'container_disk_gb': 50,
    'network_volume_id': 'se3zdnb5o4',
    'volume_mount_path': '/workspace',
    'min_vcpu': 4,
    'min_memory_gb': 16,
    'image_name': 'jgrusewski/foxhunt:latest',
}

# Model training commands
TRAINING_COMMANDS = {
    'tft': 'cargo run -p ml --example train_tft_parquet --release --features cuda -- --parquet-file /workspace/test_data/ES_FUT_180d.parquet --epochs {epochs}',
    'tft-int8': 'cargo run -p ml --example train_tft_parquet --release --features cuda -- --parquet-file /workspace/test_data/ES_FUT_180d.parquet --epochs {epochs} --use-int8',
    'tft-qat': 'cargo run -p ml --example train_tft_parquet --release --features cuda -- --parquet-file /workspace/test_data/ES_FUT_180d.parquet --epochs {epochs} --use-qat',
    'mamba2': 'cargo run -p ml --example train_mamba2_dbn --release --features cuda -- --epochs {epochs}',
    'dqn': 'cargo run -p ml --example train_dqn --release --features cuda -- --epochs {epochs}',
    'ppo': 'cargo run -p ml --example train_ppo --release --features cuda -- --epochs {epochs}',
}


def get_api_key() -> Optional[str]:
    """Get RunPod API key from environment"""
    api_key = os.environ.get('RUNPOD_API_KEY')
    if not api_key:
        print("❌ ERROR: RUNPOD_API_KEY environment variable not set")
        print("Set it with: export RUNPOD_API_KEY='your_api_key'")
        return None
    return api_key


def build_pod_mutation(
    name: str,
    docker_cmd: str,
    config: Dict[str, Any]
) -> str:
    """
    Build GraphQL mutation for pod creation

    Args:
        name: Pod name
        docker_cmd: Docker CMD override string
        config: Configuration dictionary

    Returns:
        GraphQL mutation string
    """
    # Escape quotes in docker command for GraphQL
    escaped_cmd = docker_cmd.replace('"', '\\"')

    mutation = f"""
    mutation {{
      podFindAndDeployOnDemand(
        input: {{
          cloudType: {config['cloud_type']}
          gpuTypeId: "{config['gpu_type']}"
          gpuCount: {config['gpu_count']}
          name: "{name}"
          imageName: "{config['image_name']}"

          containerDiskInGb: {config['container_disk_gb']}
          networkVolumeId: "{config['network_volume_id']}"
          volumeMountPath: "{config['volume_mount_path']}"

          minVcpuCount: {config['min_vcpu']}
          minMemoryInGb: {config['min_memory_gb']}

          dockerArgs: "{escaped_cmd}"

          env: [
            {{ key: "RUST_LOG", value: "info" }},
            {{ key: "BINARY_NAME", value: "train_model" }}
          ]

          ports: "22/tcp"
          startSsh: true
        }}
      ) {{
        id
        name
        imageName
        desiredStatus
        machineId
        machine {{
          podHostId
        }}
      }}
    }}
    """

    return mutation


def deploy_pod(
    api_key: str,
    name: str,
    docker_cmd: str,
    config: Dict[str, Any],
    dry_run: bool = False
) -> Optional[Dict[str, Any]]:
    """
    Deploy a training pod to RunPod

    Args:
        api_key: RunPod API key
        name: Pod name
        docker_cmd: Docker CMD override
        config: Configuration dictionary
        dry_run: If True, print mutation but don't execute

    Returns:
        Pod information dict if successful, None otherwise
    """
    print("\n" + "="*70)
    print("🚀 Deploying Training Pod to RunPod")
    print("="*70)

    # Build mutation
    mutation = build_pod_mutation(name, docker_cmd, config)

    # Print configuration
    print("\n📋 Configuration:")
    print(f"  Pod Name: {name}")
    print(f"  GPU: {config['gpu_type']} x{config['gpu_count']}")
    print(f"  Cloud: {config['cloud_type']}")
    print(f"  Image: {config['image_name']}")
    print(f"  Resources: {config['min_vcpu']} vCPU, {config['min_memory_gb']}GB RAM")
    print(f"  Storage: {config['container_disk_gb']}GB container + network volume")
    print(f"  Network Volume: {config['network_volume_id']}")
    print(f"  Mount Path: {config['volume_mount_path']}")
    print(f"\n🔧 Training Command:")
    print(f"  {docker_cmd}")

    if dry_run:
        print("\n🔍 DRY RUN MODE - GraphQL Mutation:")
        print("="*70)
        print(mutation)
        print("="*70)
        print("\n✅ Dry run complete. No pod was created.")
        return None

    # Execute mutation
    url = f'https://api.runpod.io/graphql?api_key={api_key}'

    headers = {
        'Content-Type': 'application/json'
    }

    payload = {'query': mutation}

    print("\n📡 Sending deployment request to RunPod...")

    try:
        response = requests.post(url, json=payload, headers=headers, timeout=60)

        print(f"📥 Response Status: {response.status_code}")

        if response.status_code != 200:
            print(f"❌ HTTP Error: {response.status_code}")
            print(f"Response: {response.text}")
            return None

        data = response.json()

        # Check for errors
        if 'errors' in data:
            print(f"\n❌ GraphQL Errors:")
            for error in data['errors']:
                print(f"  - {error.get('message', str(error))}")
            return None

        # Check for successful pod creation
        if 'data' in data and data['data'].get('podFindAndDeployOnDemand'):
            pod = data['data']['podFindAndDeployOnDemand']

            print("\n" + "="*70)
            print("✅ Pod Deployed Successfully!")
            print("="*70)
            print(f"\n🆔 Pod ID: {pod['id']}")
            print(f"📛 Name: {pod['name']}")
            print(f"🖼️  Image: {pod['imageName']}")
            print(f"📊 Status: {pod['desiredStatus']}")

            if pod.get('machineId'):
                print(f"🖥️  Machine ID: {pod['machineId']}")

            if pod.get('machine') and pod['machine'].get('podHostId'):
                print(f"🏠 Host ID: {pod['machine']['podHostId']}")

            print(f"\n🌐 Access your pod:")
            print(f"  https://www.runpod.io/console/pods/{pod['id']}")

            print(f"\n💰 Cost Estimate:")
            print(f"  ~$0.50/hour for {config['gpu_type']}")
            print(f"  ~$0.05-0.10 per training run (3-5 minutes)")

            print(f"\n⚠️  IMPORTANT:")
            print(f"  - Training will start automatically")
            print(f"  - Monitor progress in RunPod console")
            print(f"  - STOP the pod when training completes to avoid charges")
            print(f"  - Models saved to network volume will persist")

            return pod
        else:
            print(f"\n❌ Pod creation returned null")
            print(f"This usually means no capacity available for {config['gpu_type']}")
            print(f"\nResponse: {json.dumps(data, indent=2)}")
            return None

    except requests.exceptions.Timeout:
        print("\n❌ Request timed out after 60 seconds")
        return None
    except requests.exceptions.RequestException as e:
        print(f"\n❌ Request failed: {e}")
        return None
    except Exception as e:
        print(f"\n❌ Unexpected error: {e}")
        import traceback
        traceback.print_exc()
        return None


def stop_pod(api_key: str, pod_id: str) -> bool:
    """Stop a running pod"""
    url = f'https://api.runpod.io/graphql?api_key={api_key}'

    headers = {
        'Content-Type': 'application/json'
    }

    mutation = f"""
    mutation {{
      podStop(input: {{podId: "{pod_id}"}}) {{
        id
        desiredStatus
      }}
    }}
    """

    payload = {'query': mutation}

    print(f"\n🛑 Stopping pod {pod_id}...")

    try:
        response = requests.post(url, json=payload, headers=headers, timeout=10)

        if response.status_code != 200:
            print(f"❌ Failed to stop pod: HTTP {response.status_code}")
            return False

        data = response.json()

        if 'errors' in data:
            print(f"❌ Errors stopping pod: {data['errors']}")
            return False

        print(f"✅ Pod stopped successfully")
        return True

    except Exception as e:
        print(f"❌ Error stopping pod: {e}")
        return False


def parse_args():
    """Parse command line arguments"""
    parser = argparse.ArgumentParser(
        description='Deploy Foxhunt training pod to RunPod',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Deploy TFT training with default settings
  %(prog)s

  # Deploy MAMBA-2 training for 100 epochs
  %(prog)s --model mamba2 --epochs 100

  # Deploy TFT with INT8 quantization
  %(prog)s --model tft-int8 --epochs 50

  # Test without creating pod
  %(prog)s --dry-run

Available models:
  tft       - Temporal Fusion Transformer (FP32)
  tft-int8  - TFT with INT8 post-training quantization
  tft-qat   - TFT with quantization-aware training (QAT)
  mamba2    - MAMBA-2 state space model
  dqn       - Deep Q-Network
  ppo       - Proximal Policy Optimization
        """
    )

    parser.add_argument(
        '--model',
        type=str,
        default='tft',
        choices=list(TRAINING_COMMANDS.keys()),
        help='Model to train (default: tft)'
    )

    parser.add_argument(
        '--epochs',
        type=int,
        default=50,
        help='Number of training epochs (default: 50)'
    )

    parser.add_argument(
        '--dry-run',
        action='store_true',
        help='Print configuration without creating pod'
    )

    parser.add_argument(
        '--gpu',
        type=str,
        default=None,
        help='GPU type (default: NVIDIA RTX A4000)'
    )

    parser.add_argument(
        '--name',
        type=str,
        default=None,
        help='Custom pod name (default: auto-generated)'
    )

    return parser.parse_args()


def main():
    """Main deployment function"""
    args = parse_args()

    print("="*70)
    print("🦊 Foxhunt Training Pod Deployment")
    print("="*70)
    print(f"\nTimestamp: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")

    # Get API key
    api_key = get_api_key()
    if not api_key:
        return 1

    # Build configuration
    config = DEFAULT_CONFIG.copy()
    if args.gpu:
        config['gpu_type'] = args.gpu

    # Generate pod name
    if args.name:
        pod_name = args.name
    else:
        timestamp = datetime.now().strftime('%Y%m%d-%H%M%S')
        pod_name = f"foxhunt-{args.model}-{timestamp}"

    # Build training command
    training_cmd = TRAINING_COMMANDS[args.model].format(epochs=args.epochs)

    # Deploy pod
    pod = deploy_pod(
        api_key=api_key,
        name=pod_name,
        docker_cmd=training_cmd,
        config=config,
        dry_run=args.dry_run
    )

    if pod:
        print("\n" + "="*70)
        print("✅ Deployment Complete!")
        print("="*70)
        print(f"\n📝 Save this Pod ID: {pod['id']}")
        print(f"\n🔧 To stop the pod later:")
        print(f"  python -c \"import sys; sys.path.insert(0, 'scripts'); from deploy_runpod_training import stop_pod; stop_pod('{api_key[:10]}...', '{pod['id']}')\"")
        print(f"\n  Or use RunPod console: https://www.runpod.io/console/pods")
        return 0
    elif not args.dry_run:
        print("\n❌ Deployment failed")
        return 1
    else:
        return 0


if __name__ == '__main__':
    sys.exit(main())