# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "runpod>=1.6.0",
#     "click>=8.0.0",
# ]
# ///
"""
Simple RunPod test script to verify API connection and GPU availability.

Usage:
    export RUNPOD_API_KEY="your-api-key"
    uv run scripts/runpod_simple_test.py
    uv run scripts/runpod_simple_test.py --list-gpus
    uv run scripts/runpod_simple_test.py --run-test
"""

import os
import click
import runpod


@click.command()
@click.option("--list-gpus", is_flag=True, help="List available GPU types")
@click.option("--run-test", is_flag=True, help="Run a quick test pod")
def main(list_gpus, run_test):
    """Test RunPod API connection and GPU availability."""
    api_key = os.environ.get("RUNPOD_API_KEY")
    if not api_key:
        raise click.ClickException(
            "Set RUNPOD_API_KEY environment variable.\n"
            "Get your key at: https://runpod.io/console/user/settings"
        )

    runpod.api_key = api_key

    # Test API connection
    click.echo("Testing RunPod API connection...")
    try:
        pods = runpod.get_pods()
        click.echo(f"  Connected! Active pods: {len(pods)}")
    except Exception as e:
        raise click.ClickException(f"API connection failed: {e}")

    # List GPUs
    if list_gpus:
        click.echo("\nAvailable GPU types:")
        try:
            gpus = runpod.get_gpus()
            for gpu in gpus:
                name = gpu.get("id", "Unknown")
                mem = gpu.get("memoryInGb", "?")
                click.echo(f"  - {name} ({mem}GB)")
        except Exception as e:
            click.echo(f"  Could not list GPUs: {e}")

    # Run test pod
    if run_test:
        click.echo("\nLaunching test pod...")
        test_script = "nvidia-smi && python3 -c 'import torch; print(f\"PyTorch: {torch.__version__}, CUDA: {torch.cuda.is_available()}\")'"

        pod = runpod.create_pod(
            name="bamboo-1-test",
            image_name="runpod/pytorch:2.1.0-py3.10-cuda11.8.0-devel-ubuntu22.04",
            gpu_type_id="NVIDIA RTX A4000",
            volume_in_gb=5,
            docker_args=f"bash -c '{test_script}; sleep 60'",
        )

        click.echo(f"  Pod ID: {pod['id']}")
        click.echo(f"  Monitor: https://runpod.io/console/pods")
        click.echo(f"\n  Terminate after checking:")
        click.echo(f"  uv run scripts/runpod_setup.py terminate {pod['id']}")

    if not list_gpus and not run_test:
        click.echo("\nUse --list-gpus to see available GPUs")
        click.echo("Use --run-test to launch a quick test pod")


if __name__ == "__main__":
    main()