#!/usr/bin/env python3
"""E2B side of the rollout benchmark: the same metrics as Cua's rollout_bench.

Usage (E2B_API_KEY in the environment; never printed):
  e2b_bench.py cold --n 10 --conc 1 --label serial
  e2b_bench.py channel
  e2b_bench.py snapshot
  e2b_bench.py rollout --n 16 --steps 50
  e2b_bench.py cleanup

One JSON line per sample on stdout; cost estimate appended to e2b-cost.log.
Every sandbox gets a 10-minute timeout backstop and is killed in `finally`.
"""
import argparse, asyncio, json, os, sys, time

import httpx

from e2b import AsyncSandbox

VCPU_S = 0.000014  # $/vCPU-second, https://e2b.dev/pricing (2026-10-19)
GIB_S = 0.0000045  # $/GiB-second
BUDGET = float(os.environ.get("E2B_BUDGET_USD", "4.5"))
COST_LOG = os.environ.get("E2B_COST_LOG", os.path.expanduser("~/e2b-bench/out/e2b-cost.log"))
SHAPE = {"cpu": 2, "mem_gib": 0.5}  # refined from get_info on the first sandbox
LIVE = {}  # sandbox_id -> (sandbox, start time)
STREAM_CMD = "yes 0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcde | head -c {n}"


def spent():
    try:
        return sum(float(l.split()[-1]) for l in open(COST_LOG) if l.strip())
    except FileNotFoundError:
        return 0.0


def charge(sid, secs, what):
    usd = secs * (SHAPE["cpu"] * VCPU_S + SHAPE["mem_gib"] * GIB_S)
    with open(COST_LOG, "a") as f:
        f.write(f"{time.strftime('%FT%T')} {what} {sid} secs={secs:.1f} cpu={SHAPE['cpu']} mem_gib={SHAPE['mem_gib']} usd {usd:.6f}\n")


def guard(extra_sandbox_seconds=0):
    est = extra_sandbox_seconds * (SHAPE["cpu"] * VCPU_S + SHAPE["mem_gib"] * GIB_S)
    if spent() + est > BUDGET:
        sys.exit(f"budget guard: spent ${spent():.3f} + planned ${est:.3f} > ${BUDGET}")


def pct(v, p):
    if not v:
        return None
    s = sorted(v)
    import math
    return s[min(len(s), max(1, math.ceil(p / 100 * len(s)))) - 1]


def stats(metric, v, **kw):
    d = {"metric": metric, "n": len(v), "p50": pct(v, 50), "p95": pct(v, 95), "p99": pct(v, 99),
         "max": pct(v, 100), "mean": (sum(v) / len(v)) if v else None}
    d.update(kw)
    return d


RETRIES = {"n": 0}


async def retry(fn, *a, **kw):
    """One retry for transient transport errors; counted in the output."""
    try:
        return await fn(*a, **kw)
    except Exception:  # noqa: BLE001
        RETRIES["n"] += 1
        return await fn(*a, **kw)


def ms(t):
    return round((time.perf_counter() - t) * 1000, 1)


async def new(template=None):
    sb = await AsyncSandbox.create(template=template, timeout=600)
    LIVE[sb.sandbox_id] = (sb, time.time())
    return sb


async def kill(sb, what="kill"):
    sid = sb.sandbox_id
    try:
        await sb.kill()
    finally:
        _, t0 = LIVE.pop(sid, (None, time.time()))
        charge(sid, time.time() - t0, what)


async def refine_shape(sb):
    try:
        info = await sb.get_info()
        cpu = getattr(info, "cpu_count", None)
        mem = getattr(info, "memory_mb", None)
        if cpu:
            SHAPE["cpu"] = cpu
        if mem:
            SHAPE["mem_gib"] = mem / 1024
        print(json.dumps({"metric": "shape", "cpu": SHAPE["cpu"], "mem_gib": SHAPE["mem_gib"],
                          "template": getattr(info, "template_id", None)}), flush=True)
    except Exception as e:  # noqa: BLE001
        print(json.dumps({"metric": "shape", "error": str(e)}), flush=True)


