"""Async wrappers around the vast.ai CLI. Cannibalized from dtrain."""

from __future__ import annotations

import asyncio
import json


async def _run(*args: str) -> str | None:
    proc = await asyncio.create_subprocess_exec(
        "vastai", *args,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
    )
    stdout, stderr = await proc.communicate()
    if proc.returncode != 0:
        return None
    return stdout.decode()


async def search_offers(min_vram: int, limit: int = 10) -> list[dict]:
    query = f"gpu_ram>={min_vram} inet_down>=100"
    output = await _run(
        "search", "offers", query, "-o", "dph_total",
        "--limit", str(limit), "--raw",
    )
    if not output:
        return []
    return json.loads(output)


async def rent_instance(offer_id: int, image: str, disk: int = 20) -> str | None:
    return await _run(
        "create", "instance", str(offer_id),
        "--image", image, "--disk", str(disk),
    )


async def get_instances() -> list[dict]:
    output = await _run("show", "instances", "--raw")
    if not output:
        return []
    return json.loads(output)


async def destroy_instance(instance_id: int) -> None:
    await _run("destroy", "instance", str(instance_id))