"""pack.py ON TYPE N CONC LABEL [extra create args]

Packing run on the elastic pool (released cua CLI): N creates of
`cua spaces create IMAGE --on ON --instance-type TYPE --name pk-LABEL-i --json [extra]`,
CONC at a time; then `cua sb exec <ref> -- echo ok` in every created Space,
counted per host; then every Space is deleted. IMAGE: env CUA_BENCH_IMAGE
(default ghcr.io/trycua/linux:24.04-slim). One JSON line per Space, then a summary.
"""
import collections
import concurrent.futures as cf
import json
import os
import subprocess
import sys
import time

ON, TYPE, N, CONC, LABEL = sys.argv[1], sys.argv[2], int(sys.argv[3]), int(sys.argv[4]), sys.argv[5]
EXTRA = sys.argv[6:]
IMAGE = os.environ.get("CUA_BENCH_IMAGE", "ghcr.io/trycua/linux:24.04-slim")


def create(i):
    name = f"pk-{LABEL}-{i}"
    t = time.perf_counter()
    p = subprocess.run(["cua", "spaces", "create", IMAGE, "--on", ON, "--instance-type", TYPE,
                        "--name", name, "--json", *EXTRA], capture_output=True, text=True, timeout=1800)
    row = {"i": i, "name": name, "create_ms": round((time.perf_counter() - t) * 1000, 1), "create_rc": p.returncode}
    try:
        s = json.loads(p.stdout.strip().splitlines()[-1])["spaces"][0]
        row.update(ref=s["id"], host=s.get("host_name"))
    except Exception:
        row["error"] = (p.stderr or p.stdout)[-400:]
    return row


def exec_(row):
    if "ref" not in row:
        row["exec_ok"] = False
        return row
    q = subprocess.run(["cua", "sb", "exec", row["ref"], "--", "echo", "ok"], capture_output=True, text=True, timeout=300)
    row["exec_ok"] = q.returncode == 0 and "ok" in q.stdout
    return row


def delete(row):
    d = subprocess.run(["cua", "spaces", "delete", row["name"], "--force"], capture_output=True, text=True, timeout=600)
    row["delete_rc"] = d.returncode
    return row


t0 = time.perf_counter()
with cf.ThreadPoolExecutor(CONC) as ex:
    rows = list(ex.map(create, range(N)))
up_s = time.perf_counter() - t0
with cf.ThreadPoolExecutor(CONC) as ex:
    rows = list(ex.map(exec_, rows))
at_full = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
hosts = subprocess.run(["cua", "cloud", "hosts", "ls"], capture_output=True, text=True).stdout
with cf.ThreadPoolExecutor(CONC) as ex:
    rows = list(ex.map(lambda r: delete(r) if (r.get("create_rc") == 0 or "ref" in r) else r, rows))
for r in rows:
    print(json.dumps(r), flush=True)
per_host = collections.Counter(r.get("host") for r in rows if r.get("exec_ok"))
print(json.dumps({"summary": True, "label": LABEL, "asked": N, "created": sum("ref" in r for r in rows),
                  "exec_ok": sum(bool(r.get("exec_ok")) for r in rows), "per_host": per_host,
                  "up_wall_s": round(up_s, 1), "at_full": at_full, "hosts_ls": hosts.strip()}), flush=True)