"""Chrome vs Lightpanda vs Obscura through the same agent-browser 0.39.3 harness, on 50 real pages.
Per page and engine: open, wait for network idle (15 s cap), then save the rendered HTML and `snapshot -i`,
time both, and sum the memory of the engine's process tree. One fresh session per page.
Usage: python3 bench.py <run> [engine ...]"""
import json, os, pathlib, re, subprocess, sys, time

AB = ["npx", "agent-browser"]
BIN = pathlib.Path(__file__).parent / "bin"
ENGINES = {
    "chrome": {},
    "lightpanda": {"AGENT_BROWSER_ENGINE": "lightpanda", "AGENT_BROWSER_EXECUTABLE_PATH": str(BIN / "lightpanda/lightpanda-aarch64-macos")},
    "obscura": {"AGENT_BROWSER_ENGINE": "obscura", "AGENT_BROWSER_EXECUTABLE_PATH": str(BIN / "obscura/obscura")},
}
MARK = {"chrome": "--remote-debugging-port=0", "lightpanda": "lightpanda-aarch64-macos", "obscura": str(BIN / "obscura")}
RUN = sys.argv[1]
ONLY = sys.argv[2:] or list(ENGINES)
PAGES = [l.rstrip("\n").split("\t") for l in open(pathlib.Path(__file__).parent / "pages.tsv")]
# Run from the folder where agent-browser is installed (npm install agent-browser).
OUT = pathlib.Path(__file__).parent / "out" / RUN


def ab(session, env, *args, timeout=60):
    try:
        p = subprocess.run(AB + ["--session", session, *args], capture_output=True, text=True, timeout=timeout, env={**os.environ, **env})
        return p.returncode, p.stdout, p.stderr
    except subprocess.TimeoutExpired:
        return "timeout", "", ""


def engine_rss_mb(engine):
    """Memory of the engine's process tree: the launched root process plus all its descendants."""
    rows = [l.split(None, 3) for l in subprocess.run(["ps", "-axo", "pid=,ppid=,rss=,command="], capture_output=True, text=True).stdout.splitlines()]
    rows = [(int(a), int(b), int(c), d) for a, b, c, d in rows if len([a, b, c, d]) == 4]
    roots = {pid for pid, _, _, cmd in rows if MARK[engine] in cmd}
    tree, changed = set(roots), True
    while changed:
        new = {pid for pid, ppid, _, _ in rows if ppid in tree} - tree
        tree |= new
        changed = bool(new)
    return round(sum(rss for pid, _, rss, _ in rows if pid in tree) / 1024)


for engine in ONLY:
    env = ENGINES[engine]
    d = OUT / engine
    d.mkdir(parents=True, exist_ok=True)
    for i, (group, url) in enumerate(PAGES, 1):
        path = d / f"{i:02d}.json"
        if path.exists():
            continue
        s = f"eng-{RUN}-{engine}-{i}"
        t0 = time.time()
        code, out, err = ab(s, env, "open", url, timeout=75)
        t_open = time.time() - t0
        w_code, _, _ = ab(s, env, "wait", "--load", "networkidle", timeout=25)
        t_ready = time.time() - t0
        _, html, _ = ab(s, env, "eval", "document.documentElement.outerHTML", timeout=40)
        _, snap, _ = ab(s, env, "snapshot", "-i", timeout=40)
        mb = engine_rss_mb(engine)
        _, title, _ = ab(s, env, "get", "title", timeout=20)
        ab(s, env, "close", timeout=30)
        try:
            html = json.loads(html)
        except Exception:
            pass
        (d / f"{i:02d}.html").write_text(html if isinstance(html, str) else "")
        (d / f"{i:02d}.snap.txt").write_text(snap)
        row = {"i": i, "group": group, "url": url, "engine": engine, "open_ok": code == 0 and "✓" in out,
               "open_error": (out + err).strip()[:300] if code != 0 else "", "network_idle": w_code == 0,
               "t_open": round(t_open, 2), "t_ready": round(t_ready, 2), "html_chars": len(html) if isinstance(html, str) else 0,
               "snap_chars": len(snap), "rss_mb": mb, "title": title.strip()[:200]}
        path.write_text(json.dumps(row, indent=1))
        print(f"{RUN} {engine:<10} {i:02d} {url[:55]:<55} ok={row['open_ok']!s:<5} idle={row['network_idle']!s:<5} {row['t_ready']:>6}s snap={row['snap_chars']:>7} mem={mb:>5}MB", flush=True)