# tractor: distributed structured concurrency.
'''
``.. d2::`` - embed d2lang_ diagrams in sphinx docs
with build-time rendering and a committed-SVG
fallback.

Rendering policy,

- when a ``d2`` binary is found (see discovery order
  below) any out-of-date SVG is (re)rendered from its
  ``.d2`` source (normally kept in ``docs/diagrams/``)
  into ``docs/_diagrams/``,
- otherwise any pre-rendered (and git committed) SVG
  already in ``docs/_diagrams/`` is used as-is,
- when neither is possible the diagram *source* is
  emitted as a literal block so no content is ever
  silently dropped.

A render that is *attempted and fails* (a ``d2`` bin is
present but errors on the source) is NOT silently
degraded to the stale committed SVG: it raises a
docutils error (so ``sphinx-build -W`` fails the
build). The render is also atomic — a failed/torn
render can never clobber a good committed SVG — so the
last-good diagram survives a bad edit.

Binary discovery order,

- the ``D2_BIN`` env var; may contain args which are
  split via `shlex`, eg.
  ``D2_BIN='nix run nixpkgs#d2 --'``,
- the ``d2_bin`` sphinx config value,
- ``shutil.which('d2')``.

Usage,

.. code:: rst

    .. d2:: diagrams/actor_tree.d2
       :caption: A tree of ``trio``-task-trees.
       :margin:

.. _d2lang: https://d2lang.com

'''
from __future__ import annotations

import enum
import hashlib
import os
from pathlib import Path
import re
import shlex
import shutil
import subprocess as sp
import tempfile

from docutils import nodes
from docutils.parsers.rst import directives
from sphinx.application import Sphinx
from sphinx.util import logging
from sphinx.util.docutils import SphinxDirective

log = logging.getLogger(__name__)

# subdir (under the sphinx srcdir) holding rendered,
# git-committed, fallback SVG outputs.
_outdir: str = '_diagrams'

# per-build map of {output-svg-name: source-relpath} used
# to detect 2 distinct `.d2` sources colliding on a single
# output stem; reset on each `builder-inited` (see `setup`).
_seen_outputs: dict[str, str] = {}


class RenderResult(enum.Enum):
    '''
    Outcome of a `render_svg()` call.

    '''
    OK = 'ok'                # fresh SVG exists on disk
    FELL_BACK = 'fell_back'  # no bin; serving committed SVG
    NO_OUTPUT = 'no_output'  # no bin AND no committed SVG
    FAILED = 'failed'        # bin present, render errored


def find_d2(
    app: Sphinx,
) -> list[str]|None:
    '''
    Resolve the d2 render command as an argv list or
    `None` when no binary can be found.

    '''
    if env_bin := os.environ.get('D2_BIN'):
        return shlex.split(env_bin)
    if cfg_bin := app.config.d2_bin:
        return shlex.split(cfg_bin)
    if path_bin := shutil.which('d2'):
        return [path_bin]
    return None


def render_svg(
    app: Sphinx,
    src: Path,
    out: Path,
) -> RenderResult:
    '''
    Maybe (re)render `src` -> `out` and report the
    outcome as a `RenderResult`.

    A render is performed only when `out` is missing or
    older than `src` AND a `d2` binary is available. The
    write is atomic (temp-file + `os.replace`) so a
    failed or torn render never clobbers an existing
    (committed) SVG.

    '''
    fresh: bool = (
        out.exists()
        and
        src.stat().st_mtime <= out.stat().st_mtime
    )
    if fresh:
        return RenderResult.OK
    d2cmd: list[str]|None = find_d2(app)
    if d2cmd is None:
        # no binary: fall back to whatever is committed,
        # else signal the caller to emit the raw source.
        if out.exists():
            log.info(
                f'no d2 binary; using committed svg '
                f'for {src.name}'
            )
            return RenderResult.FELL_BACK
        return RenderResult.NO_OUTPUT
    out.parent.mkdir(
        parents=True,
        exist_ok=True,
    )
    # render into a sibling temp-file first so a bad
    # `d2` exit (or a crash mid-write) leaves any prior
    # committed SVG fully intact.
    fd, tmpname = tempfile.mkstemp(
        dir=str(out.parent),
        prefix=f'.{out.stem}.',
        suffix='.svg.tmp',
    )
    os.close(fd)
    tmp = Path(tmpname)
    argv: list[str] = (
        d2cmd
        + list(app.config.d2_args)
        + [str(src), str(tmp)]
    )
    try:
        proc = sp.run(
            argv,
            capture_output=True,
            text=True,
            timeout=120,
        )
    except (
        OSError,
        sp.TimeoutExpired,
    ) as err:
        tmp.unlink(missing_ok=True)
        log.warning(
            f'd2 invocation failed for {src.name}: '
            f'{err}'
        )
        return RenderResult.FAILED
    if proc.returncode != 0:
        tmp.unlink(missing_ok=True)
        log.warning(
            f'd2 render error for {src.name}:\n'
            f'{proc.stderr}'
        )
        return RenderResult.FAILED
    # atomic swap-in of the freshly rendered SVG.
    os.replace(tmp, out)
    return RenderResult.OK