async def one_cold(i, label, wall, template=None):
    t0 = time.perf_counter()
    line = {"cmd": "create", "label": label, "i": i}
    sb = None
    try:
        sb = await new(template)
        line["create_ms"] = ms(t0)
        t1 = time.perf_counter()
        r = await sb.commands.run("true")
        assert r.exit_code == 0
        line["first_exec_ms"] = ms(t1)
        line["ready_ms"] = ms(t0)
        line["ok"] = True
    except Exception as e:  # noqa: BLE001
        line["ok"] = False
        line["error"] = f"{type(e).__name__}: {e}"[:300]
    line["done_at_ms"] = ms(wall)
    print(json.dumps(line), flush=True)
    return sb, line


async def cold(a):
    guard(a.n * 30)
    wall = time.perf_counter()
    sem = asyncio.Semaphore(a.conc)
    out = []

    async def task(i):
        async with sem:
            sb, line = await one_cold(i, a.label, wall, a.template)
            if sb:
                await kill(sb, "cold")
            out.append(line)

    await asyncio.gather(*(task(i) for i in range(a.n)))
    ok = [l["ready_ms"] for l in out if l.get("ok")]
    print(json.dumps(dict(stats("ready_ms", ok), summary=True, cmd="create", label=a.label,
                          failed=len(out) - len(ok), n_req=a.n, conc=a.conc, wall_ms=ms(wall),
                          errors=sorted({l.get("error", "")[:120] for l in out if not l.get("ok")}))), flush=True)


async def channel(a):
    guard(300)
    sb = await new()
    await refine_shape(sb)
    res = []
    try:
        for _ in range(5):
            await sb.commands.run("echo hi")
        v = []
        for _ in range(a.iters):
            t = time.perf_counter()
            r = await sb.commands.run("echo hi")
            v.append(ms(t))
            assert r.exit_code == 0
        res.append(stats("exec_sh_echo_ms", v))
        print(json.dumps(res[-1]), flush=True)
        for size, reps in ((1 << 20, 20), (100 << 20, 3)):
            data = os.urandom(size)
            up, down = [], []
            for r in range(reps):
                p = f"/tmp/bench-{size}-{r}.bin"
                t = time.perf_counter()
                if size >= (64 << 20):
                    # The SDK's own HTTP/2 transport is refused for large bodies
                    # (too_many_data_frames); use the SDK's signed file URLs over HTTP/1.1.
                    async with httpx.AsyncClient(timeout=600) as c:
                        rr = await c.post(sb.upload_url(p), files={"file": ("f", data)})
                        rr.raise_for_status()
                        up.append(ms(t))
                        t = time.perf_counter()
                        rr = await c.get(sb.download_url(p))
                        rr.raise_for_status()
                        got = rr.content
                        down.append(ms(t))
                else:
                    await retry(sb.files.write, p, data)
                    up.append(ms(t))
                    t = time.perf_counter()
                    got = await retry(sb.files.read, p, format="bytes")
                    down.append(ms(t))
                assert len(got) == size
                await sb.commands.run(f"rm -f {p}")
            mb = size / (1 << 20)
            for d, v in (("upload", up), ("download", down)):
                s = stats(f"{d}_{int(mb)}MB_ms", v, MBps_p50=mb / (pct(v, 50) / 1000), via="signed-url-http1" if size >= (64 << 20) else "sdk")
                res.append(s)
                print(json.dumps(s), flush=True)
        for n in (64 << 20, 512 << 20):
            got = 0
            first = None
            t = time.perf_counter()

            def on_out(chunk):
                nonlocal got, first
                if first is None:
                    first = ms(t)
                got += len(chunk)

            await sb.commands.run(STREAM_CMD.format(n=n), on_stdout=on_out, timeout=600)
            secs = time.perf_counter() - t
            s = {"metric": "stream_stdout", "bytes": got, "secs": secs, "MBps": got / (1 << 20) / secs,
                 "first_byte_ms": first}
            res.append(s)
            print(json.dumps(s), flush=True)
    finally:
        await kill(sb, "channel")
    print(json.dumps({"summary": True, "cmd": "channel", "results": res, "transport_retries": RETRIES["n"]}), flush=True)


