src/editable_image/svg_export.py

from __future__ import annotations

from urllib.parse import quote
from xml.etree import ElementTree as ET

from .fonts import font_data_uri, get_face
from .models import (
    BitmapRegion,
    DocumentSpec,
    EllipseGeometry,
    LinearGradientPaint,
    PathGeometry,
    PolygonGeometry,
    RectGeometry,
    ShapeRegion,
    SolidPaint,
    TextRegion,
)


SVG_NS = "http://www.w3.org/2000/svg"
XLINK_NS = "http://www.w3.org/1999/xlink"
ET.register_namespace("", SVG_NS)
ET.register_namespace("xlink", XLINK_NS)


def build_svg(
    document: DocumentSpec,
    background_href: str | None = None,
    background_fallback_href: str | None = None,
    embed_fonts: bool = False,
    bitmap_hrefs: dict[str, tuple[str, str]] | None = None,
) -> str:
    root = ET.Element(
        f"{{{SVG_NS}}}svg",
        {
            "width": str(document.canvas.width),
            "height": str(document.canvas.height),
            "viewBox": f"0 0 {document.canvas.width} {document.canvas.height}",
            "version": "1.1",
        },
    )
    bitmap_hrefs = bitmap_hrefs or {}
    needs_markers = any(
        isinstance(shape.geometry, PathGeometry)
        and (shape.geometry.arrow_start or shape.geometry.arrow_end)
        for shape in document.shapes
    )
    if embed_fonts or document.shapes or needs_markers:
        definitions = ET.SubElement(root, f"{{{SVG_NS}}}defs")
    else:
        definitions = None
    if embed_fonts:
        used = {
            (region.style.font_family, region.style.font_weight, region.style.italic)
            for region in document.regions
            if region.include_overlay and region.text
        }
        assert definitions is not None
        style = ET.SubElement(definitions, f"{{{SVG_NS}}}style", {"type": "text/css"})
        style.text = "\n" + "\n".join(font_rule(*item) for item in sorted(used)) + "\n"
    if definitions is not None:
        for shape in document.shapes:
            if isinstance(shape.style.fill, LinearGradientPaint):
                definitions.append(gradient_element(shape))
        if needs_markers:
            add_arrow_marker(definitions)

    if background_href:
        relative_href = encoded_image_href(background_href)
        fallback_href = encoded_image_href(background_fallback_href or background_href)
        ET.SubElement(
            root,
            f"{{{SVG_NS}}}image",
            {
                "id": "cleared-background",
                "href": relative_href,
                f"{{{XLINK_NS}}}href": fallback_href,
                "x": "0",
                "y": "0",
                "width": str(document.canvas.width),
                "height": str(document.canvas.height),
                "preserveAspectRatio": "none",
            },
        )

    layer = ET.SubElement(root, f"{{{SVG_NS}}}g", {"id": "editable-content"})
    content: list[tuple[int, str, ET.Element]] = []
    for shape in document.shapes:
        if shape.include_overlay:
            content.append((shape.z_index, shape.id, shape_element(shape)))
    for bitmap in document.bitmap_regions:
        hrefs = bitmap_hrefs.get(bitmap.id)
        if hrefs:
            content.append((bitmap.z_index, bitmap.id, bitmap_element(bitmap, *hrefs)))
    for _, _, element in sorted(content, key=lambda item: (item[0], item[1])):
        layer.append(element)
    text_layer = ET.SubElement(root, f"{{{SVG_NS}}}g", {"id": "editable-text"})
    for region in sorted(document.regions, key=lambda item: (item.z_index, item.id)):
        if region.include_overlay and region.text:
            text_layer.append(text_element(region))
    return ET.tostring(root, encoding="unicode", xml_declaration=True)


def gradient_element(shape: ShapeRegion) -> ET.Element:
    paint = shape.style.fill
    assert isinstance(paint, LinearGradientPaint)
    import math

    radians = math.radians(paint.angle)
    dx, dy = math.cos(radians) * 50, math.sin(radians) * 50
    gradient = ET.Element(
        f"{{{SVG_NS}}}linearGradient",
        {
            "id": f"gradient-{shape.id}",
            "x1": f"{compact(50 - dx)}%",
            "y1": f"{compact(50 - dy)}%",
            "x2": f"{compact(50 + dx)}%",
            "y2": f"{compact(50 + dy)}%",
        },
    )
    for stop in paint.stops:
        ET.SubElement(
            gradient,
            f"{{{SVG_NS}}}stop",
            {
                "offset": f"{compact(stop.offset * 100)}%",
                "stop-color": stop.color,
                "stop-opacity": compact(stop.opacity),
            },
        )
    return gradient


def add_arrow_marker(definitions: ET.Element) -> None:
    marker = ET.SubElement(
        definitions,
        f"{{{SVG_NS}}}marker",
        {
            "id": "editable-arrow",
            "viewBox": "0 0 10 10",
            "refX": "9",
            "refY": "5",
            "markerWidth": "6",
            "markerHeight": "6",
            "orient": "auto-start-reverse",
        },
    )
    ET.SubElement(marker, f"{{{SVG_NS}}}path", {"d": "M 0 0 L 10 5 L 0 10 z", "fill": "context-stroke"})


