"""用 Playwright 将 Mermaid 源码渲染为 PNG（复用包内 res/mermaid.min.js）。"""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any, Callable

from aicd.kit.validate_viz import validate_mermaid_body
from aicd.tools import viz_assets

MAX_MERMAID_RENDER_CHARS = 120_000


def _default_mermaid_init(
    theme: str, theme_variables: dict[str, Any] | None
) -> dict[str, Any]:
    init: dict[str, Any] = {
        "startOnLoad": False,
        "securityLevel": "loose",
        "theme": (theme or "neutral").strip() or "neutral",
    }
    if theme_variables and isinstance(theme_variables, dict):
        init["themeVariables"] = theme_variables
    return init


def build_mermaid_render_html(
    mermaid_body: str, mermaid_init: dict[str, Any]
) -> str:
    diagram_js = json.dumps(mermaid_body, ensure_ascii=False)
    init_js = json.dumps(mermaid_init, ensure_ascii=False)
    return f"""<!DOCTYPE html>
<html lang="zh-CN">
<head>
  <meta charset="utf-8"/>
  <style>html,body{{margin:0;background:#fff;}}</style>
</head>
<body>
<div id="wrap" style="display:inline-block;padding:16px 20px;box-sizing:content-box;"></div>
  <script src="res/{viz_assets.MERMAID_JS}"></script>
  <script>
  (async function () {{
    const diagram = {diagram_js};
    const initOpts = {init_js};
    mermaid.initialize(initOpts);
    const id = "aicd_mmd_" + Math.random().toString(36).slice(2, 12);
    try {{
      const out = await mermaid.render(id, diagram);
      document.getElementById("wrap").innerHTML = out.svg;
    }} catch (e) {{
      document.getElementById("wrap").textContent = "MERMAID_RENDER_ERROR: " + (e && e.message ? e.message : String(e));
    }}
  }})();
  </script>
</body>
</html>
"""


def render_mermaid_to_png(
    mermaid_body: str,
    output_png: Path,
    *,
    policy_check: Callable[[Path], None],
    theme: str = "neutral",
    theme_variables: dict[str, Any] | None = None,
    scale: float = 3.0,
    timeout_ms: int = 90_000,
) -> dict[str, Any]:
    """
    在工作目录写入临时 HTML + 复制 res/，用 Chromium 打开并截取 ``#wrap``。

    ``theme_variables`` 对应 Mermaid ``theme: base`` 时常用；也可在 init 里配合 ``theme: default`` 等。
    若参考文档为深色背景，可设 ``theme`` 为 ``dark`` 或传入 ``themeVariables``（如 primaryColor、lineColor）。
    """
    out = output_png.expanduser().resolve()
    policy_check(out)
    if out.suffix.lower() != ".png":
        return {"ok": False, "error": "output_png_path must end with .png"}

    body = (mermaid_body or "").strip()
    if len(body) > MAX_MERMAID_RENDER_CHARS:
        return {
            "ok": False,
            "error": f"mermaid_body exceeds {MAX_MERMAID_RENDER_CHARS} chars",
        }
    v = validate_mermaid_body(body)
    if not v.get("ok"):
        return {
            "ok": False,
            "error": "invalid mermaid_body",
            "validation": v,
        }

    work_dir = out.parent
    policy_check(work_dir)
    work_dir.mkdir(parents=True, exist_ok=True)

    copied = viz_assets.copy_viz_resources(work_dir)
    if not copied.get("ok"):
        return {"ok": False, "error": copied.get("error", "copy_viz_resources failed")}

    init = _default_mermaid_init(theme, theme_variables)
    html = build_mermaid_render_html(body, init)
    html_path = work_dir / "_aicd_mermaid_render.html"
    html_path.write_text(html, encoding="utf-8")
    policy_check(html_path)
    uri = html_path.resolve().as_uri()

    try:
        from playwright.sync_api import sync_playwright
    except ImportError as e:
        return {"ok": False, "error": f"playwright not importable: {e}"}

    try:
        with sync_playwright() as p:
            browser = p.chromium.launch(headless=True)
            try:
                context = browser.new_context(
                    device_scale_factor=max(1.0, min(6.0, float(scale)))
                )
                page = context.new_page()
                page.set_default_timeout(timeout_ms)
                page.goto(uri, wait_until="load")
                page.wait_for_selector("#wrap svg", timeout=timeout_ms)
                page.wait_for_timeout(400)
                err = page.locator("#wrap").inner_text()
                if err and "MERMAID_RENDER_ERROR:" in err:
                    return {"ok": False, "error": err.strip()}
                page.locator("#wrap").screenshot(path=str(out), type="png")
            finally:
                browser.close()
    except Exception as e:
        msg = str(e).lower()
        hint = "执行: python -m playwright install chromium（需已 pip install playwright）。"
        if "executable doesn't exist" in msg or "browsertype.launch" in msg:
            hint = "未找到 Playwright 自带的 Chromium，请先运行: python -m playwright install chromium"
        return {"ok": False, "error": str(e), "hint": hint}

    try:
        html_path.unlink(missing_ok=True)
    except OSError:
        pass

    if not out.is_file() or out.stat().st_size == 0:
        return {"ok": False, "error": "PNG was not written or is empty"}

    return {
        "ok": True,
        "png_path": str(out),
        "work_dir": str(work_dir),
        "mermaid_theme": init.get("theme"),
        "validation": {"char_count": len(body)},
    }