async def snapshot(a):
    guard(900)
    sb = await new()
    await refine_shape(sb)
    out = {}
    snap_id = None
    forks = []
    try:
        t = time.perf_counter()
        r = await sb.commands.run(
            "python3 -m venv /home/user/venv && /home/user/venv/bin/pip install -q numpy pandas"
            " && dd if=/dev/urandom of=/home/user/state.bin bs=1M count=300 status=none && sync"
            " && du -sh /home/user", timeout=600)
        out["setup_ms"] = ms(t)
        out["state"] = r.stdout.strip()[-60:]
        t = time.perf_counter()
        snap = await sb.create_snapshot()
        out["snapshot_ms"] = ms(t)
        snap_id = getattr(snap, "snapshot_id", None) or getattr(snap, "id", None)
        out["snapshot_id_present"] = bool(snap_id)
        # Restore = new sandbox from the snapshot, to first exec.
        for i in range(3):
            t = time.perf_counter()
            f = await new(snap_id)
            r = await f.commands.run("test -s /home/user/state.bin && /home/user/venv/bin/python -c 'import numpy'")
            out.setdefault("restore_to_exec_ms", []).append(ms(t))
            assert r.exit_code == 0
            forks.append(f)
        # Fork of the running sandbox.
        try:
            t = time.perf_counter()
            fs = await sb.fork(count=4, timeout=120)
            out["fork4_ms"] = ms(t)
            for f in fs:
                if isinstance(f, Exception):
                    out.setdefault("fork_errors", []).append(str(f)[:200])
                else:
                    LIVE[f.sandbox_id] = (f, time.time())
                    forks.append(f)
        except Exception as e:  # noqa: BLE001
            out["fork_error"] = f"{type(e).__name__}: {e}"[:300]
        # Pause / resume of the source.
        t = time.perf_counter()
        await sb.pause()
        out["pause_ms"] = ms(t)
        _, t0 = LIVE[sb.sandbox_id]
        charge(sb.sandbox_id, time.time() - t0, "snapshot-src-until-pause")
        t = time.perf_counter()
        sb2 = await AsyncSandbox.connect(sb.sandbox_id, timeout=300)
        r = await sb2.commands.run("test -s /home/user/state.bin")
        out["resume_to_exec_ms"] = ms(t)
        LIVE[sb.sandbox_id] = (sb2, time.time())
        sb = sb2
    except Exception as e:  # noqa: BLE001
        out["error"] = f"{type(e).__name__}: {e}"[:400]
    finally:
        for f in forks:
            await kill(f, "snapshot-fork")
        await kill(sb, "snapshot-src")
        if snap_id:
            try:
                await AsyncSandbox.delete_snapshot(snap_id)
                out["snapshot_deleted"] = True
            except Exception as e:  # noqa: BLE001
                out["snapshot_delete_error"] = str(e)[:200]
    print(json.dumps(dict(out, summary=True, cmd="snapshot")), flush=True)


async def rollout(a):
    guard(a.n * 120)
    sbs = []
    try:
        res = await asyncio.gather(*(new() for _ in range(a.n)), return_exceptions=True)
        sbs = [s for s in res if not isinstance(s, Exception)]
        errs = [str(e)[:120] for e in res if isinstance(e, Exception)]
        payload = b"a" * a.payload
        wall = time.perf_counter()

        async def loop(i, sb):
            st, ex, wr, rd, err = [], [], [], [], 0
            for s in range(a.steps):
                t = time.perf_counter()
                try:
                    t1 = time.perf_counter()
                    r = await sb.commands.run(
                        f"cd /tmp && ls >/dev/null && printf '{{\"reward\": {s % 2}.{s}}}' > reward-{i}.json")
                    assert r.exit_code == 0
                    ex.append(ms(t1))
                    t2 = time.perf_counter()
                    await sb.files.write(f"/tmp/action-{i}-{s}.json", payload)  # rewriting one path: permission denied on E2B
                    wr.append(ms(t2))
                    t3 = time.perf_counter()
                    b = await sb.files.read(f"/tmp/reward-{i}.json")
                    assert b
                    rd.append(ms(t3))
                    st.append(ms(t))
                except Exception:  # noqa: BLE001
                    err += 1
            return st, ex, wr, rd, err

        outs = await asyncio.gather(*(loop(i, s) for i, s in enumerate(sbs)))
        secs = time.perf_counter() - wall
        agg = [[], [], [], [], 0]
        for o in outs:
            for k in range(4):
                agg[k].extend(o[k])
            agg[4] += o[4]
        for l, v in zip(("step_ms", "exec_ms", "write_ms", "read_reward_ms"), agg[:4]):
            print(json.dumps(stats(l, v)), flush=True)
        print(json.dumps(dict(stats("step_ms", agg[0]), summary=True, cmd="rollout", label=a.label,
                              sandboxes=len(sbs), create_errors=errs, steps_ok=len(agg[0]),
                              errors=agg[4], secs=secs, steps_per_sec=len(agg[0]) / secs)), flush=True)
    finally:
        await asyncio.gather(*(kill(s, "rollout") for s in sbs), return_exceptions=True)


