#!/usr/bin/env python3
"""RunPod orchestration: status, deploy, stop, download results.

Single source of truth for all RunPod training state. Run this before and
after training to make sure nothing is left running.

Usage:
    python runpod_status.py                # Show all pods + volumes + cost
    python runpod_status.py deploy         # Deploy a new training pod
    python runpod_status.py stop           # Stop all running pods
    python runpod_status.py stop POD_ID    # Stop a specific pod
    python runpod_status.py download       # Download results from volume
    python runpod_status.py terminate      # Terminate all stopped pods
    python runpod_status.py ssh            # Print SSH command for running pod
    python runpod_status.py logs           # Tail training log on running pod
"""

import json
import os
import subprocess
import sys
import urllib.error
import urllib.request
from datetime import timedelta
from pathlib import Path

# ── Config ────────────────────────────────────────────────────────────────────

VOLUME_ID = "nxq4xshgcr"
VOLUME_NAME = "checkbox-training"
VOLUME_MOUNT = "/runpod-volume"
DATACENTER = "US-MD-1"

GPU_TYPE = "NVIDIA A100-SXM4-80GB"
CONTAINER_IMAGE = "runpod/pytorch:2.4.0-py3.11-cuda12.4.1-devel-ubuntu22.04"
CONTAINER_DISK_GB = 50
POD_NAME_PREFIX = "checkbox-train"

LOCAL_OUTPUT = Path(__file__).resolve().parent.parent / "cloud_training_output"

SSH_PUBLIC_KEY = (
    "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQDQzXqYhqDWN4uFdgNIS1w2fQB2Cq9OgCinb2uahk5K4Aj"
    "VmdQvuQdApnU/WtegEtK46eNF8NQMXN093wMKj5S6ByHEyVxLuYRPdKQy8rtPmpflos8TsJjUbt5kE8lvCd9"
    "1RLM5bCJbLORs5eT5lyTr4xruUwdrMunso1U2sUJsJgpMj4iM/baYSZG3Kz3X9nqc3neH6SeMqThUiEqeN3W"
    "Jc4Z9Kt4m3yOtAWRERtzxPOkaynWJaZhwCBpZtoYXpDDJn38bnpxml59Zq7AKG3COtdUijp0feznx8rFQUKNL"
    "ffqxLhmvFA9aM9maAq26PBR6psG2bN/IYLcOfIL/jMav RunPod-Key-Go"
)


# ── API helpers ───────────────────────────────────────────────────────────────


def _get_api_key() -> str:
    """Read RunPod API key from ~/.runpod/config.toml."""
    config = Path.home() / ".runpod" / "config.toml"
    if not config.exists():
        print("ERROR: No RunPod API key. Run: runpodctl config")
        sys.exit(1)
    for line in config.read_text().splitlines():
        if "apikey" in line.lower() and "=" in line:
            return line.split("=", 1)[1].strip().strip("'\"")
    print("ERROR: Could not parse API key from ~/.runpod/config.toml")
    sys.exit(1)


def _graphql(query: str) -> dict:
    """Execute a RunPod GraphQL query."""
    api_key = _get_api_key()
    url = f"https://api.runpod.io/graphql?api_key={api_key}"
    # Collapse whitespace — RunPod API rejects multiline queries
    clean_query = " ".join(query.split())
    req = urllib.request.Request(
        url,
        data=json.dumps({"query": clean_query}).encode(),
        headers={
            "Content-Type": "application/json",
            "User-Agent": "runpod-status/1.0",
        },
    )
    with urllib.request.urlopen(req, timeout=30) as resp:
        result = json.loads(resp.read())
    if "errors" in result:
        print(f"API ERROR: {result['errors'][0]['message']}")
        sys.exit(1)
    return result["data"]


# ── SSH helper ────────────────────────────────────────────────────────────────


def _get_ssh_info(pod: dict) -> tuple:
    """Extract (ip, port) from pod runtime info. Returns (None, None) if unavailable."""
    runtime = pod.get("runtime") or {}
    for port in runtime.get("ports") or []:
        if port.get("privatePort") == 22 and port.get("isIpPublic"):
            return port["ip"], port["publicPort"]
    return None, None


