Creates a draw.io style diagram for a data flow modelled in yaml.

diagram_generator.py

#!/usr/bin/env python3
"""Generate a .drawio diagram from a flow definition plus its registry.

Reads the Git-native registry layout described in work/Data flow design.md
(registry/components, registry/edges, registry/flows — one YAML file per
ID) and emits a draw.io (mxGraph) XML file for a single flow: one vertex
per component occurrence in the path, one edge per path step labelled with
transport/format.

Annotations render in two distinct ways depending on where they're
declared:

- Flow-level annotations (top-level `annotations` on the flow) are never
  anchored to anything — they're flow-wide notes, listed above the diagram.
- Edge-level annotations (`annotations` on a `path` entry) are always
  anchored to that specific edge occurrence — rendered as a small numbered
  badge on the edge, with the full text in a legend column down the right.

Usage:
    python generate_diagram.py <registry_dir> <flow_id> [output_path]
"""
import sys
import os
import yaml
import xml.etree.ElementTree as ET
from xml.dom import minidom

MARGIN_X = 40
LANE_Y_MIN = 140
SPACING_X = 220
VERTEX_WIDTH = 160
VERTEX_HEIGHT = 60
CIRCLE_SIZE = 26
LEGEND_GAP = 60
LEGEND_WIDTH = 260
LEGEND_ROW_HEIGHT = 70
LEGEND_ROW_GAP = 14
LEGEND_Y_START = 70

LANE_LABEL_HEIGHT = 24
LANE_PAD_X = 20
LANE_PAD_Y = 20
LANE_GAP = 20
LANE_HEIGHT = LANE_LABEL_HEIGHT + LANE_PAD_Y + VERTEX_HEIGHT + LANE_PAD_Y

LANE_PALETTE = ["#666666", "#6c8ebf", "#d79b00", "#9673a6", "#82b366", "#b85450"]

FLOW_LIST_X = MARGIN_X
FLOW_LIST_WIDTH = 900
FLOW_LIST_Y_START = 60
FLOW_ROW_HEIGHT = 26
FLOW_LIST_HEADER_HEIGHT = 14

ANNOTATION_STYLES = {
    "note": dict(fill="#f5f5f5", stroke="#666666"),
    "info": dict(fill="#dae8fc", stroke="#6c8ebf"),
    "warning": dict(fill="#fff2cc", stroke="#d6b656"),
    "callout": dict(fill="#ffe6cc", stroke="#d79b00"),
}


def load_registry(registry_dir):
    def load_dir(name):
        path = os.path.join(registry_dir, name)
        items = {}
        for fname in os.listdir(path):
            if fname.endswith((".yaml", ".yml")):
                with open(os.path.join(path, fname)) as f:
                    doc = yaml.safe_load(f)
                    items[doc["id"]] = doc
        return items

    return {
        "components": load_dir("components"),
        "edges": load_dir("edges"),
        "flows": load_dir("flows"),
    }


def cell(xml_parent, **attrs):
    return ET.SubElement(xml_parent, "mxCell", {k: str(v) for k, v in attrs.items() if v is not None})


def geometry(cell_el, **attrs):
    attrs["as"] = "geometry"
    ET.SubElement(cell_el, "mxGeometry", {k: str(v) for k, v in attrs.items()})


class IdGen:
    def __init__(self):
        self.n = 1

    def next(self, prefix):
        self.n += 1
        return f"{prefix}{self.n}"


SKETCH = "sketch=1;fontFamily=Architects Daughter;curveFitting=1;jiggle=2;"


def make_vertex(root, cell_id, label, x, y):
    c = cell(root, id=cell_id, value=label,
             style=(SKETCH + "rounded=1;whiteSpace=wrap;html=1;"
                    "fillColor=#dae8fc;strokeColor=#6c8ebf;"),
             vertex=1, parent=1)
    geometry(c, x=x, y=y, width=VERTEX_WIDTH, height=VERTEX_HEIGHT)


def make_edge(root, cell_id, source_id, target_id, label):
    c = cell(root, id=cell_id, value=label,
             style=(SKETCH + "edgeStyle=orthogonalEdgeStyle;rounded=0;html=1;fontSize=11;"),
             edge=1, parent=1, source=source_id, target=target_id)
    geometry(c, relative=1)