class D2Diagram(SphinxDirective):
    '''
    Render a ``.d2`` source file (path relative to
    the sphinx srcdir) as an SVG figure.

    '''
    required_arguments = 1
    has_content = False
    option_spec = {
        'caption': directives.unchanged,
        'alt': directives.unchanged,
        'width': (
            directives.length_or_percentage_or_unitless
        ),
        'margin': directives.flag,
        'class': directives.class_option,
        'name': directives.unchanged,
    }

    def run(self) -> list[nodes.Node]:
        relsrc: str = self.arguments[0]
        srcdir = Path(self.env.srcdir)
        src: Path = srcdir / relsrc
        self.env.note_dependency(relsrc)
        if not src.exists():
            err = self.state_machine.reporter.error(
                f'd2 source not found: {relsrc}',
                line=self.lineno,
            )
            return [err]
        out: Path = (
            srcdir
            / _outdir
            / f'{src.stem}.svg'
        )
        # collision guard: two distinct sources must not
        # map onto the same output stem within a build.
        prior: str|None = _seen_outputs.get(out.name)
        if (
            prior is not None
            and
            prior != relsrc
        ):
            err = self.state_machine.reporter.error(
                f'd2 output collision: {relsrc!r} and '
                f'{prior!r} both render to '
                f'{_outdir}/{out.name}; rename one '
                f'`.d2` stem',
                line=self.lineno,
            )
            return [err]
        _seen_outputs[out.name] = relsrc

        result: RenderResult = render_svg(
            self.env.app,
            src,
            out,
        )
        if result is RenderResult.FAILED:
            # loud, build-failing (under `-W`) signal; the
            # prior committed SVG, if any, is untouched.
            err = self.state_machine.reporter.error(
                f'd2 render failed for {relsrc} (see '
                f'build log); the committed svg, if any, '
                f'was left untouched',
                line=self.lineno,
            )
            return [err]
        if result is RenderResult.NO_OUTPUT:
            # last resort: emit the raw d2 source so no
            # content is ever silently dropped.
            log.warning(
                f'no svg available for {relsrc}; '
                f'emitting d2 source as literal block'
            )
            src_text: str = src.read_text()
            literal = nodes.literal_block(
                src_text,
                src_text,
            )
            literal['language'] = 'text'
            return [literal]
        # OK | FELL_BACK -> embed the SVG figure.
        img = nodes.image(
            uri=f'/{_outdir}/{out.name}',
            alt=self.options.get(
                'alt',
                f'd2 diagram: {src.stem}',
            ),
        )
        if width := self.options.get('width'):
            img['width'] = width
        # content-hash the rendered svg so the served
        # `<img src>` carries a `?v=<hash>` cache-buster
        # (sphinx tags css/js this way but not images);
        # without it an edited diagram keeps showing the
        # browser's stale copy at the stable
        # `_images/<stem>.svg` url on autobuild reload.
        if out.exists():
            img['d2_stem'] = out.stem
            img['d2_cachebust'] = hashlib.sha256(
                out.read_bytes()
            ).hexdigest()[:8]
        fig = nodes.figure()
        fig += img
        classes: list[str] = (
            ['d2-diagram']
            + self.options.get('class', [])
        )
        if 'margin' in self.options:
            # NB: the bare 'margin' class is what
            # book-style themes key off for
            # right-margin placement; our custom css
            # uses 'd2-margin'.
            classes += [
                'margin',
                'd2-margin',
            ]
        fig['classes'] += classes
        if caption_txt := self.options.get('caption'):
            (
                inline_nodes,
                _msgs,
            ) = self.state.inline_text(
                caption_txt,
                self.lineno,
            )
            caption = nodes.caption(
                caption_txt,
                '',
                *inline_nodes,
            )
            fig += caption
        self.add_name(fig)
        return [fig]


def _reset_seen(
    app: Sphinx,
) -> None:
    '''
    Clear the per-build output-collision map at the
    start of each build.

    '''
    _seen_outputs.clear()


def _cachebust_d2_images(
    app: Sphinx,
    pagename: str,
    templatename: str,
    context: dict,
    doctree,
) -> None:
    '''
    Append a ``?v=<content-hash>`` query to every d2
    diagram ``<img src>`` in the rendered page ``body``
    so a browser re-fetches an *edited* SVG instead of
    its stale cached copy at the stable
    ``_images/<stem>.svg`` url (sphinx cache-busts
    css/js this way, but not images).

    Scoped strictly to d2 images (those carrying the
    `d2_cachebust` attr set in `D2Diagram.run()`); all
    other markup is left untouched.

    '''
    if doctree is None or 'body' not in context:
        return
    busts: dict[str, str] = {}
    for img in doctree.findall(nodes.image):
        if ver := img.get('d2_cachebust'):
            busts[img['d2_stem']] = ver
    if not busts:
        return
    body: str = context['body']
    for stem, ver in busts.items():
        body = re.sub(
            r'(src="[^"]*?' + re.escape(stem) + r'\.svg)(")',
            rf'\1?v={ver}\2',
            body,
        )
    context['body'] = body


def setup(app: Sphinx) -> dict:
    app.add_config_value('d2_bin', None, 'env')
    app.add_config_value('d2_args', [], 'env')
    app.add_directive('d2', D2Diagram)
    app.connect('builder-inited', _reset_seen)
    app.connect('html-page-context', _cachebust_d2_images)
    return {
        'version': '0.1.0',
        # NB: run() writes the rendered SVG during the
        # READ phase; the temp-file + `os.replace` swap
        # keeps that atomic so concurrent renders of the
        # same diagram (under `sphinx-build -j`) can't
        # tear the output file.
        'parallel_read_safe': True,
        'parallel_write_safe': True,
    }