ac-docs/_pdf_build/mermaid_render.py

"""Render Mermaid diagram sources to PNG for PDF embedding (mermaid-cli + Chromium)."""

from __future__ import annotations

import hashlib
import subprocess
import tempfile
from pathlib import Path

from reportlab.lib.units import cm
from reportlab.lib.utils import ImageReader
from reportlab.platypus import Image, Paragraph, Spacer

BUILD_DIR = Path(__file__).resolve().parent
CONFIG = BUILD_DIR / "mermaid-config.json"
CACHE_DIR = BUILD_DIR / "mermaid_cache"
MERMAID_CLI = ("npx", "--yes", "@mermaid-js/mermaid-cli@11.4.0")

MAX_DIAGRAM_WIDTH_CM = 15.8
MAX_DIAGRAM_HEIGHT_CM = 17.0
CAPTION_STYLE_NAME = "DiagramCaption"


def _content_hash(content: str) -> str:
    return hashlib.sha256(content.strip().encode("utf-8")).hexdigest()[:16]


def estimate_mermaid_pixels(content: str) -> tuple[int, int, int]:
    """Return (width_px, height_px, scale) heuristics for flowchart layout."""
    lines = [ln for ln in content.splitlines() if ln.strip() and not ln.strip().startswith("%%")]
    n_edges = sum(1 for ln in lines if "-->" in ln or "---" in ln or "-.->" in ln)
    n_nodes = sum(
        1
        for ln in lines
        if "[" in ln or "(" in ln or "{" in ln or (ln.strip() and "-->" not in ln and "subgraph" not in ln.lower())
    )
    n_sub = sum(1 for ln in lines if "subgraph" in ln.lower())
    lr = any("flowchart LR" in ln or "graph LR" in ln for ln in lines)

    scale = 2
    if lr:
        width = 2200 + min(n_nodes, 12) * 120
        height = 650 + n_sub * 180 + min(n_edges, 10) * 30
    else:
        width = 1700 + min(n_sub, 4) * 200
        height = 650 + n_edges * 70 + n_sub * 220 + min(n_nodes, 14) * 25

    width = int(min(max(width, 1200), 3600))
    height = int(min(max(height, 550), 3400))
    return width, height, scale


def render_mermaid_to_png(content: str, dest: Path, *, width: int, height: int, scale: int) -> None:
    if not CONFIG.exists():
        raise FileNotFoundError(f"mermaid config missing: {CONFIG}")
    dest.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.NamedTemporaryFile("w", suffix=".mmd", delete=False, encoding="utf-8") as tmp:
        tmp.write(content.strip() + "\n")
        mmd_path = Path(tmp.name)
    try:
        subprocess.run(
            [
                *MERMAID_CLI,
                "-i",
                str(mmd_path),
                "-o",
                str(dest),
                "-c",
                str(CONFIG),
                "-b",
                "white",
                "-w",
                str(width),
                "-H",
                str(height),
                "-s",
                str(scale),
            ],
            cwd=str(BUILD_DIR.parents[1]),
            check=True,
            timeout=180,
            capture_output=True,
            text=True,
        )
    except subprocess.CalledProcessError as exc:
        msg = (exc.stderr or exc.stdout or str(exc)).strip()
        raise RuntimeError(f"mermaid-cli failed: {msg}") from exc
    finally:
        mmd_path.unlink(missing_ok=True)
    if not dest.is_file() or dest.stat().st_size == 0:
        raise RuntimeError(f"mermaid-cli produced no PNG: {dest}")


def cached_mermaid_png(content: str) -> Path:
    """Render or reuse cached PNG for Mermaid source."""
    CACHE_DIR.mkdir(parents=True, exist_ok=True)
    key = _content_hash(content)
    png = CACHE_DIR / f"{key}.png"
    if png.is_file() and png.stat().st_size > 0:
        return png
    w, h, scale = estimate_mermaid_pixels(content)
    render_mermaid_to_png(content, png, width=w, height=h, scale=scale)
    return png


def diagram_flowables(
    png_path: Path,
    styles,
    *,
    caption: str = "",
    max_width_cm: float = MAX_DIAGRAM_WIDTH_CM,
    max_height_cm: float = MAX_DIAGRAM_HEIGHT_CM,
) -> list:
    """ReportLab flowables: image (+ optional caption) sized to fit A4 body."""
    reader = ImageReader(str(png_path))
    w_px, h_px = reader.getSize()
    max_w = max_width_cm * cm
    max_h = max_height_cm * cm
    scale = min(max_w / w_px, max_h / h_px)
    img = Image(str(png_path), width=w_px * scale, height=h_px * scale)
    img.hAlign = "CENTER"
    out: list = [Spacer(1, 0.12 * cm), img, Spacer(1, 0.06 * cm)]
    if caption:
        cap = styles[CAPTION_STYLE_NAME] if CAPTION_STYLE_NAME in styles else styles["Normal"]
        out.insert(0, Spacer(1, 0.06 * cm))
        out.append(Paragraph(f"<i>{caption}</i>", cap))
    out.append(Spacer(1, 0.1 * cm))
    return out


def ensure_diagram_caption_style(styles) -> None:
    from reportlab.lib.styles import ParagraphStyle
    from reportlab.lib import colors

    if CAPTION_STYLE_NAME not in styles:
        styles.add(
            ParagraphStyle(
                CAPTION_STYLE_NAME,
                parent=styles["Normal"],
                fontSize=7,
                textColor=colors.grey,
                alignment=1,
                spaceAfter=4,
            )
        )