def make_marker(root, cell_id, number, ann_type, x, y):
    style = ANNOTATION_STYLES.get(ann_type, ANNOTATION_STYLES["note"])
    c = cell(root, id=cell_id, value=str(number),
             style=("ellipse;whiteSpace=wrap;html=1;align=center;verticalAlign=middle;"
                    f"fillColor={style['fill']};strokeColor={style['stroke']};"
                    "fontSize=12;fontStyle=1;"),
             vertex=1, parent=1)
    geometry(c, x=x, y=y, width=CIRCLE_SIZE, height=CIRCLE_SIZE)


def make_lane_box(root, cell_id, label, color, x, y, width):
    c = cell(root, id=cell_id, value=label,
             style=("rounded=0;whiteSpace=wrap;html=1;align=left;verticalAlign=top;"
                    "dashed=1;fillColor=none;fontStyle=1;fontSize=12;spacing=8;"
                    f"strokeColor={color};fontColor={color};"),
             vertex=1, parent=1)
    geometry(c, x=x, y=y, width=width, height=LANE_HEIGHT)


def make_legend_row(root, cell_id, number, text, ann_type, x, y):
    style = ANNOTATION_STYLES.get(ann_type, ANNOTATION_STYLES["note"])
    c = cell(root, id=cell_id, value=f"{number}. {text}",
             style=("rounded=0;whiteSpace=wrap;html=1;align=left;verticalAlign=top;"
                    f"spacing=8;fontSize=11;fillColor={style['fill']};strokeColor={style['stroke']};"),
             vertex=1, parent=1)
    geometry(c, x=x, y=y, width=LEGEND_WIDTH, height=LEGEND_ROW_HEIGHT)


