render.py

"""
mermaid-renderer — Convert Mermaid diagrams to images using @mermaid-js/mermaid-cli
Author: algorembrant
License: MIT

USAGE COMMANDS
==============

  Single file:
    python render.py diagram.mmd
    python render.py diagram.mmd -o output.png
    python render.py diagram.mmd -o output.svg
    python render.py diagram.mmd -o output.pdf

  Inline diagram:
    python render.py --text "graph TD; A --> B --> C"
    python render.py --text "sequenceDiagram; Alice->>Bob: Hello"

  Themes (15 built-in from beautiful-mermaid):
    python render.py diagram.mmd --theme tokyo-night
    python render.py diagram.mmd --theme catppuccin-mocha
    python render.py diagram.mmd --theme github-light
    python render.py diagram.mmd --theme dracula
    python render.py diagram.mmd --theme nord
    python render.py diagram.mmd --theme one-dark

  Custom colors:
    python render.py diagram.mmd --bg "#1a1b26" --fg "#a9b1d6"
    python render.py diagram.mmd --bg "#ffffff" --fg "#333333" --accent "#0969da"

  Background:
    python render.py diagram.mmd --background white
    python render.py diagram.mmd --background transparent
    python render.py diagram.mmd --background "#1e1e2e"

  Scale / quality:
    python render.py diagram.mmd --scale 2
    python render.py diagram.mmd --scale 3
    python render.py diagram.mmd --width 1920

  Batch (directory or glob):
    python render.py --batch ./diagrams/
    python render.py --batch ./diagrams/ -o ./output/ --format png
    python render.py --batch "./diagrams/*.mmd" --theme tokyo-night --scale 2

  Batch parallel workers:
    python render.py --batch ./diagrams/ --workers 4

  Watch mode (re-render on change):
    python render.py diagram.mmd --watch
    python render.py --batch ./diagrams/ --watch

  List built-in themes:
    python render.py --list-themes

  Check environment / dependencies:
    python render.py --check

  Verbose logging:
    python render.py diagram.mmd -v
    python render.py --batch ./diagrams/ -v --workers 2

  Help:
    python render.py --help
"""

from __future__ import annotations

import argparse
import glob
import json
import logging
import os
import shutil
import subprocess
import sys
import tempfile
import time
from concurrent.futures import ProcessPoolExecutor, as_completed
from pathlib import Path
from typing import Optional

# ---------------------------------------------------------------------------
# Logging
# ---------------------------------------------------------------------------
logging.basicConfig(
    level=logging.WARNING,
    format="[%(levelname)s] %(message)s",
)
log = logging.getLogger("mermaid-renderer")