def shape_element(shape: ShapeRegion) -> ET.Element:
    geometry = shape.geometry
    style = shape.style
    fill = (
        style.fill.color
        if isinstance(style.fill, SolidPaint)
        else f"url(#gradient-{shape.id})"
    )
    attributes = {
        "id": shape.id,
        "data-role": shape.semantic_role,
        "fill": fill,
        "fill-opacity": compact(style.fill_opacity),
        "stroke": style.stroke,
        "stroke-width": compact(style.stroke_width),
        "stroke-opacity": compact(style.stroke_opacity),
        "stroke-linecap": style.linecap,
        "stroke-linejoin": style.linejoin,
        "data-confidence": compact(shape.confidence),
        "data-source-contours": " | ".join(
            " ".join(f"{compact(point.x)},{compact(point.y)}" for point in contour)
            for contour in shape.source_contours
        ),
    }
    if isinstance(geometry, RectGeometry):
        box = geometry.box
        attributes.update(
            {
                "x": compact(box.x),
                "y": compact(box.y),
                "width": compact(box.width),
                "height": compact(box.height),
                "rx": compact(geometry.rx),
                "ry": compact(geometry.ry),
            }
        )
        if abs(geometry.rotation) > 0.01:
            cx, cy = box.x + box.width / 2, box.y + box.height / 2
            attributes["transform"] = f"rotate({compact(geometry.rotation)} {compact(cx)} {compact(cy)})"
        return ET.Element(f"{{{SVG_NS}}}rect", attributes)
    if isinstance(geometry, EllipseGeometry):
        attributes.update(
            {
                "cx": compact(geometry.cx),
                "cy": compact(geometry.cy),
                "rx": compact(geometry.rx),
                "ry": compact(geometry.ry),
            }
        )
        if abs(geometry.rotation) > 0.01:
            attributes["transform"] = (
                f"rotate({compact(geometry.rotation)} {compact(geometry.cx)} {compact(geometry.cy)})"
            )
        return ET.Element(f"{{{SVG_NS}}}ellipse", attributes)
    if isinstance(geometry, PolygonGeometry):
        attributes["points"] = " ".join(
            f"{compact(point.x)},{compact(point.y)}" for point in geometry.points
        )
        return ET.Element(f"{{{SVG_NS}}}polygon", attributes)
    assert isinstance(geometry, PathGeometry)
    attributes["d"] = path_data(geometry)
    if geometry.arrow_start:
        attributes["marker-start"] = "url(#editable-arrow)"
    if geometry.arrow_end:
        attributes["marker-end"] = "url(#editable-arrow)"
    if not geometry.closed and fill != "none":
        attributes["fill"] = "none"
    return ET.Element(f"{{{SVG_NS}}}path", attributes)


def path_data(geometry: PathGeometry) -> str:
    points = geometry.points
    if not geometry.smooth or len(points) < 3:
        output = f"M {compact(points[0].x)} {compact(points[0].y)}"
        output += "".join(f" L {compact(point.x)} {compact(point.y)}" for point in points[1:])
    else:
        output = f"M {compact(points[0].x)} {compact(points[0].y)}"
        for current, following in zip(points[1:-1], points[2:]):
            midpoint_x = (current.x + following.x) / 2
            midpoint_y = (current.y + following.y) / 2
            output += (
                f" Q {compact(current.x)} {compact(current.y)}"
                f" {compact(midpoint_x)} {compact(midpoint_y)}"
            )
        output += f" T {compact(points[-1].x)} {compact(points[-1].y)}"
    if geometry.closed:
        output += " Z"
    return output


def bitmap_element(bitmap: BitmapRegion, href: str, fallback_href: str) -> ET.Element:
    box = bitmap.box
    return ET.Element(
        f"{{{SVG_NS}}}image",
        {
            "id": bitmap.id,
            "href": encoded_image_href(href),
            "x": compact(box.x),
            "y": compact(box.y),
            "width": compact(box.width),
            "height": compact(box.height),
            "opacity": compact(bitmap.opacity),
            "preserveAspectRatio": "none",
        },
    )


def font_rule(family: str, weight: int, italic: bool) -> str:
    face = get_face(family, weight, italic)
    style = "italic" if italic else "normal"
    return (
        f"@font-face {{font-family:'{face.family}';font-style:{style};"
        f"font-weight:{face.weight};src:url('{font_data_uri(face.filename)}') format('woff2');}}"
    )


def text_element(region: TextRegion) -> ET.Element:
    box = region.overlay_box
    style = region.style
    anchor_x = {
        "start": box.x,
        "middle": box.x + box.width / 2,
        "end": box.x + box.width,
    }[style.anchor]
    baseline = box.y + box.height * 0.82
    attributes = {
        "id": region.id,
        "x": compact(anchor_x),
        "y": compact(baseline),
        "font-family": style.font_family,
        "font-size": compact(style.font_size),
        "font-weight": str(style.font_weight),
        "font-style": "italic" if style.italic else "normal",
        "fill": style.fill,
        "fill-opacity": compact(style.opacity),
        "stroke": style.stroke,
        "stroke-width": compact(style.stroke_width),
        "letter-spacing": compact(style.letter_spacing),
        "text-anchor": style.anchor,
        "data-confidence": compact(region.confidence),
        "data-source-quad": " ".join(f"{compact(p.x)},{compact(p.y)}" for p in region.source_quad),
    }
    if abs(region.rotation) > 0.01:
        attributes["transform"] = f"rotate({compact(region.rotation)} {compact(anchor_x)} {compact(baseline)})"
    element = ET.Element(f"{{{SVG_NS}}}text", attributes)
    element.text = region.text
    return element


def encoded_image_href(href: str) -> str:
    if href.startswith("data:image/png;base64,"):
        return href
    return quote(href, safe="/:%-._~")


def compact(value: float) -> str:
    return f"{value:.3f}".rstrip("0").rstrip(".")