def _ssh_cmd(ip: str, port: int) -> str:
    return f"ssh -o StrictHostKeyChecking=no -p {port} root@{ip}"


# ── Commands ──────────────────────────────────────────────────────────────────


def cmd_status():
    """Show all pods, volumes, and estimated cost."""
    data = _graphql("""
    {
        myself {
            pods {
                id name desiredStatus costPerHr
                gpuCount
                runtime {
                    uptimeInSeconds
                    ports { ip isIpPublic privatePort publicPort type }
                    gpus { gpuUtilPercent memoryUtilPercent }
                }
            }
            networkVolumes {
                id name size dataCenterId
            }
        }
    }
    """)

    pods = data["myself"]["pods"]
    volumes = data["myself"]["networkVolumes"]

    print("=" * 60)
    print("  RunPod Status")
    print("=" * 60)

    # Volumes
    print(f"\nVolumes ({len(volumes)}):")
    if not volumes:
        print("  (none)")
    for vol in volumes:
        cost = vol["size"] * 0.05
        print(f"  {vol['name']} ({vol['id']}): {vol['size']}GB in {vol['dataCenterId']} — ${cost:.2f}/mo")

    # Pods
    print(f"\nPods ({len(pods)}):")
    if not pods:
        print("  (none)")

    total_hourly = 0
    for pod in pods:
        status = pod["desiredStatus"]
        name = pod["name"]
        cost = pod["costPerHr"] or 0

        runtime = pod.get("runtime") or {}
        uptime_s = runtime.get("uptimeInSeconds") or 0
        uptime = str(timedelta(seconds=int(uptime_s))) if uptime_s else "-"
        spent = cost * uptime_s / 3600 if uptime_s else 0

        gpu_util = "-"
        gpus = runtime.get("gpus") or []
        if gpus:
            gpu_util = f"{gpus[0].get('gpuUtilPercent', 0)}%"

        ip, port = _get_ssh_info(pod)
        ssh = f"ssh -p {port} root@{ip}" if ip else "-"

        marker = " *** RUNNING ***" if status == "RUNNING" else ""
        print(f"\n  {name} ({pod['id']}){marker}")
        print(f"    Status: {status}  |  Cost: ${cost:.2f}/hr  |  Uptime: {uptime}  |  Spent: ${spent:.2f}")
        print(f"    GPU util: {gpu_util}  |  SSH: {ssh}")

        if status == "RUNNING":
            total_hourly += cost

    if total_hourly > 0:
        print(f"\n  >>> ACTIVE COST: ${total_hourly:.2f}/hr <<<")
    else:
        print(f"\n  No running pods.")

    print()


def cmd_deploy():
    """Deploy a new training pod with network volume."""
    # Check for existing running pods first
    data = _graphql('{ myself { pods { id name desiredStatus costPerHr } } }')
    running = [p for p in data["myself"]["pods"] if p["desiredStatus"] == "RUNNING"]
    if running:
        print("WARNING: You already have running pod(s):")
        for p in running:
            print(f"  {p['name']} ({p['id']}) — ${p['costPerHr']:.2f}/hr")
        resp = input("Deploy another? [y/N] ")
        if resp.lower() != "y":
            print("Aborted.")
            return

    print(f"Deploying {GPU_TYPE} with volume {VOLUME_NAME}...")

    # Escape the SSH key for GraphQL
    escaped_key = SSH_PUBLIC_KEY.replace('"', '\\"')

    data = _graphql(f"""
    mutation {{
        podFindAndDeployOnDemand(input: {{
            name: "{POD_NAME_PREFIX}"
            imageName: "{CONTAINER_IMAGE}"
            gpuTypeId: "{GPU_TYPE}"
            gpuCount: 1
            containerDiskInGb: {CONTAINER_DISK_GB}
            networkVolumeId: "{VOLUME_ID}"
            volumeMountPath: "{VOLUME_MOUNT}"
            ports: "8888/http,22/tcp"
            env: [{{key: "PUBLIC_KEY", value: "{escaped_key}"}}]
        }}) {{
            id desiredStatus costPerHr
        }}
    }}
    """)

    pod = data["podFindAndDeployOnDemand"]
    print(f"Pod deployed: {pod['id']} (${pod['costPerHr']:.2f}/hr)")
    print("Waiting for SSH...")

    # Poll for SSH to come up
    import time
    for _ in range(30):
        time.sleep(5)
        info = _graphql(f"""
        {{ pod(input: {{podId: "{pod['id']}"}}) {{
            runtime {{ ports {{ ip isIpPublic privatePort publicPort }} }}
        }} }}
        """)
        ip, port = _get_ssh_info(info["pod"])
        if ip:
            print(f"\nReady! SSH: ssh -p {port} root@{ip}")
            return
    print("Pod deployed but SSH not ready yet. Run 'python runpod_status.py' to check.")