def _escape_html(text):
    return (text.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;"))


def make_flow_notes_box(root, cell_id, annotations, x, y):
    items = []
    for ann in annotations:
        style = ANNOTATION_STYLES.get(ann.get("type", "note"), ANNOTATION_STYLES["note"])
        items.append(f'<li style="color:{style["stroke"]};">{_escape_html(ann["text"])}</li>')
    value = "<ul style=\"margin:4px 0 4px 16px;padding:0;\">" + "".join(items) + "</ul>"
    height = FLOW_LIST_HEADER_HEIGHT + len(annotations) * FLOW_ROW_HEIGHT
    c = cell(root, id=cell_id, value=value,
             style=("rounded=0;whiteSpace=wrap;html=1;align=left;verticalAlign=top;"
                    "spacing=8;fontSize=11;fillColor=#f5f5f5;strokeColor=#666666;"),
             vertex=1, parent=1)
    geometry(c, x=x, y=y, width=FLOW_LIST_WIDTH, height=height)
    return height


def generate(registry, flow_id):
    flow = registry["flows"][flow_id]
    edges = registry["edges"]
    components = registry["components"]
    path = flow["path"]

    ids = IdGen()

    mxfile = ET.Element("mxfile", host="app.diagrams.net")
    diagram = ET.SubElement(mxfile, "diagram", id=flow_id, name=flow["display_name"])
    model = ET.SubElement(diagram, "mxGraphModel",
                           dx="800", dy="600", grid="1", gridSize="10", guides="1",
                           tooltips="1", connect="1", arrows="1", fold="1", page="1",
                           pageScale="1", pageWidth="1400", pageHeight="900",
                           math="0", shadow="0")
    root = ET.SubElement(model, "root")
    cell(root, id="0")
    cell(root, id="1", parent="0")

    header_text = (f"{flow['id']} · v{flow['version']} · {flow['status']} "
                    f"· architect: {flow['architect']}")
    header = cell(root, id="header", value=header_text,
                  style="text;html=1;align=left;verticalAlign=middle;fontStyle=1;fontSize=14;",
                  vertex=1, parent=1)
    geometry(header, x=MARGIN_X, y=20, width=800, height=30)

    # Flow-level annotations: never anchored, rendered as one bulleted list
    # above the diagram
    flow_anns = flow.get("annotations", [])
    flow_list_y = FLOW_LIST_Y_START
    if flow_anns:
        box_height = make_flow_notes_box(root, ids.next("flowann"), flow_anns,
                                          FLOW_LIST_X, flow_list_y)
        flow_list_y += box_height

    lanes_y_start = max(LANE_Y_MIN, flow_list_y + 20)

    # Occurrence sequence (source, then each path step's target), used both
    # for x position (path order) and for assigning each occurrence to its
    # component's environment lane (y position) — a component can repeat
    # non-contiguously and still land in the same lane.
    first_edge = edges[path[0]["edge"]]
    occurrence_components = [components[first_edge["source"]]]
    for step in path:
        edge_def = edges[step["edge"]]
        occurrence_components.append(components[edge_def["target"]])

    lane_index = {}
    for comp in occurrence_components:
        env = comp.get("environment", "unspecified")
        if env not in lane_index:
            lane_index[env] = len(lane_index)

    last_x = MARGIN_X + (len(occurrence_components) - 1) * SPACING_X
    lane_width = last_x + VERTEX_WIDTH - MARGIN_X + 2 * LANE_PAD_X
    lane_top = {
        env: lanes_y_start + i * (LANE_HEIGHT + LANE_GAP)
        for env, i in lane_index.items()
    }
    for env, i in lane_index.items():
        color = LANE_PALETTE[i % len(LANE_PALETTE)]
        make_lane_box(root, ids.next("lane"), env, color,
                      MARGIN_X - LANE_PAD_X, lane_top[env], lane_width)

    def vertex_y(comp):
        return lane_top[comp.get("environment", "unspecified")] + LANE_LABEL_HEIGHT + LANE_PAD_Y

    # Edge-level annotations: always anchored to their path step, numbered
    # in path order, rendered as a badge plus a right-hand legend entry
    edge_annotations = []  # (step_edge_id, text, type)

    x = MARGIN_X
    prev_vertex = ids.next("v")
    prev_x, prev_y = x, vertex_y(occurrence_components[0])
    make_vertex(root, prev_vertex, occurrence_components[0]["display_name"], x, prev_y)
    x += SPACING_X

    edge_mid = {}
    for i, step in enumerate(path):
        edge_def = edges[step["edge"]]
        target_component = occurrence_components[i + 1]
        y = vertex_y(target_component)

        v_id = ids.next("v")
        make_vertex(root, v_id, target_component["display_name"], x, y)

        e_id = ids.next("e")
        label = f"{edge_def['transport']} / {edge_def['source_produces']}"
        make_edge(root, e_id, prev_vertex, v_id, label)
        mid_x = (prev_x + VERTEX_WIDTH + x) / 2
        mid_y = (prev_y + y) / 2 + VERTEX_HEIGHT / 2
        edge_mid[step["edge"]] = (mid_x, mid_y)

        for ann in step.get("annotations", []):
            edge_annotations.append((step["edge"], ann["text"], ann.get("type", "note")))

        prev_vertex, prev_x, prev_y = v_id, x, y
        x += SPACING_X

    last_vertex_x = prev_x
    lanes_bottom = lanes_y_start + len(lane_index) * (LANE_HEIGHT + LANE_GAP) - LANE_GAP
    legend_x = last_vertex_x + VERTEX_WIDTH + LEGEND_GAP
    legend_y = LEGEND_Y_START
    marker_count_by_edge = {}

    for number, (step_edge_id, text, ann_type) in enumerate(edge_annotations, start=1):
        mid_x, mid_y = edge_mid[step_edge_id]
        stack_offset = marker_count_by_edge.get(step_edge_id, 0)
        marker_count_by_edge[step_edge_id] = stack_offset + 1

        marker_x = mid_x - CIRCLE_SIZE / 2 + stack_offset * (CIRCLE_SIZE + 4)
        marker_y = mid_y - CIRCLE_SIZE - 4
        make_marker(root, ids.next("marker"), number, ann_type, marker_x, marker_y)

        make_legend_row(root, ids.next("legend"), number, text, ann_type, legend_x, legend_y)
        legend_y += LEGEND_ROW_HEIGHT + LEGEND_ROW_GAP

    model.set("pageWidth", str(int(legend_x + LEGEND_WIDTH + MARGIN_X)))
    model.set("pageHeight", str(int(max(lanes_bottom + 50, legend_y + 50))))

    return mxfile


def prettify(elem):
    rough = ET.tostring(elem, "utf-8")
    return minidom.parseString(rough).toprettyxml(indent="  ")


def main():
    if len(sys.argv) < 3:
        print("Usage: generate_diagram.py <registry_dir> <flow_id> [output_path]")
        sys.exit(1)

    registry_dir = sys.argv[1]
    flow_id = sys.argv[2]
    output_path = sys.argv[3] if len(sys.argv) > 3 else f"{flow_id}.drawio"

    registry = load_registry(registry_dir)
    mxfile = generate(registry, flow_id)

    with open(output_path, "w") as f:
        f.write(prettify(mxfile))

    print(f"Wrote {output_path}")


if __name__ == "__main__":
    main()

use_case_diagram_generator.py

#!/usr/bin/env python3
"""Generate one merged .drawio diagram for a use case (multiple flows).

This is "approach 2" from the use-case diagramming discussion: instead of
rendering each flow separately, union every flow's components and edges
into a single graph. A component shared by more than one flow (same
registry ID) is drawn once. An edge shared by more than one flow is drawn
as one parallel line per flow, fanned out and colour-coded so they run
alongside each other rather than overlapping, with a flow legend on the
right.

Components still group into horizontal environment lanes, same idea as
generate_diagram.py, but x position now comes from a topological rank in
the merged graph rather than path order — a use case's flows can branch
and reconverge, so there's no single linear path to lay out left to right.

Known gaps (prototype scope, see README):
- Per-edge and flow-level annotations aren't rendered here yet.
- If a single flow revisits a component (e.g. router -> enrichment ->
  router), the merged graph has a genuine cycle once components are
  deduped by ID. Rank computation tolerates this by ignoring back edges,
  but the resulting left-to-right order for components inside a cycle is
  a heuristic, not a meaningful "depth".

Usage:
    python generate_use_case_diagram.py <registry_dir> <use_case_id> [output_path]
"""
import sys
import os
import yaml
import xml.etree.ElementTree as ET
from xml.dom import minidom
from collections import defaultdict

MARGIN_X = 40
MARGIN_Y = 20
SPACING_X = 220
VERTEX_WIDTH = 160
VERTEX_HEIGHT = 60
ROW_GAP = 20

LANE_LABEL_HEIGHT = 24
LANE_PAD_X = 20
LANE_PAD_Y = 20
LANE_GAP = 20

LEGEND_GAP = 60
LEGEND_WIDTH = 240
LEGEND_ROW_HEIGHT = 60
LEGEND_ROW_GAP = 12
LEGEND_Y_START = 70

SKETCH = "sketch=1;fontFamily=Architects Daughter;curveFitting=1;jiggle=2;"

LANE_PALETTE = ["#666666", "#6c8ebf", "#d79b00", "#9673a6", "#82b366", "#b85450"]
FLOW_PALETTE = ["#6c8ebf", "#d79b00", "#82b366", "#9673a6", "#b85450", "#666666"]

# When an edge is shared by multiple flows, each flow's line is drawn
# separately and fanned out vertically across this fraction of the source/
# target box height, instead of collapsing into one overlapping line.
EDGE_FAN_SPAN = 0.3


def load_registry(registry_dir):
    def load_dir(name):
        path = os.path.join(registry_dir, name)
        items = {}
        for fname in os.listdir(path):
            if fname.endswith((".yaml", ".yml")):
                with open(os.path.join(path, fname)) as f:
                    doc = yaml.safe_load(f)
                    items[doc["id"]] = doc
        return items

    return {
        "components": load_dir("components"),
        "edges": load_dir("edges"),
        "flows": load_dir("flows"),
        "use_cases": load_dir("use-cases"),
    }


def cell(xml_parent, **attrs):
    return ET.SubElement(xml_parent, "mxCell", {k: str(v) for k, v in attrs.items() if v is not None})


def geometry(cell_el, **attrs):
    attrs["as"] = "geometry"
    ET.SubElement(cell_el, "mxGeometry", {k: str(v) for k, v in attrs.items()})


class IdGen:
    def __init__(self):
        self.n = 1

    def next(self, prefix):
        self.n += 1
        return f"{prefix}{self.n}"


def make_vertex(root, cell_id, label, x, y):
    c = cell(root, id=cell_id, value=label,
             style=(SKETCH + "rounded=1;whiteSpace=wrap;html=1;"
                    "fillColor=#dae8fc;strokeColor=#6c8ebf;"),
             vertex=1, parent=1)
    geometry(c, x=x, y=y, width=VERTEX_WIDTH, height=VERTEX_HEIGHT)


def make_edge(root, cell_id, source_id, target_id, label, color, width,
              exit_x=None, exit_y=None, entry_x=None, entry_y=None):
    style = (SKETCH + "edgeStyle=orthogonalEdgeStyle;rounded=0;html=1;fontSize=11;"
             f"strokeColor={color};strokeWidth={width};")
    if exit_x is not None:
        style += f"exitX={exit_x};exitY={exit_y};exitDx=0;exitDy=0;"
    if entry_x is not None:
        style += f"entryX={entry_x};entryY={entry_y};entryDx=0;entryDy=0;"
    c = cell(root, id=cell_id, value=label, style=style,
             edge=1, parent=1, source=source_id, target=target_id)
    geometry(c, relative=1)


def make_lane_box(root, cell_id, label, color, x, y, width, height):
    c = cell(root, id=cell_id, value=label,
             style=("rounded=0;whiteSpace=wrap;html=1;align=left;verticalAlign=top;"
                    "dashed=1;fillColor=none;fontStyle=1;fontSize=12;spacing=8;"
                    f"strokeColor={color};fontColor={color};"),
             vertex=1, parent=1)
    geometry(c, x=x, y=y, width=width, height=height)


def make_legend_row(root, cell_id, color, label, x, y, width, height):
    c = cell(root, id=cell_id, value=label,
             style=("rounded=0;whiteSpace=wrap;html=1;align=left;verticalAlign=middle;"
                    f"spacing=8;fontSize=11;fillColor=#f5f5f5;strokeColor={color};strokeWidth=3;"),
             vertex=1, parent=1)
    geometry(c, x=x, y=y, width=width, height=height)


def build_merged_graph(registry, use_case):
    flows = registry["flows"]
    edges = registry["edges"]

    flow_ids = [f["flow_id"] for f in use_case["flows"]]

    edge_flows = defaultdict(set)  # edge_id -> set of flow_ids that traverse it
    component_ids = set()
    incoming = defaultdict(list)  # target -> [source, ...]
    for flow_id in flow_ids:
        for step in flows[flow_id]["path"]:
            edge_id = step["edge"]
            edge_flows[edge_id].add(flow_id)
            edge_def = edges[edge_id]
            component_ids.add(edge_def["source"])
            component_ids.add(edge_def["target"])
            incoming[edge_def["target"]].append(edge_def["source"])

    # Longest-path rank from each component's predecessors. A flow that
    # revisits a component (router -> enrichment -> router) creates a real
    # cycle once components are deduped by ID -- back edges (predecessors
    # still mid-computation) are ignored rather than recursing forever, so
    # this always terminates, but rank inside a cycle is a heuristic.
    rank = {}
    state = {}

    def compute_rank(comp_id):
        if comp_id in rank:
            return rank[comp_id]
        state[comp_id] = "visiting"
        best = 0
        for pred in incoming.get(comp_id, []):
            if state.get(pred) == "visiting":
                continue
            best = max(best, 1 + compute_rank(pred))
        rank[comp_id] = best
        state[comp_id] = "done"
        return best

    for comp_id in component_ids:
        compute_rank(comp_id)

    return flow_ids, component_ids, edge_flows, rank


def generate(registry, use_case_id):
    use_case = registry["use_cases"][use_case_id]
    components = registry["components"]
    edges = registry["edges"]

    flow_ids, component_ids, edge_flows, rank = build_merged_graph(registry, use_case)
    flow_color = {fid: FLOW_PALETTE[i % len(FLOW_PALETTE)] for i, fid in enumerate(flow_ids)}

    ids = IdGen()

    mxfile = ET.Element("mxfile", host="app.diagrams.net")
    diagram = ET.SubElement(mxfile, "diagram", id=use_case_id, name=use_case["display_name"])
    model = ET.SubElement(diagram, "mxGraphModel",
                           dx="800", dy="600", grid="1", gridSize="10", guides="1",
                           tooltips="1", connect="1", arrows="1", fold="1", page="1",
                           pageScale="1", pageWidth="1400", pageHeight="900",
                           math="0", shadow="0")
    root = ET.SubElement(model, "root")
    cell(root, id="0")
    cell(root, id="1", parent="0")

    header_text = (f"{use_case['id']} · v{use_case['version']} · {use_case['status']} "
                   f"· architect: {use_case['architect']}")
    header = cell(root, id="header", value=header_text,
                  style="text;html=1;align=left;verticalAlign=middle;fontStyle=1;fontSize=14;",
                  vertex=1, parent=1)
    geometry(header, x=MARGIN_X, y=MARGIN_Y, width=900, height=30)

    lanes_y_start = MARGIN_Y + 50

    # Order components by rank (then ID, for determinism) to decide lane
    # appearance order and to assign sub-row slots within a lane for
    # components that land on the same rank in the same environment.
    ordered_components = sorted(component_ids, key=lambda c: (rank[c], c))

    lane_order = []
    for comp_id in ordered_components:
        env = components[comp_id].get("environment", "unspecified")
        if env not in lane_order:
            lane_order.append(env)

    rank_slot_count = defaultdict(int)
    component_subrow = {}
    for comp_id in ordered_components:
        env = components[comp_id].get("environment", "unspecified")
        key = (env, rank[comp_id])
        component_subrow[comp_id] = rank_slot_count[key]
        rank_slot_count[key] += 1

    lane_subrows = defaultdict(lambda: 1)
    for (env, _r), count in rank_slot_count.items():
        lane_subrows[env] = max(lane_subrows[env], count)

    max_rank = max(rank.values())
    lane_width = max_rank * SPACING_X + VERTEX_WIDTH + 2 * LANE_PAD_X

    lane_top = {}
    y_cursor = lanes_y_start
    for i, env in enumerate(lane_order):
        subrows = lane_subrows[env]
        height = (LANE_LABEL_HEIGHT + LANE_PAD_Y
                  + subrows * VERTEX_HEIGHT + (subrows - 1) * ROW_GAP
                  + LANE_PAD_Y)
        lane_top[env] = y_cursor
        color = LANE_PALETTE[i % len(LANE_PALETTE)]
        make_lane_box(root, ids.next("lane"), env, color,
                      MARGIN_X - LANE_PAD_X, y_cursor, lane_width, height)
        y_cursor += height + LANE_GAP

    vertex_id = {}
    for comp_id in ordered_components:
        env = components[comp_id].get("environment", "unspecified")
        x = MARGIN_X + rank[comp_id] * SPACING_X
        y = (lane_top[env] + LANE_LABEL_HEIGHT + LANE_PAD_Y
             + component_subrow[comp_id] * (VERTEX_HEIGHT + ROW_GAP))
        v_id = ids.next("v")
        vertex_id[comp_id] = v_id
        make_vertex(root, v_id, components[comp_id]["display_name"], x, y)

    lanes_bottom = y_cursor - LANE_GAP
    legend_x = MARGIN_X + lane_width + LEGEND_GAP

    for edge_id, used_by in edge_flows.items():
        edge_def = edges[edge_id]
        label = f"{edge_def['transport']} / {edge_def['source_produces']}"
        source_id = edge_def["source"]
        target_id = edge_def["target"]

        # Side to exit/enter on: forward edges run left-to-right (rank
        # increases), back edges (target rank no higher than source) run
        # right-to-left, so their fan-out sits on the opposite sides.
        forward = rank[target_id] >= rank[source_id]
        exit_x, entry_x = (1, 0) if forward else (0, 1)

        used_in_order = [fid for fid in flow_ids if fid in used_by]
        n = len(used_in_order)
        for i, flow_id in enumerate(used_in_order):
            if n == 1:
                y_frac = 0.5
            else:
                y_frac = 0.5 - EDGE_FAN_SPAN / 2 + i * (EDGE_FAN_SPAN / (n - 1))
            edge_label = label if i == 0 else None
            make_edge(root, ids.next("e"), vertex_id[source_id], vertex_id[target_id],
                      edge_label, flow_color[flow_id], 1,
                      exit_x=exit_x, exit_y=y_frac, entry_x=entry_x, entry_y=y_frac)

    legend_y = LEGEND_Y_START
    flow_entry_by_id = {f["flow_id"]: f for f in use_case["flows"]}
    for flow_id in flow_ids:
        entry = flow_entry_by_id[flow_id]
        label = f"{entry.get('description', flow_id)} ({entry['status']})"
        make_legend_row(root, ids.next("legend"), flow_color[flow_id], label,
                         legend_x, legend_y, LEGEND_WIDTH, LEGEND_ROW_HEIGHT)
        legend_y += LEGEND_ROW_HEIGHT + LEGEND_ROW_GAP

    model.set("pageWidth", str(int(legend_x + LEGEND_WIDTH + MARGIN_X)))
    model.set("pageHeight", str(int(max(lanes_bottom + 50, legend_y + 50))))

    return mxfile


def prettify(elem):
    rough = ET.tostring(elem, "utf-8")
    return minidom.parseString(rough).toprettyxml(indent="  ")


def main():
    if len(sys.argv) < 3:
        print("Usage: generate_use_case_diagram.py <registry_dir> <use_case_id> [output_path]")
        sys.exit(1)

    registry_dir = sys.argv[1]
    use_case_id = sys.argv[2]
    output_path = sys.argv[3] if len(sys.argv) > 3 else f"{use_case_id}.drawio"

    registry = load_registry(registry_dir)
    mxfile = generate(registry, use_case_id)

    with open(output_path, "w") as f:
        f.write(prettify(mxfile))

    print(f"Wrote {output_path}")


if __name__ == "__main__":
    main()