#!/usr/bin/env python3
"""Render local HTML (especially Mermaid diagrams) to a high-quality PNG."""

from __future__ import annotations

import argparse
import asyncio
import json
import sys
from pathlib import Path

from playwright.async_api import async_playwright

INJECT_STYLE = """
html, body { margin:0; overflow:visible !important; background:#fff !important; }
body { padding:24px 28px 80px !important; width:max-content !important; min-width:max-content !important; }
header { margin-bottom:18px !important; text-align:center !important; }
h1 { font-size:28px !important; font-weight:800 !important; color:#0f172a !important; }
.legend {
    display:flex !important; justify-content:center !important; gap:32px !important;
    font-size:16px !important; font-weight:600 !important; color:#334155 !important;
    margin:16px auto 22px !important; padding:12px 24px !important;
    background:#f8fafc !important; border:1px solid #e2e8f0 !important; border-radius:10px !important;
    width:max-content !important; max-width:none !important;
}
.leg i { width:20px !important; height:20px !important; border-width:2px !important; border-radius:4px !important; }
.wrap { overflow:visible !important; padding:24px 20px 40px !important; box-shadow:none !important; width:max-content !important; }
.mermaid { min-width:auto !important; overflow:visible !important; width:max-content !important; }
.mermaid svg { max-width:none !important; overflow:visible !important; display:block !important; }
"""

FIX_SVG_JS = """
({ padX, padTop, padBottom }) => {
    const svg = document.querySelector('.mermaid svg');
    if (!svg) return { fixed: false };

    function fullBBox(root) {
        let minX = Infinity, minY = Infinity, maxX = -Infinity, maxY = -Infinity;
        root.querySelectorAll('g, path, rect, text, line, polygon, ellipse, foreignObject').forEach(el => {
            if (!el.getBBox) return;
            try {
                const b = el.getBBox();
                if (!b.width && !b.height) return;
                minX = Math.min(minX, b.x);
                minY = Math.min(minY, b.y);
                maxX = Math.max(maxX, b.x + b.width);
                maxY = Math.max(maxY, b.y + b.height);
            } catch (e) {}
        });
        return { x: minX, y: minY, width: maxX - minX, height: maxY - minY };
    }

    const full = fullBBox(svg);
    const vbX = full.x - padX;
    const vbY = full.y - padTop;
    const vbW = full.width + padX * 2;
    const vbH = full.height + padTop + padBottom;
    svg.setAttribute('viewBox', `${vbX} ${vbY} ${vbW} ${vbH}`);
    svg.removeAttribute('height');
    svg.removeAttribute('width');
    svg.style.width = vbW + 'px';
    svg.style.height = vbH + 'px';
    return { fixed: true, vbW, vbH };
}
"""

MEASURE_JS = """
(selectors) => {
    const parts = selectors
        .map(sel => document.querySelector(sel))
        .filter(Boolean)
        .map(el => el.getBoundingClientRect());
    if (!parts.length) {
        const body = document.body.getBoundingClientRect();
        return { totalW: Math.ceil(body.width + 48), totalH: Math.ceil(body.height + 48) };
    }
    const left = Math.min(...parts.map(r => r.left));
    const right = Math.max(...parts.map(r => r.right));
    const top = Math.min(...parts.map(r => r.top));
    const bottom = Math.max(...parts.map(r => r.bottom));
    return {
        totalW: Math.ceil(right - left + 48),
        totalH: Math.ceil(bottom - top + 48)
    };
}
"""

DEFAULT_SELECTORS = ["header", ".legend", ".wrap"]


async def render_html_to_png(
    html_path: Path,
    out_path: Path,
    *,
    dpr: float = 2.0,
    selectors: list[str] | None = None,
    pad_x: int = 32,
    pad_top: int = 48,
    pad_bottom: int = 72,
    initial_viewport: tuple[int, int] = (10000, 2600),
    wait_ms: int = 3000,
    verbose: bool = False,
) -> dict:
    selectors = selectors or DEFAULT_SELECTORS
    html_path = html_path.resolve()
    out_path = out_path.resolve()
    out_path.parent.mkdir(parents=True, exist_ok=True)

    async with async_playwright() as p:
        browser = await p.chromium.launch()
        page = await browser.new_page(device_scale_factor=dpr)
        await page.set_viewport_size({"width": initial_viewport[0], "height": initial_viewport[1]})
        await page.goto(html_path.as_uri(), wait_until="networkidle")

        has_mermaid = await page.query_selector(".mermaid svg")
        if has_mermaid:
            await page.wait_for_selector(".mermaid svg", timeout=45000)
        await page.wait_for_timeout(wait_ms)

        await page.add_style_tag(content=INJECT_STYLE)

        svg_info = None
        if has_mermaid:
            svg_info = await page.evaluate(
                FIX_SVG_JS, {"padX": pad_x, "padTop": pad_top, "padBottom": pad_bottom}
            )

        dims = await page.evaluate(MEASURE_JS, selectors)
        w, h = dims["totalW"], dims["totalH"]

        if verbose:
            print(json.dumps({"svg": svg_info, "dims": dims, "dpr": dpr}, ensure_ascii=False, indent=2))

        await page.set_viewport_size({"width": w, "height": h})
        await page.wait_for_timeout(600)
        await page.screenshot(
            path=str(out_path),
            type="png",
            clip={"x": 0, "y": 0, "width": w, "height": h},
        )
        await browser.close()

    return {"input": str(html_path), "output": str(out_path), "logical_size": [w, h], "dpr": dpr}


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Convert local HTML to high-quality PNG via Playwright")
    parser.add_argument("html", type=Path, help="Input HTML file path")
    parser.add_argument("-o", "--output", type=Path, help="Output PNG path (default: <html_stem>@2x.png)")
    parser.add_argument("--dpr", type=float, default=2.0, help="Device pixel ratio (default: 2)")
    parser.add_argument(
        "--selector",
        action="append",
        dest="selectors",
        help="CSS selector for capture bounds (repeatable; default: header, .legend, .wrap)",
    )
    parser.add_argument("--pad-bottom", type=int, default=72, help="Extra SVG bottom padding")
    parser.add_argument("-v", "--verbose", action="store_true", help="Print measurement info")
    return parser.parse_args(argv)


def main(argv: list[str] | None = None) -> int:
    args = parse_args(argv)
    html_path: Path = args.html
    if not html_path.is_file():
        print(f"Error: HTML not found: {html_path}", file=sys.stderr)
        return 1

    out_path = args.output or html_path.with_name(f"{html_path.stem}@2x.png")

    try:
        result = asyncio.run(
            render_html_to_png(
                html_path,
                out_path,
                dpr=args.dpr,
                selectors=args.selectors,
                pad_bottom=args.pad_bottom,
                verbose=args.verbose,
            )
        )
    except Exception as exc:
        print(f"Error: {exc}", file=sys.stderr)
        return 1

    w, h = result["logical_size"]
    px_w, px_h = int(w * args.dpr), int(h * args.dpr)
    print(f"Saved: {result['output']}")
    print(f"Logical: {w}x{h}, pixels: ~{px_w}x{px_h} (dpr={args.dpr})")
    return 0


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