async def forkmany(a):
    """Fork a running sandbox (venv + numpy + 300 MB state) into N copies in one
    call, then time until every fork answers an exec that sees the state."""
    guard(a.n * 120 + 600)
    sb = await new()
    await refine_shape(sb)
    out = {"n": a.n}
    forks = []
    try:
        t = time.perf_counter()
        await sb.commands.run(
            "python3 -m venv /home/user/venv && /home/user/venv/bin/pip install -q numpy"
            " && dd if=/dev/urandom of=/home/user/state.bin bs=1M count=300 status=none && sync", timeout=600)
        out["setup_ms"] = ms(t)
        for rep in range(a.iters):
            t = time.perf_counter()
            res = await sb.fork(count=a.n, timeout=300)
            call_ms = ms(t)
            ok = [f for f in res if not isinstance(f, Exception)]
            for f in ok:
                LIVE[f.sandbox_id] = (f, time.time())
                forks.append(f)

            async def first_exec(f):
                r = await f.commands.run("test -s /home/user/state.bin && echo ok")
                return r.exit_code == 0

            checks = await asyncio.gather(*(first_exec(f) for f in ok), return_exceptions=True)
            all_ms = ms(t)
            row = {"rep": rep, "fork_call_ms": call_ms, "all_ready_ms": all_ms, "ok": sum(1 for c in checks if c is True),
                   "errors": [str(e)[:160] for e in res if isinstance(e, Exception)][:3]}
            out.setdefault("reps", []).append(row)
            print(json.dumps(dict(row, cmd="forkmany", n=a.n)), flush=True)
            await asyncio.gather(*(kill(f, "forkmany") for f in ok), return_exceptions=True)
            forks.clear()
    except Exception as e:  # noqa: BLE001
        out["error"] = f"{type(e).__name__}: {e}"[:400]
    finally:
        await asyncio.gather(*(kill(f, "forkmany") for f in forks), return_exceptions=True)
        await kill(sb, "forkmany-src")
    print(json.dumps(dict(out, summary=True, cmd="forkmany")), flush=True)


async def cleanup(_a):
    n = 0
    paginator = AsyncSandbox.list()
    while True:
        items = await paginator.next_items()
        for it in items:
            await AsyncSandbox.kill(it.sandbox_id)
            n += 1
        if not paginator.has_next:
            break
    print(json.dumps({"cmd": "cleanup", "killed": n, "spent_estimate_usd": round(spent(), 4)}))


def main():
    p = argparse.ArgumentParser()
    p.add_argument("cmd")
    p.add_argument("--n", type=int, default=1)
    p.add_argument("--conc", type=int, default=1)
    p.add_argument("--iters", type=int, default=200)
    p.add_argument("--steps", type=int, default=50)
    p.add_argument("--payload", type=int, default=4096)
    p.add_argument("--label", default="")
    p.add_argument("--template", default=None)
    a = p.parse_args()
    if not os.environ.get("E2B_API_KEY"):
        sys.exit("E2B_API_KEY not set")
    fn = {"cold": cold, "channel": channel, "snapshot": snapshot, "rollout": rollout, "cleanup": cleanup, "forkmany": forkmany}[a.cmd]
    try:
        asyncio.run(fn(a))
    finally:
        print(json.dumps({"spent_estimate_usd": round(spent(), 4)}), file=sys.stderr)


if __name__ == "__main__":
    main()