md_to_gost/md2gost/diagram_renderer.py

"""Render UML/BPMN/C4 fenced blocks to PNG (PlantUML jar / Kroki local / remote)."""

from __future__ import annotations

import hashlib
import logging
import os
import shutil
import subprocess
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import Literal

import requests

from . import package_dir

DIAGRAM_LANGS = frozenset({"uml", "plantuml", "bpmn", "c4"})
FallbackMode = Literal["local", "remote", "off"]

DEFAULT_KROKI_URL = "http://localhost:8000"
REMOTE_KROKI_URL = "https://kroki.io"

_log = logging.getLogger(__name__)


@dataclass
class DiagramConfig:
    plantuml_jar: str | None = None
    kroki_url: str | None = None
    fallback: FallbackMode = "remote"
    cache_dir: str | None = None


_CONFIG = DiagramConfig()


def configure_diagrams(
    plantuml_jar: str | None = None,
    kroki_url: str | None = None,
    fallback: FallbackMode = "remote",
    cache_dir: str | None = None,
) -> None:
    global _CONFIG
    _CONFIG = DiagramConfig(
        plantuml_jar=plantuml_jar or os.environ.get("PLANTUML_JAR"),
        kroki_url=kroki_url or os.environ.get("KROKI_URL"),
        fallback=fallback,
        cache_dir=cache_dir,
    )


def diagrams_dir() -> Path:
    return Path(package_dir()) / "diagrams"


def prepare_source(lang: str, source: str) -> tuple[str, str]:
    """Return (prepared_source, kroki_diagram_type)."""
    lang = (lang or "uml").lower().strip()
    text = source.strip()

    if lang == "bpmn":
        if "@startbpmn" not in text.lower() and "<definitions" not in text.lower():
            # PlantUML BPMN dialect
            if not text.startswith("@start"):
                text = "@startbpmn\n" + text + "\n@endbpmn"
        return text, "bpmn"

    if lang == "c4":
        if "!include" not in text and "!includeurl" not in text.lower():
            # Prefer PlantUML stdlib; also ship local stubs for jar -I path
            includes = (
                f"!include {diagrams_dir() / 'C4_Container.puml'}\n"
                if (diagrams_dir() / "C4_Container.puml").exists()
                else "!include <C4/C4_Container>\n"
            )
            body = text
            if body.lower().startswith("@startuml"):
                lines = body.splitlines()
                text = lines[0] + "\n" + includes + "\n".join(lines[1:])
            else:
                text = f"@startuml\n{includes}{body}\n@enduml"
        elif not text.lower().startswith("@start"):
            text = f"@startuml\n{text}\n@enduml"
        return text, "plantuml"

    # uml / plantuml
    if not text.lower().startswith("@start"):
        text = f"@startuml\n{text}\n@enduml"
    return text, "plantuml"


def _cache_path(source: str, cache_dir: str | None = None) -> Path:
    root = cache_dir or _CONFIG.cache_dir
    if not root:
        wd = os.environ.get("WORKING_DIR", ".")
        root = os.path.join(wd, ".md2gost-cache")
    Path(root).mkdir(parents=True, exist_ok=True)
    digest = hashlib.sha256(source.encode("utf-8")).hexdigest()[:24]
    return Path(root) / f"{digest}.png"


def _java_available() -> bool:
    return shutil.which("java") is not None


def _render_plantuml_jar(source: str, out_png: Path, jar: str) -> bool:
    if not jar or not os.path.isfile(jar) or not _java_available():
        return False
    with tempfile.TemporaryDirectory() as tmp:
        src = Path(tmp) / "diagram.puml"
        src.write_text(source, encoding="utf-8")
        cmd = [
            "java", "-jar", jar,
            "-tpng",
            "-charset", "UTF-8",
            f"-I{diagrams_dir()}",
            "-o", tmp,
            str(src),
        ]
        try:
            proc = subprocess.run(
                cmd, capture_output=True, text=True, timeout=120, check=False,
            )
        except (OSError, subprocess.TimeoutExpired) as e:
            _log.warning("PlantUML jar failed: %s", e)
            return False
        produced = Path(tmp) / "diagram.png"
        if proc.returncode != 0 or not produced.is_file():
            _log.warning("PlantUML jar error: %s", proc.stderr or proc.stdout)
            return False
        shutil.copyfile(produced, out_png)
        return True


def _render_kroki(source: str, diagram_type: str, base_url: str, out_png: Path) -> bool:
    url = base_url.rstrip("/") + f"/{diagram_type}/png"
    try:
        resp = requests.post(
            url,
            data=source.encode("utf-8"),
            headers={"Content-Type": "text/plain"},
            timeout=60,
        )
        if resp.status_code != 200 or not resp.content.startswith(b"\x89PNG"):
            _log.warning("Kroki %s → HTTP %s", url, resp.status_code)
            return False
        out_png.write_bytes(resp.content)
        return True
    except requests.RequestException as e:
        _log.warning("Kroki request failed (%s): %s", url, e)
        return False


def render_diagram(
    lang: str,
    source: str,
    *,
    plantuml_jar: str | None = None,
    kroki_url: str | None = None,
    fallback: FallbackMode | None = None,
    cache_dir: str | None = None,
) -> str:
    """Render diagram to PNG path. Raises RuntimeError if all backends fail."""
    prepared, kroki_type = prepare_source(lang, source)
    out = _cache_path(prepared, cache_dir)
    if out.is_file() and out.stat().st_size > 0:
        return str(out)

    jar = plantuml_jar if plantuml_jar is not None else _CONFIG.plantuml_jar
    jar = jar or os.environ.get("PLANTUML_JAR")
    local_kroki = kroki_url if kroki_url is not None else _CONFIG.kroki_url
    local_kroki = local_kroki or os.environ.get("KROKI_URL") or DEFAULT_KROKI_URL
    mode: FallbackMode = fallback if fallback is not None else _CONFIG.fallback

    # 1) PlantUML jar
    if _render_plantuml_jar(prepared, out, jar or ""):
        return str(out)

    # 2) Local Kroki
    if _render_kroki(prepared, kroki_type, local_kroki, out):
        return str(out)

    if mode == "off":
        raise RuntimeError(
            "Не удалось отрендерить диаграмму локально "
            "(задайте --plantuml-jar или KROKI_URL; remote fallback отключён)"
        )
    if mode == "local":
        raise RuntimeError(
            "Локальный рендер диаграммы недоступен "
            "(Java+plantuml.jar или локальный Kroki)"
        )

    # 3) Remote fallback
    _log.warning(
        "Диаграмма: локальный рендер недоступен, использую remote %s",
        REMOTE_KROKI_URL,
    )
    if _render_kroki(prepared, kroki_type, REMOTE_KROKI_URL, out):
        return str(out)

    raise RuntimeError("Не удалось отрендерить диаграмму (PlantUML/Kroki)")