def cmd_stop(pod_id=None):
    """Stop running pod(s)."""
    if pod_id:
        data = _graphql(f'mutation {{ podStop(input: {{podId: "{pod_id}"}}) {{ id desiredStatus }} }}')
        print(f"Stopped: {pod_id}")
        return

    # Stop all running pods
    data = _graphql('{ myself { pods { id name desiredStatus } } }')
    running = [p for p in data["myself"]["pods"] if p["desiredStatus"] == "RUNNING"]
    if not running:
        print("No running pods.")
        return

    for pod in running:
        _graphql(f'mutation {{ podStop(input: {{podId: "{pod["id"]}"}}) {{ id }} }}')
        print(f"Stopped: {pod['name']} ({pod['id']})")


def cmd_terminate(pod_id=None):
    """Terminate (delete) stopped pod(s)."""
    data = _graphql('{ myself { pods { id name desiredStatus } } }')

    if pod_id:
        targets = [p for p in data["myself"]["pods"] if p["id"] == pod_id]
    else:
        targets = [p for p in data["myself"]["pods"] if p["desiredStatus"] == "EXITED"]

    if not targets:
        print("No pods to terminate.")
        return

    for pod in targets:
        if pod["desiredStatus"] == "RUNNING":
            print(f"Skipping {pod['name']} ({pod['id']}) — still running. Stop it first.")
            continue
        _graphql(f'mutation {{ podTerminate(input: {{podId: "{pod["id"]}"}}) }}')
        print(f"Terminated: {pod['name']} ({pod['id']})")


def cmd_ssh():
    """Print SSH command for the running pod."""
    data = _graphql('{ myself { pods { id name desiredStatus runtime { ports { ip isIpPublic privatePort publicPort } } } } }')
    for pod in data["myself"]["pods"]:
        if pod["desiredStatus"] == "RUNNING":
            ip, port = _get_ssh_info(pod)
            if ip:
                print(_ssh_cmd(ip, port))
                return
    print("No running pods with SSH.")


def cmd_progress():
    """Show training progress: epoch, loss, cost."""
    data = _graphql("""{ myself { pods { id name desiredStatus costPerHr
        runtime { uptimeInSeconds ports { ip isIpPublic privatePort publicPort }
        gpus { gpuUtilPercent memoryUtilPercent } } } } }""")
    for pod in data["myself"]["pods"]:
        if pod["desiredStatus"] != "RUNNING":
            continue
        ip, port = _get_ssh_info(pod)
        if not ip:
            continue

        uptime_s = (pod.get("runtime") or {}).get("uptimeInSeconds") or 0
        cost = (pod.get("costPerHr") or 0) * uptime_s / 3600
        gpus = (pod.get("runtime") or {}).get("gpus") or []
        gpu_util = f"{gpus[0]['gpuUtilPercent']}%" if gpus else "?"

        result = subprocess.run(
            ["ssh", "-o", "StrictHostKeyChecking=no", "-o", "ConnectTimeout=10",
             "-p", str(port), f"root@{ip}",
             # Epoch line (completed) + last validation "all" line
             r"grep -P '^\s+\d+/\d+.*\d+:\d+$' /workspace/train.log | tail -1; "
             r"grep -P '^\s+all\s' /workspace/train.log | tail -1"],
            capture_output=True, text=True, timeout=15,
        )
        lines = [l.strip() for l in result.stdout.strip().split("\n") if l.strip()]

        epoch = "starting..."
        metrics = ""
        for line in lines:
            parts = line.split()
            if parts and "/" in parts[0] and parts[0][0].isdigit():
                epoch = parts[0]
            elif parts and parts[0] == "all":
                # "all  4026  30946  0.853  0.707  0.778  0.416"
                try:
                    p, r, map50, map50_95 = parts[3], parts[4], parts[5], parts[6]
                    metrics = f"P={p} R={r} mAP50={map50} mAP50-95={map50_95}"
                except IndexError:
                    pass

        # Get WandB run URL from log
        wandb_result = subprocess.run(
            ["ssh", "-o", "StrictHostKeyChecking=no", "-o", "ConnectTimeout=10",
             "-p", str(port), f"root@{ip}",
             r"grep -oP 'https://wandb\.ai/\S+' /workspace/train.log | tail -1"],
            capture_output=True, text=True, timeout=15,
        )
        wandb_url = wandb_result.stdout.strip()

        print(f"Epoch: {epoch}  |  GPU: {gpu_util}  |  Spent: ${cost:.2f}")
        if metrics:
            print(f"Val:   {metrics}")
        if wandb_url:
            print(f"WandB: {wandb_url}")
        return
        return

    print("No running pods.")


