"""Minimal RunPod GraphQL client (API key from ~/.runpod/config.toml).

    python scripts/runpod.py balance
    python scripts/runpod.py gpus
    python scripts/runpod.py create NAME --gpu "NVIDIA H100 80GB HBM3" --count 1 [--cloud COMMUNITY] [--volume 300]
    python scripts/runpod.py status POD_ID        # prints ssh target when running
    python scripts/runpod.py stop POD_ID | terminate POD_ID
"""
import argparse
import json
import re
import sys
import time
import urllib.request
from pathlib import Path

IMAGE = 'runpod/pytorch:1.4.1-cu1281-torch2130-ubuntu2404'
KEY_FILE = Path.home() / '.runpod' / 'ssh' / 'runpodctl-ssh-key'


def api_key():
    text = (Path.home() / '.runpod' / 'config.toml').read_text()
    return re.search(r"apikey\s*=\s*['\"]([^'\"]+)", text).group(1)


def gql(query, variables=None):
    req = urllib.request.Request('https://api.runpod.io/graphql',
                                 data=json.dumps({'query': query, 'variables': variables or {}}).encode(),
                                 headers={'Content-Type': 'application/json', 'Authorization': f'Bearer {api_key()}',
                                          'User-Agent': 'rapid-moss'})
    with urllib.request.urlopen(req, timeout=60) as r:
        out = json.load(r)
    if out.get('errors'):
        raise SystemExit(json.dumps(out['errors'], indent=1))
    return out['data']


def pod(pod_id):
    return gql('query($id: String!) { pod(input: {podId: $id}) { id name desiredStatus costPerHr gpuCount '
               'machine { gpuDisplayName location } runtime { uptimeInSeconds ports { ip isIpPublic privatePort '
               'publicPort type } } } }', {'id': pod_id})['pod']


def ssh_target(p):
    for port in (p.get('runtime') or {}).get('ports') or []:
        if port['privatePort'] == 22 and port['isIpPublic']:
            return port['ip'], port['publicPort']
    return None


def main():
    ap = argparse.ArgumentParser()
    sub = ap.add_subparsers(dest='cmd', required=True)
    sub.add_parser('balance')
    sub.add_parser('gpus')
    c = sub.add_parser('create')
    c.add_argument('name')
    c.add_argument('--gpu', default='NVIDIA H100 80GB HBM3')
    c.add_argument('--count', type=int, default=1)
    c.add_argument('--cloud', default='COMMUNITY')
    c.add_argument('--volume', type=int, default=300)
    c.add_argument('--disk', type=int, default=60)
    for name in ('status', 'stop', 'terminate', 'wait'):
        sub.add_parser(name).add_argument('pod_id')
    a = ap.parse_args()
    if a.cmd == 'balance':
        print(json.dumps(gql('query { myself { clientBalance currentSpendPerHr pods { id name desiredStatus costPerHr } } }')['myself'], indent=1))
    elif a.cmd == 'gpus':
        data = gql('query { gpuTypes { id memoryInGb securePrice communityPrice lowestPrice(input:{gpuCount:1}) { stockStatus } } }')
        for g in sorted(data['gpuTypes'], key=lambda g: -(g['memoryInGb'] or 0)):
            if (g['memoryInGb'] or 0) >= 80:
                print(f"{g['id']:<60} {g['memoryInGb']:>4} GB secure {g['securePrice']} community {g['communityPrice']} stock {(g['lowestPrice'] or {}).get('stockStatus')}")
    elif a.cmd == 'create':
        pub = KEY_FILE.with_suffix('.pub').read_text().strip()
        q = ('mutation($in: PodFindAndDeployOnDemandInput) { podFindAndDeployOnDemand(input: $in) '
             '{ id costPerHr machine { gpuDisplayName location } } }')
        inp = {'cloudType': a.cloud, 'gpuCount': a.count, 'volumeInGb': a.volume, 'containerDiskInGb': a.disk,
               'gpuTypeId': a.gpu, 'name': a.name, 'imageName': IMAGE, 'ports': '22/tcp',
               'volumeMountPath': '/workspace', 'startSsh': True, 'supportPublicIp': True,
               'minCudaVersion': '12.8', 'env': [{'key': 'PUBLIC_KEY', 'value': pub}]}
        print(json.dumps(gql(q, {'in': inp})['podFindAndDeployOnDemand'], indent=1))
    elif a.cmd == 'status':
        p = pod(a.pod_id)
        print(json.dumps(p, indent=1))
        t = ssh_target(p)
        if t:
            print(f'ssh -i {KEY_FILE} -p {t[1]} root@{t[0]}')
    elif a.cmd == 'wait':
        for _ in range(120):
            t = ssh_target(pod(a.pod_id))
            if t:
                print(f'{t[0]} {t[1]}')
                return
            time.sleep(10)
        sys.exit('pod did not expose ssh in 20 min')
    elif a.cmd == 'stop':
        print(gql('mutation($id: String!) { podStop(input: {podId: $id}) { id desiredStatus } }', {'id': a.pod_id}))
    elif a.cmd == 'terminate':
        print(gql('mutation($id: String!) { podTerminate(input: {podId: $id}) }', {'id': a.pod_id}))


if __name__ == '__main__':
    main()