#!/usr/bin/env python3
"""End-to-end eval renders: each page as the user sees it with a given build of the extension (or none).

usage: ext_render.py --urls URLS.jsonl --out OUT --variant TAG=EXT_DIR [--variant TAG=EXT_DIR ...] [--light]
URLS.jsonl lines: {"id", "url"}. Writes OUT/<tag>/<id>.jpg (viewport, after the extension settled) and
OUT/<tag>/<id>.json (Nadir status). With --light also OUT/light/<id>.jpg without any extension.
Each variant runs in its own persistent profile, so first-visit behaviour (no site cache) is what gets judged.
"""
import argparse, asyncio, json, os, shutil, time
from pathlib import Path
from playwright.async_api import async_playwright

STATUS_JS = "async () => { const a = await chrome.storage.session.get(null); return Object.values(a).pop() || null; }"


async def render_variant(p, tag, ext, urls, out, args):
    od = out / tag
    od.mkdir(parents=True, exist_ok=True)
    prof = out / f"_profile_{tag}"
    shutil.rmtree(prof, ignore_errors=True)
    # the whole browser (extension pages too) presents a normal Chrome UA: some image hosts refuse "HeadlessChrome"
    launch_args = ["--hide-scrollbars", "--mute-audio", "--disable-features=DisableLoadExtensionCommandLineSwitch", f"--user-agent={args.ua}"]
    if ext:
        launch_args += [f"--disable-extensions-except={ext}", f"--load-extension={ext}"]
    # extensions need the full browser in new-headless mode, not Playwright's headless shell
    ctx = await p.chromium.launch_persistent_context(str(prof), executable_path=args.exe or None, channel=None if args.exe else "chromium", headless=True,
                                                     ignore_default_args=["--disable-extensions"] if ext else [],
                                                     viewport={"width": 1280, "height": 800}, args=launch_args, user_agent=args.ua)
    sw = None
    if ext:
        sw = ctx.service_workers[0] if ctx.service_workers else await ctx.wait_for_event("serviceworker", timeout=20000)
    sem = asyncio.Semaphore(args.conc)

    async def one(rec):
        dst = od / f"{rec['id']}.jpg"
        if dst.exists():
            return
        async with sem:
            pg = await ctx.new_page()
            st = {"url": rec["url"], "t": time.time()}
            try:
                await pg.goto(rec["url"], wait_until="domcontentloaded", timeout=30000)
                try:
                    await pg.wait_for_load_state("load", timeout=10000)
                except Exception:
                    pass
                await pg.wait_for_timeout(args.wait * 1000)
                await pg.screenshot(path=str(dst), type="jpeg", quality=82)
                if sw:
                    try:
                        st["nadir"] = await sw.evaluate(STATUS_JS)
                    except Exception as e:
                        st["nadir_err"] = str(e)[:120]
            except Exception as e:
                st["error"] = str(e)[:200]
            finally:
                st["ok"] = dst.exists() and "error" not in st
                (od / f"{rec['id']}.json").write_text(json.dumps(st))
                await pg.close()
    await asyncio.gather(*[one(r) for r in urls])
    await ctx.close()
    shutil.rmtree(prof, ignore_errors=True)


async def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--urls", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--variant", action="append", default=[], help="TAG=EXT_DIR")
    ap.add_argument("--light", action="store_true")
    ap.add_argument("--exe", default="")
    ap.add_argument("--wait", type=float, default=7.0)
    ap.add_argument("--conc", type=int, default=4)
    ap.add_argument("--ua", default="Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/141.0.0.0 Safari/537.36")
    args = ap.parse_args()
    out = Path(args.out)
    urls = [json.loads(l) for l in open(args.urls) if l.strip()]
    async with async_playwright() as p:
        jobs = [("light", None)] if args.light else []
        jobs += [tuple(v.split("=", 1)) for v in args.variant]
        # variants side by side (separate browsers): the same live page state as far as possible
        await asyncio.gather(*[render_variant(p, tag, ext, urls, out, args) for tag, ext in jobs])
    print("done", len(urls), [t for t, _ in jobs], flush=True)


if __name__ == "__main__":
    asyncio.run(main())