# ---------------------------------------------------------------------------
# Theme definitions (mirrors beautiful-mermaid THEMES)
# ---------------------------------------------------------------------------
THEMES: dict[str, dict] = {
    "default": {
        "bg": "#FFFFFF",
        "fg": "#27272A",
        "mermaid_theme": "default",
    },
    "zinc-light": {
        "bg": "#FFFFFF",
        "fg": "#18181B",
        "mermaid_theme": "default",
    },
    "zinc-dark": {
        "bg": "#18181B",
        "fg": "#E4E4E7",
        "mermaid_theme": "dark",
    },
    "tokyo-night": {
        "bg": "#1a1b26",
        "fg": "#a9b1d6",
        "accent": "#7aa2f7",
        "mermaid_theme": "dark",
    },
    "tokyo-night-storm": {
        "bg": "#24283b",
        "fg": "#c0caf5",
        "accent": "#7aa2f7",
        "mermaid_theme": "dark",
    },
    "tokyo-night-light": {
        "bg": "#d5d6db",
        "fg": "#343b58",
        "accent": "#34548a",
        "mermaid_theme": "default",
    },
    "catppuccin-mocha": {
        "bg": "#1e1e2e",
        "fg": "#cdd6f4",
        "accent": "#cba6f7",
        "mermaid_theme": "dark",
    },
    "catppuccin-latte": {
        "bg": "#eff1f5",
        "fg": "#4c4f69",
        "accent": "#8839ef",
        "mermaid_theme": "default",
    },
    "nord": {
        "bg": "#2e3440",
        "fg": "#d8dee9",
        "accent": "#88c0d0",
        "mermaid_theme": "dark",
    },
    "nord-light": {
        "bg": "#eceff4",
        "fg": "#2e3440",
        "accent": "#5e81ac",
        "mermaid_theme": "default",
    },
    "dracula": {
        "bg": "#282a36",
        "fg": "#f8f8f2",
        "accent": "#bd93f9",
        "mermaid_theme": "dark",
    },
    "github-light": {
        "bg": "#ffffff",
        "fg": "#1F2328",
        "accent": "#0969da",
        "mermaid_theme": "default",
    },
    "github-dark": {
        "bg": "#0d1117",
        "fg": "#e6edf3",
        "accent": "#4493f8",
        "mermaid_theme": "dark",
    },
    "solarized-light": {
        "bg": "#fdf6e3",
        "fg": "#657b83",
        "accent": "#268bd2",
        "mermaid_theme": "default",
    },
    "solarized-dark": {
        "bg": "#002b36",
        "fg": "#839496",
        "accent": "#268bd2",
        "mermaid_theme": "dark",
    },
    "one-dark": {
        "bg": "#282c34",
        "fg": "#abb2bf",
        "accent": "#c678dd",
        "mermaid_theme": "dark",
    },
}


# ---------------------------------------------------------------------------
# Environment detection
# ---------------------------------------------------------------------------

def find_mmdc() -> Optional[str]:
    """Return path to mmdc binary or None."""
    # Prefer the npm-global install used by this environment
    candidates = [
        shutil.which("mmdc"),
        shutil.which("mmdc.cmd"),
        os.path.expanduser("~/.npm-global/bin/mmdc"),
        os.path.join(os.environ.get("APPDATA", ""), "npm", "mmdc.cmd"),
        "/usr/local/bin/mmdc",
        "/usr/bin/mmdc",
    ]
    for c in candidates:
        if c and os.path.isfile(c):
            return c
    return None


def find_chrome() -> Optional[str]:
    """Return path to a Chromium/Chrome binary or None."""
    candidates = [
        # Windows paths
        r"C:\Program Files\Google\Chrome\Application\chrome.exe",
        r"C:\Program Files (x86)\Google\Chrome\Application\chrome.exe",
        os.path.join(os.environ.get("LOCALAPPDATA", ""), r"Google\Chrome\Application\chrome.exe"),
        # Puppeteer-managed chrome (most reliable in this env)
        os.path.expanduser(
            "~/.cache/puppeteer/chrome/linux-131.0.6778.204/chrome-linux64/chrome"
        ),
        "/opt/google/chrome/chrome",
        "/opt/pw-browsers/chromium-1194/chrome-linux/chrome",
        shutil.which("google-chrome"),
        shutil.which("chromium"),
        shutil.which("chromium-browser"),
        shutil.which("chrome"),
    ]
    for c in candidates:
        if c and os.path.isfile(c):
            return c
    return None


def build_puppeteer_config(chrome_path: str) -> str:
    """Write a puppeteer config JSON to a temp file, return its path."""
    cfg = {
        "executablePath": chrome_path,
        "args": [
            "--no-sandbox",
            "--disable-setuid-sandbox",
            "--disable-dev-shm-usage",
            "--disable-gpu",
        ],
    }
    fd, path = tempfile.mkstemp(suffix=".json", prefix="puppeteer_cfg_")
    with os.fdopen(fd, "w") as f:
        json.dump(cfg, f)
    return path


