"""Render HTML figures to PNG via headless Chromium (Playwright).

Each figure is a self-contained HTML file in /data/cameron/para/paper/figs/html/.
Renders at 2x device pixel ratio for retina-crisp output, then saves to both
paper/figs/generated/ and .agents/reports/paper_figures/media/.

Usage:
    python3 render_html_figures.py --figs 1
    python3 render_html_figures.py --figs all
"""
import argparse
import sys
from pathlib import Path
from playwright.sync_api import sync_playwright

HTML_DIR = Path("/data/cameron/para/paper/figs/html")
OUT_DIRS = [
    Path("/data/cameron/para/paper/figs/generated"),
    Path("/data/cameron/para/.agents/reports/paper_figures/media"),
]

# Per-figure rendering config: HTML file, output PNG name, viewport size
FIGURES = {
    "1": {
        "html": "fig1_overview.html",
        "png":  "fig1_overview.png",
        "width":  1400,
        "height": 640,   # 20% taller for paper-column readability
    },
    "2": {
        "html": "fig2_method.html",
        "png":  "fig2_method.png",
        "width":  1400,
        "height": 520,
    },
    "5": {
        "html": "fig5_video.html",
        "png":  "fig5_video.png",
        "width":  1400,
        "height": 520,
    },
}


def render(fig_key):
    cfg = FIGURES[fig_key]
    html_path = HTML_DIR / cfg["html"]
    if not html_path.exists():
        print(f"  skip {fig_key}: {html_path} not found")
        return

    with sync_playwright() as p:
        browser = p.chromium.launch(headless=True)
        context = browser.new_context(
            viewport={"width": cfg["width"], "height": cfg["height"]},
            device_scale_factor=2,   # retina-crisp
        )
        page = context.new_page()
        page.goto(f"file://{html_path}")
        # Wait for webfonts to load
        page.wait_for_load_state("networkidle")
        page.evaluate("document.fonts.ready")
        page.wait_for_timeout(500)   # safety margin

        # Screenshot just the .figure div (so we get tight cropping)
        figure = page.locator(".figure").first
        png_bytes = figure.screenshot(omit_background=False)

        browser.close()

    for d in OUT_DIRS:
        d.mkdir(parents=True, exist_ok=True)
        (d / cfg["png"]).write_bytes(png_bytes)
    print(f"  wrote {cfg['png']} to {len(OUT_DIRS)} locations")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--figs", default="all")
    args = parser.parse_args()

    if args.figs == "all":
        keys = list(FIGURES.keys())
    else:
        keys = [k.strip() for k in args.figs.split(",")]

    for k in keys:
        if k not in FIGURES:
            print(f"unknown figure: {k}")
            sys.exit(1)
        print(f"=== render fig {k} (HTML) ===")
        render(k)


if __name__ == "__main__":
    main()