mnemosyne/mnemosyne/library/services/svg_raster.py

"""SVG rasterization for the ingest pipeline.

The vision stage sends each extracted image to a vision LLM as a ``data:`` URI,
which cannot carry ``image/svg+xml`` — so an SVG has to become raster before it
can be described, OCR'd, or embedded. PyMuPDF (already a dependency for PDF
parsing) renders SVG natively, so this needs no cairo/rsvg system libraries.

The non-obvious part is page splitting. Handwritten notes from the Write
(Stylus Labs) app are a single SVG document holding one or more
``<svg class="write-page">`` children stacked vertically via x/y offsets. Two
root formats exist in the wild and *both* must be split per page:

- Older files carry no width/height on the root ``<svg>``. Rendering the
  document as-is makes PyMuPDF fall back to US-Letter and emit the top-left
  corner only — ruled lines and no handwriting, in a perfectly valid PNG.
- Newer files (Write commit eeab021) do carry root width/height, spanning the
  full stacked extent. Rendering those as-is yields one very tall strip; capped
  to a sane longest edge, a 5-page note squashes to ~209px wide, well past
  illegible.

Write never rewrites existing files, so the old format is permanent rather than
a migration window. Page geometry lives on the ``write-page`` element and is
identical across both formats, so splitting ignores the root dimensions
entirely and needs no format detection.

.. note::
   Daedalus carries a twin of this module at
   ``backend/daedalus/extraction/svg.py``, which returns base64 for direct
   chat attachment. The two are deliberately duplicated rather than shared —
   the repos ship no common package — so fixes belong in both.
"""

from __future__ import annotations

import copy
import logging
import re

from lxml import etree

logger = logging.getLogger(__name__)

SVG_NS = "http://www.w3.org/2000/svg"
XLINK_NS = "http://www.w3.org/1999/xlink"

_SVG = f"{{{SVG_NS}}}"
_PAGE_CLASS = "write-page"

# Write documents render on a grey backdrop; pages themselves are transparent,
# so each rendered page needs an explicit white underlay or strokes land on
# black in the flattened PNG.
_PAGE_BACKGROUND = "#ffffff"

#: Longest edge of a rendered page, in pixels. Enough for a vision model to
#: read handwriting without spending tokens on unusable resolution.
DEFAULT_TARGET_PX = 1568

#: Cap on pages rendered from one document.
DEFAULT_MAX_PAGES = 20

# Unquoted attribute value in the root tag, e.g. Write's `width=auto`. Matches
# only bare alphabetic values so numeric or already-quoted attributes are left
# alone.
_UNQUOTED_ATTR = re.compile(rb"(\s[-\w:]+)=([A-Za-z][-\w]*)(?=[\s>])")
_ROOT_TAG = re.compile(rb"<svg[^>]*>")


class SvgRenderError(Exception):
    """Raised when an SVG cannot be parsed or contains no renderable page."""


def _repair_root_tag(data: bytes) -> bytes:
    """Quote unquoted attribute values in the root ``<svg>`` tag.

    Some Write files emit ``width=auto height=auto``, which is not valid XML.
    A malformation *inside* the root tag defeats ``recover=True`` differently
    from one in the body: rather than dropping a subtree, libxml2 abandons the
    whole document and yields a bare root, so every page becomes invisible.
    Quoting the values first recovers the full tree.
    """
    match = _ROOT_TAG.search(data)
    if not match:
        return data
    repaired = _UNQUOTED_ATTR.sub(rb'\1="\2"', match.group(0))
    if repaired == match.group(0):
        return data
    return data[: match.start()] + repaired + data[match.end() :]


def _parser() -> etree.XMLParser:
    """Build the hardened parser used for all untrusted SVG input.

    ``resolve_entities=False`` blocks XXE — ingest content is untrusted.
    ``huge_tree`` is required because handwriting path data runs to megabytes.
    ``recover`` salvages the handful of Write files that emit unescaped
    attribute content and are not well-formed XML.
    """
    return etree.XMLParser(
        huge_tree=True,
        resolve_entities=False,
        no_network=True,
        recover=True,
    )


def _dimension(element: etree._Element, name: str) -> float | None:
    """Read a CSS-pixel dimension attribute, tolerating a ``px`` suffix."""
    raw = element.get(name)
    if not raw:
        return None
    try:
        return float(raw.strip().removesuffix("px"))
    except ValueError:
        return None