def build_mermaid_config(theme_cfg: dict, bg: Optional[str]) -> str:
    """Write a mermaid config JSON, return its path."""
    mermaid_theme = theme_cfg.get("mermaid_theme", "default")
    cfg: dict = {
        "theme": mermaid_theme,
    }
    if mermaid_theme == "dark":
        cfg["themeVariables"] = {
            "darkMode": True,
            "background": theme_cfg.get("bg", "#1e1e2e"),
            "primaryColor": theme_cfg.get("accent", theme_cfg.get("fg", "#88c0d0")),
            "primaryTextColor": theme_cfg.get("fg", "#d8dee9"),
            "lineColor": theme_cfg.get("accent", "#88c0d0"),
        }
    else:
        cfg["themeVariables"] = {
            "darkMode": False,
            "background": theme_cfg.get("bg", "#ffffff"),
            "primaryColor": theme_cfg.get("accent", "#ececff"),
            "primaryTextColor": theme_cfg.get("fg", "#333333"),
            "lineColor": theme_cfg.get("accent", "#333333"),
        }

    if bg:
        cfg["themeVariables"]["background"] = bg

    fd, path = tempfile.mkstemp(suffix=".json", prefix="mermaid_cfg_")
    with os.fdopen(fd, "w") as f:
        json.dump(cfg, f)
    return path


# ---------------------------------------------------------------------------
# Core render function
# ---------------------------------------------------------------------------

def render_diagram(
    source: str,
    output_path: str,
    *,
    mmdc: str,
    puppeteer_cfg: str,
    theme_name: str = "default",
    bg: Optional[str] = None,
    scale: float = 2.0,
    width: Optional[int] = None,
    fmt: str = "png",
    verbose: bool = False,
) -> dict:
    """
    Render a Mermaid diagram string to an image file.

    Parameters
    ----------
    source       : Mermaid diagram text
    output_path  : Destination file path
    mmdc         : Path to mmdc binary
    puppeteer_cfg: Path to puppeteer config JSON
    theme_name   : One of THEMES keys or 'default'
    bg           : Override background color (hex or 'transparent')
    scale        : Device pixel ratio (default 2 = 2x)
    width        : Optional canvas width in pixels
    fmt          : Output format: png | svg | pdf
    verbose      : Log mmdc stdout/stderr

    Returns
    -------
    dict with keys: success, output, error, duration_ms
    """
    theme_cfg = THEMES.get(theme_name, THEMES["default"])
    t0 = time.perf_counter()

    # Write source to temp .mmd file
    fd, src_path = tempfile.mkstemp(suffix=".mmd", prefix="mermaid_src_")
    mermaid_cfg_path = None
    try:
        with os.fdopen(fd, "w") as f:
            f.write(source)

        mermaid_cfg_path = build_mermaid_config(theme_cfg, bg)

        cmd = [
            mmdc,
            "-i", src_path,
            "-o", output_path,
            "--puppeteerConfigFile", puppeteer_cfg,
            "--configFile", mermaid_cfg_path,
            "--scale", str(scale),
        ]

        if bg == "transparent":
            cmd += ["-b", "transparent"]
        elif bg:
            cmd += ["-b", bg]
        elif "bg" in theme_cfg:
            cmd += ["-b", theme_cfg["bg"]]

        if width:
            cmd += ["-w", str(width)]

        if fmt == "pdf":
            cmd += ["--pdfFit"]

        log.debug("CMD: %s", " ".join(cmd))

        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=60,
        )

        duration_ms = int((time.perf_counter() - t0) * 1000)

        if verbose:
            if result.stdout.strip():
                print(f"  [mmdc stdout] {result.stdout.strip()}")
            if result.stderr.strip():
                print(f"  [mmdc stderr] {result.stderr.strip()}")

        success = result.returncode == 0 and os.path.isfile(output_path)
        error = None
        if not success:
            error = (result.stderr or result.stdout or "mmdc exited with code "
                     + str(result.returncode)).strip()

        return {
            "success": success,
            "output": output_path if success else None,
            "error": error,
            "duration_ms": duration_ms,
        }

    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "output": None,
            "error": "Render timed out (60s)",
            "duration_ms": int((time.perf_counter() - t0) * 1000),
        }
    except Exception as exc:
        return {
            "success": False,
            "output": None,
            "error": str(exc),
            "duration_ms": int((time.perf_counter() - t0) * 1000),
        }
    finally:
        try:
            os.unlink(src_path)
        except OSError:
            pass
        if mermaid_cfg_path:
            try:
                os.unlink(mermaid_cfg_path)
            except OSError:
                pass


