r"""把 docs/intro.html 渲染成动画预览图。

做法：用 headless Chrome 打开 `intro.html?t=<秒>`，逐帧截图，
再用 Pillow 拼成动画 WebP。

之所以可行，是因为动画的时间轴是纯函数 render(t)——
任意一帧都能独立渲染，不需要真实播放，也不会有录制抖动。

用法：
    python tools/render_preview.py                    # 默认 10fps / 800x450
    python tools/render_preview.py --fps 15 --width 1200
    python tools/render_preview.py --stills           # 只导出关键帧静图

依赖：Pillow（pip install pillow）

Windows 注意：
    在某些环境下，从 Python 里 spawn Chrome 会以 0x80000003 直接退出
    （从 PowerShell 直接调同样的命令却正常）。遇到这种情况，用
    `--assemble-only` 跳过渲染，改在 PowerShell 里逐帧截图：

        $chrome = "C:\Program Files\Google\Chrome\Application\chrome.exe"
        $url = ([uri]"$PWD\\docs\\intro.html").AbsoluteUri
        for ($i=0; $i -lt 220; $i++) {
          $f = "._frames\\f{0:d4}.png" -f $i
          & $chrome --headless=new --disable-gpu --window-size=800,450 `
            --virtual-time-budget=1200 "--screenshot=$PWD\\$f" `
            "$url`?t=$([math]::Round($i/10.0,3))"
        }

    然后再跑 `--assemble-only`。
"""

from __future__ import annotations

import argparse
import os
import subprocess
import sys
import tempfile
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

try:
    from PIL import Image
except ImportError:
    sys.exit("需要 Pillow：pip install pillow")

REPO = Path(__file__).resolve().parents[1]
INTRO = REPO / "docs" / "intro.html"

CHROME_CANDIDATES = [
    os.environ.get("CHROME_PATH", ""),
    r"C:\Program Files\Google\Chrome\Application\chrome.exe",
    r"C:\Program Files (x86)\Google\Chrome\Application\chrome.exe",
    r"C:\Program Files (x86)\Microsoft\Edge\Application\msedge.exe",
    r"C:\Program Files\Microsoft\Edge\Application\msedge.exe",
    "/usr/bin/google-chrome",
    "/usr/bin/chromium",
    "/Applications/Google Chrome.app/Contents/MacOS/Google Chrome",
]

#: 关键帧：每一幕的最佳瞬间
STILLS = {
    "01-membrane.png": 3.2,
    "02-heartbeat.png": 11.4,
    "03-memory.png": 15.8,
    "04-title.png": 19.8,
}


def find_chrome() -> str:
    for path in CHROME_CANDIDATES:
        if path and Path(path).exists():
            return path
    sys.exit("找不到 Chrome/Edge；设 CHROME_PATH 环境变量指向可执行文件")


def file_url(t: float) -> str:
    return INTRO.as_uri() + f"?t={t:.3f}"


def shoot(chrome: str, t: float, out: Path, w: int, h: int, profile: Path) -> Path:
    """渲染单帧。"""
    cmd = [
        chrome,
        "--headless=new",
        "--disable-gpu",
        "--hide-scrollbars",
        "--no-first-run",
        "--no-default-browser-check",
        f"--user-data-dir={profile}",
        f"--window-size={w},{h}",
        "--virtual-time-budget=1500",
        f"--screenshot={out}",
        file_url(t),
    ]
    subprocess.run(
        cmd,
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
        timeout=60,
        check=False,
    )
    if not out.exists():
        raise RuntimeError(f"截图失败: t={t}")
    return out


def render_frames(
    chrome: str, duration: float, fps: float, w: int, h: int, outdir: Path, workers: int
) -> list[Path]:
    n = int(round(duration * fps))
    times = [i / fps for i in range(n)]
    outdir.mkdir(parents=True, exist_ok=True)

    def work(i: int) -> Path:
        profile = Path(tempfile.mkdtemp(prefix=f"cg-{i}-"))
        try:
            return shoot(chrome, times[i], outdir / f"f{i:04d}.png", w, h, profile)
        finally:
            import shutil

            shutil.rmtree(profile, ignore_errors=True)

    print(f"渲染 {n} 帧 @ {fps}fps，{workers} 并发 …")
    paths: list[Path] = [None] * n  # type: ignore[list-item]
    done = 0
    with ThreadPoolExecutor(max_workers=workers) as pool:
        for i, p in zip(range(n), pool.map(work, range(n)), strict=True):
            paths[i] = p
            done += 1
            if done % 20 == 0 or done == n:
                print(f"  {done}/{n}")
    return paths


def assemble(paths: list[Path], out: Path, fps: float, quality: int) -> None:
    frames = [Image.open(p).convert("RGB") for p in paths]
    frames[0].save(
        out,
        save_all=True,
        append_images=frames[1:],
        duration=int(round(1000 / fps)),
        loop=0,
        quality=quality,
        method=6,
    )
    size_mb = out.stat().st_size / 1024 / 1024
    print(f"写出 {out.relative_to(REPO)}  ({size_mb:.2f} MB, {len(frames)} 帧)")


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--fps", type=float, default=10)
    ap.add_argument("--width", type=int, default=800)
    ap.add_argument("--height", type=int, default=450)
    ap.add_argument("--duration", type=float, default=22.0)
    ap.add_argument("--workers", type=int, default=6)
    ap.add_argument("--quality", type=int, default=82)
    ap.add_argument("--out", default="docs/preview.webp")
    ap.add_argument("--frames-dir", default=".frames")
    ap.add_argument("--stills", action="store_true", help="只导出关键帧静图")
    ap.add_argument(
        "--assemble-only",
        action="store_true",
        help="跳过渲染，直接把 --frames-dir 里已有的 PNG 拼成动图",
    )
    ap.add_argument("--keep-frames", action="store_true")
    args = ap.parse_args()

    if args.assemble_only:
        work = REPO / args.frames_dir
        paths = sorted(work.glob("f*.png"))
        if not paths:
            sys.exit(f"{work} 里没有帧")
        assemble(paths, REPO / args.out, args.fps, args.quality)
        return

    chrome = find_chrome()
    work = REPO / args.frames_dir
    work.mkdir(parents=True, exist_ok=True)

    if args.stills:
        for name, t in STILLS.items():
            profile = Path(tempfile.mkdtemp(prefix="cg-s-"))
            try:
                out = REPO / "docs" / name
                shoot(chrome, t, out, args.width, args.height, profile)
                print(f"写出 {out.relative_to(REPO)}")
            finally:
                import shutil

                shutil.rmtree(profile, ignore_errors=True)
        return

    paths = render_frames(
        chrome, args.duration, args.fps, args.width, args.height, work, args.workers
    )
    assemble(paths, REPO / args.out, args.fps, args.quality)

    if not args.keep_frames:
        import shutil

        shutil.rmtree(work, ignore_errors=True)


if __name__ == "__main__":
    main()