"""Click through the UI in headless Chrome and save screenshots of every tab after a run.

Usage: python scripts/ui_check.py [--url http://127.0.0.1:7860] [--out shots/] [--dark]
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

from playwright.sync_api import sync_playwright

TABS = ["① Contrastive", "② Guided", "③ Parallel"]


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--url", default="http://127.0.0.1:7860")
    ap.add_argument("--out", default="shots")
    ap.add_argument("--dark", action="store_true")
    ap.add_argument("--width", type=int, default=1440)
    ap.add_argument("--timeout", type=int, default=240)
    ap.add_argument("--no-run", action="store_true", help="only screenshot the gallery traces (no GPU use)")
    args = ap.parse_args()
    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)
    suffix = "-dark" if args.dark else ""
    url = args.url + ("/?__theme=dark" if args.dark else "/?__theme=light")
    errors: list[str] = []
    with sync_playwright() as p:
        browser = p.chromium.launch(channel="chrome", headless=True)
        page = browser.new_page(viewport={"width": args.width, "height": 1000})
        page.on("console", lambda m: errors.append(f"console.{m.type}: {m.text}") if m.type == "error" else None)
        page.on("pageerror", lambda e: errors.append(f"pageerror: {e}"))
        page.goto(url, wait_until="load", timeout=120_000)
        page.get_by_role("tab", name="① Contrastive").wait_for(timeout=120_000)
        page.wait_for_timeout(2500)
        page.screenshot(path=str(out / f"0-overview{suffix}.png"), full_page=True)
        for i, tab in enumerate(TABS, start=1):
            page.get_by_role("tab", name=tab).click()
            page.wait_for_timeout(500)
            panel = page.locator(".tabitem:visible")
            if not args.no_run:
                panel.get_by_role("button", name="Run", exact=True).click()
            panel.locator(".dd-card").first.wait_for(state="visible", timeout=args.timeout * 1000)
            page.wait_for_timeout(1200)
            page.screenshot(path=str(out / f"{i}-{tab.split()[-1].lower()}{suffix}.png"), full_page=True)
            # click a token in the method strip to drive the inspector
            strips = panel.locator(".token[role=button]")
            if strips.count() > 4:
                strips.nth(strips.count() - 3).click()
                page.wait_for_timeout(800)
                page.screenshot(path=str(out / f"{i}-{tab.split()[-1].lower()}-clicked{suffix}.png"), full_page=True)
        browser.close()
    for e in errors:
        print(e)
    print(f"screenshots in {out.resolve()}  ({len(errors)} browser errors)")
    return 0


if __name__ == "__main__":
    sys.exit(main())