# ---------------------------------------------------------------------------
# Batch helpers
# ---------------------------------------------------------------------------

def _batch_worker(args: tuple) -> dict:
    """Top-level function for ProcessPoolExecutor (must be picklable)."""
    (
        src_file, out_file, mmdc, puppeteer_cfg,
        theme, bg, scale, width, fmt, verbose,
    ) = args
    source = Path(src_file).read_text(encoding="utf-8")
    return {
        "input": src_file,
        **render_diagram(
            source, out_file,
            mmdc=mmdc, puppeteer_cfg=puppeteer_cfg,
            theme_name=theme, bg=bg, scale=scale,
            width=width, fmt=fmt, verbose=verbose,
        ),
    }


def resolve_batch_inputs(batch_path: str, fmt: str, out_dir: Optional[str]) -> list[tuple[str, str]]:
    """Return list of (input_path, output_path) tuples for a batch run."""
    p = Path(batch_path)
    pairs: list[tuple[str, str]] = []

    if "*" in batch_path or "?" in batch_path:
        inputs = sorted(glob.glob(batch_path))
    elif p.is_dir():
        inputs = sorted(
            str(f) for f in p.rglob("*")
            if f.suffix.lower() in {".mmd", ".mermaid", ".txt"}
            and f.is_file()
        )
    elif p.is_file():
        inputs = [str(p)]
    else:
        print(f"[ERROR] Batch path not found: {batch_path}", file=sys.stderr)
        return []

    for inp in inputs:
        in_path = Path(inp)
        if out_dir:
            out_base = Path(out_dir)
            out_base.mkdir(parents=True, exist_ok=True)
            out_file = str(out_base / (in_path.stem + f".{fmt}"))
        else:
            out_file = str(in_path.with_suffix(f".{fmt}"))
        pairs.append((inp, out_file))

    return pairs


# ---------------------------------------------------------------------------
# Watch mode
# ---------------------------------------------------------------------------

def watch_file(path: str) -> float:
    try:
        return os.path.getmtime(path)
    except OSError:
        return 0.0