def _viewbox_size(element: etree._Element) -> tuple[float, float] | None:
    """Derive width/height from a viewBox extent."""
    raw = element.get("viewBox")
    if not raw:
        return None
    parts = raw.replace(",", " ").split()
    if len(parts) != 4:
        return None
    try:
        width, height = float(parts[2]), float(parts[3])
    except ValueError:
        return None
    return (width, height) if width > 0 and height > 0 else None


def _element_size(element: etree._Element) -> tuple[float, float] | None:
    """Resolve an element's rendered size from width/height, else viewBox."""
    width = _dimension(element, "width")
    height = _dimension(element, "height")
    if width and height:
        return width, height
    return _viewbox_size(element)


def _standalone_page(
    page: etree._Element,
    defs: etree._Element | None,
    width: float,
    height: float,
) -> bytes:
    """Wrap one ``write-page`` element as its own renderable SVG document.

    The page is repositioned to the origin (its x/y place it within the stacked
    parent) and given an explicit viewBox so the renderer has an unambiguous
    size. ``defs`` is copied in because pen and ruling definitions live on the
    root and are referenced by page content.
    """
    root = etree.Element(f"{_SVG}svg", nsmap={None: SVG_NS, "xlink": XLINK_NS})
    root.set("width", str(width))
    root.set("height", str(height))
    root.set("viewBox", f"0 0 {width} {height}")

    background = etree.SubElement(root, f"{_SVG}rect")
    background.set("width", "100%")
    background.set("height", "100%")
    background.set("fill", _PAGE_BACKGROUND)

    if defs is not None:
        root.append(copy.deepcopy(defs))

    element = copy.deepcopy(page)
    for positional in ("x", "y"):
        element.attrib.pop(positional, None)
    element.set("viewBox", f"0 0 {width} {height}")
    root.append(element)

    return etree.tostring(root)


def split_svg_pages(data: bytes) -> list[bytes]:
    """Split an SVG into one standalone document per renderable page.

    Write multi-page notes yield one document per ``write-page`` child. Any
    other SVG (a diagram, a logo) yields a single document — the input itself,
    which the renderer sizes from its own width/height or viewBox.

    :param data: Raw SVG bytes.
    :returns: One or more standalone SVG documents, in page order.
    :raises SvgRenderError: If the SVG cannot be parsed or has no usable size.
    """
    try:
        root = etree.fromstring(_repair_root_tag(data), _parser())
    except etree.XMLSyntaxError as exc:
        raise SvgRenderError(f"Could not parse SVG: {exc}") from exc

    if root is None:
        raise SvgRenderError("Could not parse SVG: no root element.")

    defs = root.find(f"{_SVG}defs")

    pages: list[bytes] = []
    for element in root.iter(f"{_SVG}svg"):
        if element is root:
            continue
        if _PAGE_CLASS not in (element.get("class") or "").split():
            continue
        size = _element_size(element)
        if not size:
            continue
        pages.append(_standalone_page(element, defs, *size))

    if pages:
        return pages

    # Generic SVG: render as-is. It must still be sizeable, or the renderer
    # would silently substitute a default page box.
    if not _element_size(root):
        raise SvgRenderError(
            "SVG has no width/height or viewBox, so its size is undefined."
        )
    return [data]


def render_svg_pages(
    data: bytes,
    max_pages: int = DEFAULT_MAX_PAGES,
    target_px: int = DEFAULT_TARGET_PX,
) -> list[tuple[bytes, int, int]]:
    """Rasterize an SVG to one PNG per page.

    :param data: Raw SVG bytes.
    :param max_pages: Cap on pages rendered.
    :param target_px: Longest edge of each rendered page, in pixels.
    :returns: One ``(png_bytes, width, height)`` tuple per page, in page order.
    :raises SvgRenderError: If nothing could be rendered.
    """
    import fitz

    pages = split_svg_pages(data)

    rendered: list[tuple[bytes, int, int]] = []
    for index, page_svg in enumerate(pages[:max_pages]):
        try:
            with fitz.open(stream=page_svg, filetype="svg") as document:
                page = document[0]
                longest = max(page.rect.width, page.rect.height)
                zoom = target_px / longest if longest else 1.0
                pixmap = page.get_pixmap(matrix=fitz.Matrix(zoom, zoom))
                rendered.append(
                    (pixmap.tobytes("png"), pixmap.width, pixmap.height)
                )
        except Exception as exc:
            raise SvgRenderError(
                f"Could not render SVG page {index + 1}: {exc}"
            ) from exc

    if not rendered:
        raise SvgRenderError("No renderable pages in SVG.")

    logger.info(
        "Rasterized SVG pages_rendered=%d total_pages=%d target_px=%d",
        len(rendered),
        len(pages),
        target_px,
    )

    return rendered