def cmd_logs():
    """Tail training log on the running pod."""
    data = _graphql('{ myself { pods { id name desiredStatus runtime { ports { ip isIpPublic privatePort publicPort } } } } }')
    for pod in data["myself"]["pods"]:
        if pod["desiredStatus"] == "RUNNING":
            ip, port = _get_ssh_info(pod)
            if ip:
                os.execvp("ssh", [
                    "ssh", "-o", "StrictHostKeyChecking=no",
                    "-p", str(port), f"root@{ip}",
                    "tail -f /workspace/train.log",
                ])
    print("No running pods.")


def cmd_download():
    """Download results from the network volume via a running pod."""
    data = _graphql('{ myself { pods { id name desiredStatus runtime { ports { ip isIpPublic privatePort publicPort } } } } }')
    running = [p for p in data["myself"]["pods"] if p["desiredStatus"] in ("RUNNING", "EXITED")]

    # Try running pods first, then stopped ones
    for status in ("RUNNING", "EXITED"):
        for pod in running:
            if pod["desiredStatus"] != status:
                continue
            ip, port = _get_ssh_info(pod)
            if not ip:
                continue

            LOCAL_OUTPUT.mkdir(parents=True, exist_ok=True)
            print(f"Downloading from {pod['name']} ({pod['id']})...")
            subprocess.run([
                "scp", "-o", "StrictHostKeyChecking=no",
                "-P", str(port), "-r",
                f"root@{ip}:{VOLUME_MOUNT}/output/.",
                str(LOCAL_OUTPUT),
            ])
            print(f"Results saved to: {LOCAL_OUTPUT}")

            # Also grab the training log
            subprocess.run([
                "scp", "-o", "StrictHostKeyChecking=no",
                "-P", str(port),
                f"root@{ip}:/workspace/train.log",
                str(LOCAL_OUTPUT / "train.log"),
            ], capture_output=True)

            print(f"\nFiles:")
            for f in sorted(LOCAL_OUTPUT.iterdir()):
                size = f.stat().st_size
                if size > 1_000_000:
                    print(f"  {f.name:40s} {size/1_000_000:.1f} MB")
                else:
                    print(f"  {f.name:40s} {size/1_000:.1f} KB")
            return

    print("No reachable pods. Start one to access the volume.")


# ── Main ──────────────────────────────────────────────────────────────────────

COMMANDS = {
    "status": cmd_status,
    "progress": cmd_progress,
    "deploy": cmd_deploy,
    "stop": cmd_stop,
    "terminate": cmd_terminate,
    "ssh": cmd_ssh,
    "logs": cmd_logs,
    "download": cmd_download,
}

if __name__ == "__main__":
    args = sys.argv[1:]
    cmd = args[0] if args else "status"

    if cmd in ("-h", "--help", "help"):
        print(__doc__)
        sys.exit(0)

    if cmd not in COMMANDS:
        print(f"Unknown command: {cmd}")
        print(f"Available: {', '.join(COMMANDS)}")
        sys.exit(1)

    if cmd == "stop" and len(args) > 1:
        cmd_stop(args[1])
    elif cmd == "terminate" and len(args) > 1:
        cmd_terminate(args[1])
    else:
        COMMANDS[cmd]()