def run_watch(
    input_files: list[str],
    output_map: dict[str, str],
    render_kwargs: dict,
    poll_interval: float = 0.8,
) -> None:
    """Poll input files and re-render on modification."""
    mtimes = {f: watch_file(f) for f in input_files}
    print(f"Watching {len(input_files)} file(s). Press Ctrl+C to stop.")
    try:
        while True:
            time.sleep(poll_interval)
            for f in input_files:
                mtime = watch_file(f)
                if mtime != mtimes[f]:
                    mtimes[f] = mtime
                    print(f"  Changed: {f} — re-rendering...")
                    source = Path(f).read_text(encoding="utf-8")
                    result = render_diagram(source, output_map[f], **render_kwargs)
                    if result["success"]:
                        print(f"  OK: {result['output']} ({result['duration_ms']}ms)")
                    else:
                        print(f"  FAIL: {result['error']}", file=sys.stderr)
    except KeyboardInterrupt:
        print("\nWatch mode stopped.")


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        prog="mermaid-renderer",
        description="Convert Mermaid diagrams to PNG, SVG, or PDF images.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )

    p.add_argument(
        "input",
        nargs="?",
        help="Path to .mmd / .mermaid file (omit when using --text or --batch)",
    )
    p.add_argument(
        "-o", "--output",
        default=None,
        help="Output file path (default: same name as input, auto-extension)",
    )
    p.add_argument(
        "--text",
        default=None,
        help="Inline Mermaid diagram text (alternative to file input)",
    )
    p.add_argument(
        "--batch",
        default=None,
        metavar="DIR_OR_GLOB",
        help="Render all .mmd/.mermaid files in a directory or matching a glob",
    )
    p.add_argument(
        "--format", "-f",
        default="png",
        choices=["png", "svg", "pdf"],
        help="Output format (default: png)",
    )
    p.add_argument(
        "--theme", "-t",
        default="default",
        choices=list(THEMES.keys()),
        help="Named theme (default: default)",
    )
    p.add_argument(
        "--bg", "--background",
        default=None,
        dest="bg",
        help='Background color hex or "transparent"',
    )
    p.add_argument(
        "--fg", "--foreground",
        default=None,
        dest="fg",
        help="Foreground color hex (overrides theme fg — informational only)",
    )
    p.add_argument(
        "--accent",
        default=None,
        help="Accent color hex (overrides theme accent — informational only)",
    )
    p.add_argument(
        "--scale",
        type=float,
        default=2.0,
        help="Device pixel ratio / scale factor (default: 2.0)",
    )
    p.add_argument(
        "--width", "-W",
        type=int,
        default=None,
        help="Canvas width in pixels",
    )
    p.add_argument(
        "--workers", "-j",
        type=int,
        default=1,
        help="Parallel workers for batch mode (default: 1)",
    )
    p.add_argument(
        "--watch", "-w",
        action="store_true",
        help="Watch input files and re-render on change",
    )
    p.add_argument(
        "--list-themes",
        action="store_true",
        help="Print all available themes and exit",
    )
    p.add_argument(
        "--check",
        action="store_true",
        help="Check environment dependencies and exit",
    )
    p.add_argument(
        "-v", "--verbose",
        action="store_true",
        help="Show mmdc stdout/stderr",
    )

    return p


def cmd_check() -> None:
    """Print environment diagnostic info."""
    mmdc = find_mmdc()
    chrome = find_chrome()

    print("mermaid-renderer — environment check")
    print("-" * 40)
    print(f"  Python      : {sys.version.split()[0]}")
    print(f"  mmdc        : {mmdc or 'NOT FOUND'}")
    print(f"  Chrome/Chrom: {chrome or 'NOT FOUND'}")

    if mmdc:
        try:
            r = subprocess.run([mmdc, "--version"], capture_output=True, text=True, timeout=5)
            print(f"  mmdc version: {r.stdout.strip()}")
        except Exception:
            print("  mmdc version: (could not query)")

    if mmdc and chrome:
        print("\n  All dependencies OK. Rendering should work.")
    else:
        missing = []
        if not mmdc:
            missing.append("mmdc (install: npm install -g @mermaid-js/mermaid-cli)")
        if not chrome:
            missing.append("Chrome/Chromium")
        print("\n  MISSING:", ", ".join(missing))
        sys.exit(1)


def cmd_list_themes() -> None:
    col_w = 24
    print(f"\n{'Theme':<{col_w}} {'Type':<8} {'Background':<12} {'Accent'}")
    print("-" * 64)
    for name, cfg in THEMES.items():
        theme_type = "dark" if cfg.get("mermaid_theme") == "dark" else "light"
        accent = cfg.get("accent", "(derived)")
        print(f"  {name:<{col_w-2}} {theme_type:<8} {cfg.get('bg',''):<12} {accent}")
    print()


