#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Render CMGUIBENCH HTML predictions to PNG screenshots with Playwright.

Default input:
  html_outputs/pair_0001.html

Default output:
  outputs/pair_0001.png
"""

from __future__ import annotations

import argparse
import asyncio
import json
import os
import re
import uuid
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

from PIL import Image

DATA_ROOT = Path(__file__).resolve().parents[1]
DEFAULT_META = DATA_ROOT / "cmgui_test_500_screenshot_pairs.json"
DEFAULT_IMAGES = DATA_ROOT / "images"
DEFAULT_HTML_DIR = DATA_ROOT / "html_outputs"
DEFAULT_OUT_DIR = DATA_ROOT / "outputs"
DEFAULT_FONTS_ROOT = Path(
    "/data/shared/public/weikai.xu/AAAI2026-AppDeltaWorld/training/training_data/annotation-main/assets/fonts"
)
_EMBEDDED_FONT_CSS: Optional[str] = None


def read_pairs(path: Path) -> List[Dict[str, Any]]:
    obj = json.loads(path.read_text(encoding="utf-8"))
    if isinstance(obj, dict) and isinstance(obj.get("pairs"), list):
        return obj["pairs"]
    if isinstance(obj, list):
        return obj
    raise ValueError(f"Unsupported metadata format: {path}")


def image_size(pair: Dict[str, Any], images_root: Path) -> Tuple[int, int]:
    before = images_root / str(pair["image_before"])
    with Image.open(before) as im:
        return im.size


def wrap_html_if_fragment(html: str) -> str:
    if re.search(r"<html[\s>]", html, re.IGNORECASE):
        return html
    return "<!DOCTYPE html><html><head><meta charset='utf-8'></head><body>" + html + "</body></html>"


def viewport_css(width: int, height: int) -> str:
    return (
        f"html, body {{ margin: 0; padding: 0; width: {width}px; min-height: {height}px; "
        f"overflow: hidden; background: white; }}\n"
        "* { box-sizing: border-box; }"
    )


def inject_viewport_css(html: str, width: int, height: int) -> str:
    # Strip any existing rogue CSS that forces incorrect sizing
    html = re.sub(r"<style>html,body\{width:\d+px!important;height:\d+px!important;max-width:\d+px!important;max-height:\d+px!important;margin:0;padding:0;overflow:hidden;\}#render-target\{width:\d+px!important;height:\d+px!important;\}</style>", "", html)
    
    css = f"\n<style id=\"cmgui-render-css\">\n{viewport_css(width, height)}\n</style>\n"
    if re.search(r"</head>", html, re.IGNORECASE):
        return re.sub(r"</head>", css + "</head>", html, count=1, flags=re.IGNORECASE)
    return css + html


CHROMIUM_ARGS = [
    "--allow-file-access-from-files",
    "--disable-web-security",
    "--no-sandbox",
    "--disable-dev-shm-usage",
]


def strip_file_font_faces(html: str) -> str:
    """Remove injected @font-face rules without deleting layout <style> blocks."""
    html = re.sub(r"@font-face\s*\{[^}]*\}", "", html, flags=re.IGNORECASE)
    # Drop style blocks that only contain @font-face / font-family overrides.
    html = re.sub(
        r"<style[^>]*>\s*(?:@font-face\s*\{[^}]*\}\s*)+"
        r"(?:html,body[^{]*\{[^}]*\}\s*)*</style>",
        "",
        html,
        flags=re.IGNORECASE,
    )
    return html


def inject_fallback_font_stack(html: str) -> str:
    if "cmgui-fallback-fonts" in html:
        return html
    stack = "'Noto Sans SC','WenQuanYi Micro Hei','DejaVu Sans',sans-serif"
    css = f"<style id=\"cmgui-fallback-fonts\">html,body,*{{font-family:{stack}!important}}</style>\n"
    if re.search(r"</head>", html, re.IGNORECASE):
        return re.sub(r"</head>", css + "</head>", html, count=1, flags=re.IGNORECASE)
    return css + html


def should_use_embedded_css() -> bool:
    mode = os.environ.get("CMGUI_RENDER_FONT_MODE", "file").strip().lower()
    return mode in {"base64", "embedded", "data"}


def embedded_font_css_path() -> Optional[Path]:
    for cand in [
        os.environ.get("CMGUI_FONTS_EMBEDDED_CSS", "").strip(),
        str(Path(os.environ.get("CMGUI_FONTS_ROOT", "")) / "embedded.css"),
        str(DATA_ROOT / "third_party" / "cmgui-render-fonts" / "embedded.css"),
    ]:
        if cand and Path(cand).is_file():
            return Path(cand)
    return None


def resolve_fonts_root(preferred: Path) -> Path:
    env_root = os.environ.get("CMGUI_FONTS_ROOT", "").strip()
    if env_root:
        p = Path(env_root)
        if (p / "NotoSansSC-Regular.otf").is_file():
            return p.resolve()
    if (preferred / "NotoSansSC-Regular.otf").is_file():
        return preferred.resolve()
    return preferred.resolve()


def cjk_font_css(fonts_root: Path) -> str:
    """Embed Noto/Inter from local /tmp fonts for Linux headless Chromium."""
    fonts = {
        "Noto Sans SC": [
            ("400", "normal", fonts_root / "NotoSansSC-Regular.otf"),
            ("500", "normal", fonts_root / "NotoSansSC-Medium.otf"),
            ("700", "normal", fonts_root / "NotoSansSC-Bold.otf"),
        ],
        "Inter": [
            ("400", "normal", fonts_root / "Inter-Regular.otf"),
            ("500", "normal", fonts_root / "Inter-Medium.otf"),
            ("600", "normal", fonts_root / "Inter-SemiBold.otf"),
            ("700", "normal", fonts_root / "Inter-Bold.otf"),
        ],
        "AnnotationCJK": [
            ("400", "normal", fonts_root / "NotoSansSC-Regular.otf"),
            ("500", "normal", fonts_root / "NotoSansSC-Medium.otf"),
            ("700", "normal", fonts_root / "NotoSansSC-Bold.otf"),
        ],
        "AnnotationLatin": [
            ("400", "normal", fonts_root / "Inter-Regular.otf"),
            ("500", "normal", fonts_root / "Inter-Medium.otf"),
            ("600", "normal", fonts_root / "Inter-SemiBold.otf"),
            ("700", "normal", fonts_root / "Inter-Bold.otf"),
        ],
    }
    lines = []
    seen_uri: set[str] = set()
    for family, entries in fonts.items():
        for weight, style, path in entries:
            if not path.is_file():
                continue
            uri = path.resolve().as_uri()
            key = f"{family}|{weight}|{style}|{uri}"
            if key in seen_uri:
                continue
            seen_uri.add(key)
            lines.append(
                f"@font-face{{font-family:'{family}';font-weight:{weight};font-style:{style};"
                f"src:url('{uri}') format('opentype');font-display:swap}}"
            )
    if not lines:
        return ""
    stack = "'Noto Sans SC','Inter','AnnotationCJK','AnnotationLatin','PingFang SC',sans-serif"
    lines.append(f"html,body,.screen,.text-node,.search-input-text,.suggestion-item{{font-family:{stack}!important}}")
    return "\n".join(lines)


def get_embedded_font_css() -> Optional[str]:
    global _EMBEDDED_FONT_CSS
    if _EMBEDDED_FONT_CSS is None:
        css_path = embedded_font_css_path()
        if css_path is not None:
            _EMBEDDED_FONT_CSS = css_path.read_text(encoding="utf-8")
    return _EMBEDDED_FONT_CSS


def inject_cjk_fonts_if_needed(html: str, fonts_root: Path) -> str:
    html = strip_file_font_faces(html)
    if should_use_embedded_css():
        css = get_embedded_font_css()
        if css:
            tag = f"<style id=\"cmgui-cjk-fonts\">\n{css}\n</style>\n"
            if re.search(r"</head>", html, re.IGNORECASE):
                return re.sub(r"</head>", tag + "</head>", html, count=1, flags=re.IGNORECASE)
            return tag + html
    local_root = resolve_fonts_root(fonts_root)
    css = cjk_font_css(local_root)
    if css:
        tag = f"<style id=\"cmgui-cjk-fonts\">\n{css}\n</style>\n"
        if re.search(r"</head>", html, re.IGNORECASE):
            return re.sub(r"</head>", tag + "</head>", html, count=1, flags=re.IGNORECASE)
        return tag + html
    css = get_embedded_font_css()
    if css:
        tag = f"<style id=\"cmgui-cjk-fonts\">\n{css}\n</style>\n"
        if re.search(r"</head>", html, re.IGNORECASE):
            return re.sub(r"</head>", tag + "</head>", html, count=1, flags=re.IGNORECASE)
        return tag + html
    return inject_fallback_font_stack(html)


async def wait_for_render_ready(page: Any) -> None:
    try:
        await asyncio.wait_for(page.evaluate("() => document.fonts.ready"), timeout=5.0)
    except Exception:
        pass
    await page.evaluate(
        """async () => {
          const timeout = (ms) => new Promise((res) => setTimeout(res, ms));
          const waits = Array.from(document.images).map((img) =>
            img.complete ? Promise.resolve() : Promise.race([
              new Promise((res) => {
                img.addEventListener('load', res, { once: true });
                img.addEventListener('error', res, { once: true });
              }),
              timeout(60000),
            ])
          );
          await Promise.all(waits);
          await new Promise((r) => setTimeout(r, 300));
        }"""
    )


def inject_base_href(html: str, base_dir: Path) -> str:
    base_tag = f'<base href="{base_dir.resolve().as_uri()}/">\n'
    if re.search(r"<head[\s>]", html, re.IGNORECASE):
        return re.sub(r"(<head[^>]*>)", r"\1\n" + base_tag, html, count=1, flags=re.IGNORECASE)
    return base_tag + html


async def render_one(
    page: Any,
    html_path: Path,
    out_png: Path,
    width: int,
    height: int,
    *,
    force: bool,
    fonts_root: Path,
) -> str:
    if out_png.is_file() and out_png.stat().st_size > 0 and not force:
        return "skipped"
    await page.set_viewport_size({"width": width, "height": height})
    html_text = html_path.read_text(encoding="utf-8", errors="replace")
    html_text = wrap_html_if_fragment(html_text)
    html_text = inject_cjk_fonts_if_needed(html_text, fonts_root)
    html_text = inject_viewport_css(html_text, width, height)
    html_dir = html_path.resolve().parent
    html_text = inject_base_href(html_text, html_dir)
    tmp_root = Path(os.environ.get("TMPDIR", "/tmp")) / "cmgui_render_html"
    tmp_root.mkdir(parents=True, exist_ok=True)
    tmp_html = tmp_root / f"{html_path.stem}.{uuid.uuid4().hex[:8]}.render.html"
    try:
        tmp_html.write_text(html_text, encoding="utf-8")
        await page.goto(tmp_html.resolve().as_uri(), wait_until="load", timeout=120_000)
        await wait_for_render_ready(page)
        await page.screenshot(path=str(out_png), full_page=False)
    finally:
        tmp_html.unlink(missing_ok=True)
    return "rendered"


async def process_pair(
    browser_pool: "BrowserPool",
    semaphore: asyncio.Semaphore,
    print_lock: asyncio.Lock,
    pair: Dict[str, Any],
    idx: int,
    end: int,
    args: argparse.Namespace,
) -> str:
    pair_id = str(pair["pair_id"])
    html_path = args.html_dir / f"{pair_id}.html"
    out_png = args.output_dir / f"{pair_id}.png"
    if not html_path.is_file():
        async with print_lock:
            print(f"[{idx}/{end}] missing html {pair_id}", flush=True)
        return "missing"
    if out_png.is_file() and out_png.stat().st_size > 0 and not args.force:
        async with print_lock:
            print(f"[{idx}/{end}] skipped {pair_id}", flush=True)
        return "skipped"
    async with semaphore:
        browser = None
        for attempt in range(2):
            page = None
            try:
                browser = await browser_pool.acquire()
                page = await browser.new_page(device_scale_factor=args.device_scale_factor)
                if args.width and args.height:
                    width, height = int(args.width), int(args.height)
                else:
                    width, height = image_size(pair, args.images_root)
                status = await render_one(
                    page, html_path, out_png, width, height,
                    force=args.force, fonts_root=args.fonts_root,
                )
                async with print_lock:
                    print(f"[{idx}/{end}] {status} {pair_id} -> {out_png.name} ({width}x{height})", flush=True)
                return status
            except Exception as exc:
                exc_name = type(exc).__name__
                if browser is not None:
                    await browser_pool.discard(browser)
                    browser = None
                if attempt == 0 and exc_name in {"TargetClosedError", "Error"}:
                    continue
                (args.output_dir / f"{pair_id}.render_error.txt").write_text(repr(exc), encoding="utf-8")
                async with print_lock:
                    print(f"[{idx}/{end}] FAIL {pair_id}: {exc!r}", flush=True)
                return "failed"
            finally:
                if page is not None:
                    try:
                        await page.close()
                    except Exception:
                        pass
                if browser is not None:
                    await browser_pool.release(browser)
                    browser = None
        return "failed"


class BrowserPool:
    def __init__(self, playwright: Any, launch_kwargs: Dict[str, Any], size: int) -> None:
        self._playwright = playwright
        self._launch_kwargs = launch_kwargs
        self._size = max(1, size)
        self._browsers: List[Any] = []
        self._available: asyncio.Queue[Any] = asyncio.Queue()
        self._lock = asyncio.Lock()

    async def _launch(self) -> Any:
        browser = await self._playwright.chromium.launch(**self._launch_kwargs)
        self._browsers.append(browser)
        return browser

    async def start(self) -> None:
        for _ in range(self._size):
            await self._available.put(await self._launch())

    async def acquire(self) -> Any:
        return await self._available.get()

    async def release(self, browser: Any) -> None:
        try:
            _ = browser.contexts
            await self._available.put(browser)
        except Exception:
            await self.discard(browser)

    async def discard(self, browser: Any) -> None:
        try:
            await browser.close()
        except Exception:
            pass
        if browser in self._browsers:
            self._browsers.remove(browser)
        async with self._lock:
            await self._available.put(await self._launch())

    async def close(self) -> None:
        for browser in list(self._browsers):
            try:
                await browser.close()
            except Exception:
                pass
        self._browsers.clear()


async def async_main(args: argparse.Namespace) -> int:
    deps_lib = Path.home() / "lib" / "playwright-deps" / "usr" / "lib" / "x86_64-linux-gnu"
    if deps_lib.is_dir():
        prev = os.environ.get("LD_LIBRARY_PATH", "")
        os.environ["LD_LIBRARY_PATH"] = str(deps_lib) + (os.pathsep + prev if prev else "")

    try:
        from playwright.async_api import async_playwright
    except ImportError as e:
        raise ImportError("Please install playwright and browsers: pip install playwright && playwright install chromium") from e

    pairs = read_pairs(args.metadata_json)
    args.html_dir = args.html_dir.resolve()
    args.output_dir = args.output_dir.resolve()
    args.images_root = args.images_root.resolve()
    end = len(pairs) if args.max_samples is None else min(len(pairs), args.start + args.max_samples)
    subset = pairs[args.start:end]
    args.output_dir.mkdir(parents=True, exist_ok=True)
    if should_use_embedded_css():
        get_embedded_font_css()

    rendered = skipped = missing = failed = 0
    workers = max(1, int(args.workers))
    async with async_playwright() as p:
        launch_kwargs = {"headless": True, "args": CHROMIUM_ARGS}
        if args.browser_channel:
            launch_kwargs["channel"] = args.browser_channel
        pool = BrowserPool(p, launch_kwargs, workers)
        await pool.start()
        semaphore = asyncio.Semaphore(workers)
        print_lock = asyncio.Lock()
        print(
            f"[render] workers={workers} browsers={workers} "
            f"fonts_root={resolve_fonts_root(args.fonts_root)} "
            f"samples={len(subset)}",
            flush=True,
        )
        try:
            results = await asyncio.gather(
                *[
                    process_pair(pool, semaphore, print_lock, pair, i, end, args)
                    for i, pair in enumerate(subset, start=args.start + 1)
                ]
            )
        finally:
            await pool.close()

    for status in results:
        if status == "rendered":
            rendered += 1
        elif status == "skipped":
            skipped += 1
        elif status == "missing":
            missing += 1
        elif status == "failed":
            failed += 1

    summary = {
        "total": len(subset),
        "rendered": rendered,
        "skipped_existing": skipped,
        "missing_html": missing,
        "failed": failed,
        "workers": workers,
        "html_dir": str(args.html_dir),
        "output_dir": str(args.output_dir),
    }
    (args.output_dir / "render_summary.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")
    print(json.dumps(summary, ensure_ascii=False, indent=2))
    return 0 if rendered or skipped else 1


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--metadata-json", type=Path, default=DEFAULT_META)
    ap.add_argument("--images-root", type=Path, default=DEFAULT_IMAGES)
    ap.add_argument("--html-dir", type=Path, default=DEFAULT_HTML_DIR)
    ap.add_argument("--output-dir", type=Path, default=DEFAULT_OUT_DIR)
    ap.add_argument("--start", type=int, default=0)
    ap.add_argument("--max-samples", type=int, default=None)
    ap.add_argument("--width", type=int, default=None)
    ap.add_argument("--height", type=int, default=None)
    ap.add_argument("--device-scale-factor", type=float, default=1.0)
    ap.add_argument("--browser-channel", type=str, default=None, help="Optional Playwright browser channel, e.g. chrome or msedge")
    ap.add_argument("--fonts-root", type=Path, default=DEFAULT_FONTS_ROOT)
    ap.add_argument("--workers", type=int, default=1, help="Parallel Playwright pages (default 1)")
    ap.add_argument("--force", action="store_true")
    args = ap.parse_args()
    return asyncio.run(async_main(args))


if __name__ == "__main__":
    raise SystemExit(main())