"""modal_start.py N [IMAGE] -- Modal Sandbox create -> first exec, serial, in $MODAL_ENVIRONMENT.

Timed per sample (wall clock on the client, time.perf_counter):
  t0 before modal.Sandbox.create(...)
  create_ms       = Sandbox.create returned (sandbox scheduled + running)
  first_exec_ms   = sb.exec("echo", "ok") output read back (from t0)
One warmup sample (index -1, image build/cache fill) is excluded. Each sandbox
is terminated right after its sample; the app is stopped at the end by the caller.
Prints one JSON line per sample, then a summary line.
"""
import json
import os
import sys
import time

ENV = os.environ["MODAL_ENVIRONMENT"]  # your Modal environment
import modal  # noqa: E402

N = int(sys.argv[1]) if len(sys.argv) > 1 else 30
IMG = sys.argv[2] if len(sys.argv) > 2 else "debian_slim"
APP = os.environ.get("MODAL_APP", "cua-bench")

app = modal.App.lookup(APP, create_if_missing=True, environment_name=ENV)
image = modal.Image.debian_slim() if IMG == "debian_slim" else modal.Image.from_registry(IMG)


def one(i):
    sb = None
    t0 = time.perf_counter()
    try:
        sb = modal.Sandbox.create(app=app, image=image, timeout=120, cpu=1.0, memory=1024)
        t1 = time.perf_counter()
        p = sb.exec("echo", "ok")
        out = p.stdout.read()
        p.wait()
        t2 = time.perf_counter()
        return {"i": i, "create_ms": round((t1 - t0) * 1000, 1),
                "first_exec_ms": round((t2 - t0) * 1000, 1),
                "ok": out.strip() == "ok", "sandbox_id": sb.object_id,
                "at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())}
    except Exception as e:  # record, never hide
        return {"i": i, "ok": False, "error": repr(e)[:300]}
    finally:
        if sb is not None:
            try:
                sb.terminate()
            except Exception:
                pass


print(json.dumps({"modal_client": modal.__version__, "image": IMG, "n": N}), flush=True)
w = one(-1)
w["warmup"] = True
print(json.dumps(w), flush=True)
res = []
for i in range(N):
    r = one(i)
    res.append(r)
    print(json.dumps(r), flush=True)
ok = [r for r in res if r.get("ok")]
fe = sorted(r["first_exec_ms"] for r in ok)
print(json.dumps({"summary": True, "n": N, "ok": len(ok),
                  "first_exec_p50": fe[len(fe) // 2] if fe else None}), flush=True)