def main() -> None:
    parser = build_parser()
    args = parser.parse_args()

    if args.verbose:
        log.setLevel(logging.DEBUG)

    if args.check:
        cmd_check()
        return

    if args.list_themes:
        cmd_list_themes()
        return

    # Resolve dependencies
    mmdc = find_mmdc()
    chrome = find_chrome()

    if not mmdc:
        print(
            "[ERROR] mmdc not found. Install with:\n"
            "  npm install -g @mermaid-js/mermaid-cli\n"
            "Then run:  python render.py --check",
            file=sys.stderr,
        )
        sys.exit(1)

    if not chrome:
        print(
            "[ERROR] No Chrome/Chromium found. See README for install instructions.",
            file=sys.stderr,
        )
        sys.exit(1)

    puppeteer_cfg = build_puppeteer_config(chrome)

    # Apply per-run theme overrides
    theme_cfg = dict(THEMES.get(args.theme, THEMES["default"]))
    if args.fg:
        theme_cfg["fg"] = args.fg
    if args.accent:
        theme_cfg["accent"] = args.accent

    render_kwargs = dict(
        mmdc=mmdc,
        puppeteer_cfg=puppeteer_cfg,
        theme_name=args.theme,
        bg=args.bg,
        scale=args.scale,
        width=args.width,
        fmt=args.format,
        verbose=args.verbose,
    )

    try:
        # -----------------------------------------------------------------------
        # BATCH MODE
        # -----------------------------------------------------------------------
        if args.batch:
            pairs = resolve_batch_inputs(args.batch, args.format, args.output)
            if not pairs:
                print("[ERROR] No input files found.", file=sys.stderr)
                sys.exit(1)

            print(f"Batch: {len(pairs)} file(s) | theme={args.theme} | "
                  f"format={args.format} | workers={args.workers}")

            if args.watch:
                input_files = [p[0] for p in pairs]
                output_map = {p[0]: p[1] for p in pairs}
                run_watch(input_files, output_map, render_kwargs)
                return

            if args.workers > 1:
                worker_args = [
                    (inp, out, mmdc, puppeteer_cfg,
                     args.theme, args.bg, args.scale, args.width,
                     args.format, args.verbose)
                    for inp, out in pairs
                ]
                results = []
                with ProcessPoolExecutor(max_workers=args.workers) as ex:
                    futures = {ex.submit(_batch_worker, a): a[0] for a in worker_args}
                    for fut in as_completed(futures):
                        results.append(fut.result())
            else:
                results = []
                for inp, out in pairs:
                    source = Path(inp).read_text(encoding="utf-8")
                    r = render_diagram(source, out, **render_kwargs)
                    results.append({"input": inp, **r})
                    status = "OK" if r["success"] else "FAIL"
                    print(f"  [{status}] {inp} -> {out} ({r['duration_ms']}ms)")
                    if not r["success"]:
                        print(f"         Error: {r['error']}", file=sys.stderr)

            ok = sum(1 for r in results if r["success"])
            fail = len(results) - ok
            print(f"\nDone: {ok} succeeded, {fail} failed.")
            if fail > 0:
                sys.exit(1)
            return

        # -----------------------------------------------------------------------
        # SINGLE FILE OR INLINE TEXT
        # -----------------------------------------------------------------------
        if args.text:
            source = args.text.replace("; ", "\n").replace(";", "\n")
            default_stem = "diagram"
        elif args.input:
            in_path = Path(args.input)
            if not in_path.is_file():
                print(f"[ERROR] Input file not found: {args.input}", file=sys.stderr)
                sys.exit(1)
            source = in_path.read_text(encoding="utf-8")
            default_stem = in_path.stem
        else:
            parser.print_help()
            sys.exit(0)

        # Determine output path
        if args.output:
            out_path = args.output
        else:
            suffix = f".{args.format}"
            out_path = str(Path(default_stem).with_suffix(suffix))

        print(f"Rendering: {out_path} | theme={args.theme} | format={args.format} | scale={args.scale}")

        result = render_diagram(source, out_path, **render_kwargs)

        if result["success"]:
            size = os.path.getsize(result["output"])
            print(f"  OK: {result['output']} ({size:,} bytes, {result['duration_ms']}ms)")
        else:
            print(f"  FAIL: {result['error']}", file=sys.stderr)
            sys.exit(1)

        # Watch mode for single file
        if args.watch and args.input:
            run_watch([args.input], {args.input: out_path}, render_kwargs)

    finally:
        try:
            os.unlink(puppeteer_cfg)
        except OSError:
            pass


if __name__ == "__main__":
    main()