# kernbench/topology/visualizer.py
"""
SVG diagram generator for TopologyGraph views.

Produces mm-accurate, deterministic SVG files for each view level
(system, SIP, cube, PE) per ADR-0005 and ADR-0006.
"""
from __future__ import annotations

from pathlib import Path

from .types import Edge, Node, TopologyGraph, ViewGraph

# ── Color palette by component kind ─────────────────────────────────

_KIND_COLORS: dict[str, str] = {
    "switch":       "#6366f1",  # indigo
    "sip":          "#e0e7ff",  # light indigo
    "iochiplet":    "#0ea5e9",  # sky blue
    "pcie_ep":      "#0ea5e9",
    "io_cpu":       "#0ea5e9",
    "ucie_port":    "#3b82f6",  # blue
    "noc":          "#a78bfa",  # purple
    "m_cpu":        "#f59e0b",  # amber
    "noc_router":   "#f97316",  # orange
    "hbm_ctrl":     "#10b981",  # emerald
    "pe":           "#94a3b8",  # slate
    "pe_cpu":       "#ef4444",  # red
    "pe_scheduler": "#f59e0b",  # amber
    "pe_dma":       "#3b82f6",  # blue
    "pe_gemm":      "#8b5cf6",  # violet
    "pe_math":      "#ec4899",  # pink
    "pe_tcm":       "#10b981",  # emerald
    "sram":         "#f59e0b",  # amber
    "cube":         "#cbd5e1",  # slate-300
}

_EDGE_COLORS: dict[str, str] = {
    "pcie":         "#6366f1",
    "io_internal":  "#0ea5e9",
    "io_to_cube":   "#0ea5e9",
    "ucie_mesh":    "#3b82f6",
    "pe_to_router":   "#f97316",
    "router_to_hbm":  "#10b981",
    "hbm_to_router":  "#10b981",
    "router_mesh":    "#a78bfa",
    "router_to_sram": "#a78bfa",
    "noc_to_ucie":  "#a78bfa",
    "pe_to_noc":    "#a78bfa",
    "noc_to_sram":  "#f59e0b",
    "command":      "#f59e0b",
    "pe_internal":  "#94a3b8",
}

# ── Node sizing ──────────────────────────────────────────────────────

_DEFAULT_NODE_W = 2.0  # mm
_DEFAULT_NODE_H = 1.2  # mm

_KIND_SIZE: dict[str, tuple[float, float]] = {
    "sip":  (60.0, 50.0),
    "cube": (6.0, 4.0),
    "iochiplet": (4.0, 1.5),
    "switch": (5.0, 1.5),
    "noc_router": (1.0, 0.7),
    "ucie_port": (1.2, 0.7),
    "ucie_conn": (0.8, 0.5),
    "sram": (1.4, 0.7),
    "m_cpu": (1.4, 0.7),
    "hbm_ctrl": (1.8, 0.8),
}


# ── Public API ───────────────────────────────────────────────────────


def emit_diagrams(graph: TopologyGraph, out_dir: Path) -> list[Path]:
    """Generate SVG diagrams for all views. Returns list of created file paths."""
    out_dir.mkdir(parents=True, exist_ok=True)
    created: list[Path] = []

    views = [
        ("system_view", graph.system_view),
        ("sip_view", graph.sip_view),
        ("cube_view", graph.cube_view),
        ("pe_view", graph.pe_view),
    ]

    for name, view in views:
        if view is None:
            continue
        if name == "cube_view":
            svg = _render_cube_view_svg(view, graph.spec)
        else:
            svg = _render_view_svg(view)
        path = out_dir / f"{name}.svg"
        path.write_text(svg, encoding="utf-8")
        created.append(path)

    return created


# ── SVG rendering ────────────────────────────────────────────────────


def _render_view_svg(view: ViewGraph) -> str:
    """Render a ViewGraph to an SVG string."""
    scale = _pick_scale(view)
    pad = 40  # px padding
    node_sizes = _compute_node_sizes(view, scale)

    # Canvas size in px
    w_px = int(view.width_mm * scale + 2 * pad)
    h_px = int(view.height_mm * scale + 2 * pad)

    parts: list[str] = []
    parts.append(_svg_header(w_px, h_px, view.name))

    # Background
    parts.append(f'  <rect width="{w_px}" height="{h_px}" fill="#f8fafc"/>')

    # Title
    parts.append(
        f'  <text x="{w_px // 2}" y="18" text-anchor="middle" '
        f'font-family="monospace" font-size="14" font-weight="bold" fill="#1e293b">'
        f'{view.name.upper()} VIEW</text>'
    )

    # Special: draw cube boundary + HBM block background in cube view
    if view.name == "cube":
        _draw_cube_boundary(parts, view, scale, pad)
        _draw_hbm_block(parts, view, scale, pad)

    # Edges (draw before nodes so nodes are on top)
    # Track fan-out edges to assign per-edge offsets
    fanout_counter: dict[str, int] = {}
    for edge in view.edges:
        if edge.src in view.nodes and edge.dst in view.nodes:
            _draw_edge(parts, edge, view, node_sizes, scale, pad, fanout_counter)

    # Nodes
    for node in view.nodes.values():
        _draw_node(parts, node, node_sizes, scale, pad)

    parts.append("</svg>")
    return "\n".join(parts)


def _pick_scale(view: ViewGraph) -> float:
    """Pixels per mm, chosen per view type."""
    return {
        "system": 4.0,
        "sip": 8.0,
        "cube": 28.0,
        "pe": 35.0,
    }.get(view.name, 10.0)


def _compute_node_sizes(
    view: ViewGraph, scale: float,
) -> dict[str, tuple[float, float]]:
    """Returns (w_px, h_px) for each node."""
    sizes: dict[str, tuple[float, float]] = {}
    for nid, node in view.nodes.items():
        w_mm, h_mm = _KIND_SIZE.get(node.kind, (_DEFAULT_NODE_W, _DEFAULT_NODE_H))
        # For cube view, use smaller PE nodes
        if view.name == "cube" and node.kind == "pe":
            w_mm, h_mm = 1.4, 0.7
        if view.name == "pe":
            w_mm, h_mm = 2.5, 1.4
        sizes[nid] = (w_mm * scale, h_mm * scale)
    return sizes


def _svg_header(w: int, h: int, title: str) -> str:
    return (
        f'<svg xmlns="http://www.w3.org/2000/svg" '
        f'width="{w}" height="{h}" viewBox="0 0 {w} {h}">\n'
        f'  <title>{title}</title>'
    )


def _draw_cube_boundary(
    parts: list[str], view: ViewGraph, scale: float, pad: int,
) -> None:
    """Draw the cube die outline as a dashed rectangle."""
    bx = pad
    by = pad
    bw = view.width_mm * scale
    bh = view.height_mm * scale
    parts.append(
        f'  <rect x="{bx:.1f}" y="{by:.1f}" '
        f'width="{bw:.1f}" height="{bh:.1f}" '
        f'rx="6" fill="none" stroke="#475569" stroke-width="2" '
        f'stroke-dasharray="8,4"/>'
    )


def _draw_hbm_block(
    parts: list[str], view: ViewGraph, scale: float, pad: int,
) -> None:
    """Draw HBM area as a filled rectangle in cube view."""
    # HBM area: centered at (8.5, 7.0), size 9x5 -> x=[4.0,13.0], y=[4.5,9.5]
    hbm_x = 4.0 * scale + pad
    hbm_y = 4.5 * scale + pad
    hbm_w = 9.0 * scale
    hbm_h = 5.0 * scale
    parts.append(
        f'  <rect x="{hbm_x:.1f}" y="{hbm_y:.1f}" '
        f'width="{hbm_w:.1f}" height="{hbm_h:.1f}" '
        f'rx="4" fill="#d1fae5" stroke="#10b981" stroke-width="1.5" '
        f'stroke-dasharray="6,3" opacity="0.5"/>'
    )
    cx = 8.5 * scale + pad
    cy = 8.5 * scale + pad
    parts.append(
        f'  <text x="{cx:.1f}" y="{cy:.1f}" text-anchor="middle" '
        f'font-family="monospace" font-size="11" fill="#047857" opacity="0.7">'
        f'HBM</text>'
    )


def _draw_node(
    parts: list[str],
    node: Node,
    sizes: dict[str, tuple[float, float]],
    scale: float,
    pad: int,
) -> None:
    """Draw a single node as a rounded rectangle with label."""
    if node.pos_mm is None:
        return
    px = node.pos_mm[0] * scale + pad
    py = node.pos_mm[1] * scale + pad
    w, h = sizes.get(node.id, (40, 24))

    x = px - w / 2
    y = py - h / 2
    fill = _KIND_COLORS.get(node.kind, "#e2e8f0")
    text_color = "#ffffff" if _is_dark(fill) else "#1e293b"

    parts.append(
        f'  <rect x="{x:.1f}" y="{y:.1f}" width="{w:.1f}" height="{h:.1f}" '
        f'rx="4" fill="{fill}" stroke="#475569" stroke-width="1"/>'
    )

    label = node.label or node.id
    font_size = _label_font_size(w, label)
    parts.append(
        f'  <text x="{px:.1f}" y="{py + 4:.1f}" text-anchor="middle" '
        f'font-family="monospace" font-size="{font_size}" fill="{text_color}">'
        f'{_escape(label)}</text>'
    )


# ── Fan-out edge kinds that need offset routing ─────────────────────

_FANOUT_KINDS = {"pe_to_router", "command", "router_to_ucie_conn", "ucie_conn_to_router"}


def _draw_edge(
    parts: list[str],
    edge: Edge,
    view: ViewGraph,
    sizes: dict[str, tuple[float, float]],
    scale: float,
    pad: int,
    fanout_counter: dict[str, int],
) -> None:
    """Draw an edge with orthogonal (90-degree) routing for fan-out kinds."""
    nodes = view.nodes
    src_node = nodes[edge.src]
    dst_node = nodes[edge.dst]
    if src_node.pos_mm is None or dst_node.pos_mm is None:
        return

    x1 = src_node.pos_mm[0] * scale + pad
    y1 = src_node.pos_mm[1] * scale + pad
    x2 = dst_node.pos_mm[0] * scale + pad
    y2 = dst_node.pos_mm[1] * scale + pad

    color = _EDGE_COLORS.get(edge.kind, "#94a3b8")
    width = "1.5" if edge.kind == "pe_internal" else "1"
    opacity = "0.6" if edge.kind in ("command", "noc_to_ucie") else "0.8"
    # HBM links: thin and faint to reduce clutter
    if edge.kind in ("router_to_hbm", "hbm_to_router"):
        width = "0.5"
        opacity = "0.3"
    # Router mesh links: thin
    if edge.kind == "router_mesh":
        width = "0.5"
        opacity = "0.4"

    if edge.kind in _FANOUT_KINDS and view.name == "cube":
        # Orthogonal routing: src→horizontal→vertical→dst with per-edge offset.
        group_key = f"{edge.kind}:{edge.dst}"
        idx = fanout_counter.get(group_key, 0)
        fanout_counter[group_key] = idx + 1

        # Route: go vertically from src to a staggered horizontal channel,
        # then horizontally to dst x, then vertically to dst.
        mid_y = (y1 + y2) / 2 + (idx - 1.5) * 10  # spread channels vertically

        parts.append(
            f'  <polyline points="{x1:.1f},{y1:.1f} {x1:.1f},{mid_y:.1f} '
            f'{x2:.1f},{mid_y:.1f} {x2:.1f},{y2:.1f}" '
            f'fill="none" stroke="{color}" stroke-width="{width}" opacity="{opacity}"/>'
        )

        # Label on the horizontal segment
        if edge.distance_mm > 0:
            lx = (x1 + x2) / 2
            label = f"{edge.distance_mm:.1f}mm"
            if edge.bw_gbs:
                label += f" {edge.bw_gbs:.0f}GB/s"
            parts.append(
                f'  <text x="{lx:.1f}" y="{mid_y - 3:.1f}" text-anchor="middle" '
                f'font-family="monospace" font-size="7" fill="#64748b">'
                f'{label}</text>'
            )
        return

    # Non-fanout: orthogonal L-bend
    if abs(x2 - x1) > 1 and abs(y2 - y1) > 1:
        # PE view: vertical-first for left→right edges (scheduler→engines),
        # horizontal-first for right→right edges (engines→tcm)
        if view.name == "pe":
            if src_node.pos_mm[0] < view.width_mm / 2:
                # Source in left half: vertical-first (scheduler fan-out)
                parts.append(
                    f'  <polyline points="{x1:.1f},{y1:.1f} {x1:.1f},{y2:.1f} {x2:.1f},{y2:.1f}" '
                    f'fill="none" stroke="{color}" stroke-width="{width}" opacity="{opacity}"/>'
                )
            else:
                # Source in right half: horizontal-first (dma/math→tcm)
                parts.append(
                    f'  <polyline points="{x1:.1f},{y1:.1f} {x2:.1f},{y1:.1f} {x2:.1f},{y2:.1f}" '
                    f'fill="none" stroke="{color}" stroke-width="{width}" opacity="{opacity}"/>'
                )
        else:
            parts.append(
                f'  <polyline points="{x1:.1f},{y1:.1f} {x2:.1f},{y1:.1f} {x2:.1f},{y2:.1f}" '
                f'fill="none" stroke="{color}" stroke-width="{width}" opacity="{opacity}"/>'
            )
    else:
        parts.append(
            f'  <line x1="{x1:.1f}" y1="{y1:.1f}" x2="{x2:.1f}" y2="{y2:.1f}" '
            f'stroke="{color}" stroke-width="{width}" opacity="{opacity}"/>'
        )

    # Distance label at midpoint
    if edge.distance_mm > 0:
        mx = (x1 + x2) / 2
        my = (y1 + y2) / 2
        label = f"{edge.distance_mm:.1f}mm"
        if edge.bw_gbs:
            label += f" {edge.bw_gbs:.0f}GB/s"
        parts.append(
            f'  <text x="{mx:.1f}" y="{my - 4:.1f}" text-anchor="middle" '
            f'font-family="monospace" font-size="7" fill="#64748b">'
            f'{label}</text>'
        )


# ── Helpers ──────────────────────────────────────────────────────────


def _is_dark(hex_color: str) -> bool:
    """Check if a hex color is dark (for white text)."""
    h = hex_color.lstrip("#")
    r, g, b = int(h[0:2], 16), int(h[2:4], 16), int(h[4:6], 16)
    return (r * 0.299 + g * 0.587 + b * 0.114) < 140


def _label_font_size(box_width: float, label: str) -> int:
    """Choose font size to fit label in box."""
    char_w = len(label) * 7
    if char_w > box_width * 0.9:
        return max(7, int(box_width * 0.9 / len(label) * 1.4))
    return 10


def _escape(text: str) -> str:
    """Escape XML special characters."""
    return text.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")


# ── Connector helper ─────────────────────────────────────────────────


def _connector_points(
    rx: float, ry: float, cx: float, cy: float
) -> str:
    """Return SVG polyline points for a rule-based connector.

    Horizontal-dominant (|dx| >= |dy|): 45° → horizontal straight → 45°.
    Vertical-dominant (|dy| > |dx|):    45° → vertical straight → 45°.
    Near-equal or tiny distance:        single straight line.
    """
    dx = cx - rx
    dy = cy - ry
    adx, ady = abs(dx), abs(dy)

    # Trivial distance → single line
    # Near-45° diagonal for short distances only (e.g. PE↔router)
    if adx + ady < 4 or (abs(adx - ady) < 4 and adx + ady < 80):
        return f"{rx:.0f},{ry:.0f} {cx:.0f},{cy:.0f}"

    sx = 1 if dx >= 0 else -1
    sy = 1 if dy >= 0 else -1

    if adx >= ady:
        # Horizontal-dominant: stubs handle vertical, straight is horizontal
        stub = ady / 2
        if stub < 2:
            return f"{rx:.0f},{ry:.0f} {cx:.0f},{cy:.0f}"
        r45x = rx + sx * stub
        r45y = ry + sy * stub
        c45x = cx - sx * stub
        c45y = cy - sy * stub  # r45y == c45y (horizontal)
    else:
        # Vertical-dominant: stubs handle horizontal, straight is vertical
        stub = adx / 2
        if stub < 2:
            return f"{rx:.0f},{ry:.0f} {cx:.0f},{cy:.0f}"
        r45x = rx + sx * stub
        r45y = ry + sy * stub
        c45x = cx - sx * stub
        c45y = cy - sy * stub  # r45x == c45x (vertical)

    return (
        f"{rx:.0f},{ry:.0f} {r45x:.0f},{r45y:.0f} "
        f"{c45x:.0f},{c45y:.0f} {cx:.0f},{cy:.0f}"
    )


# ── Cube-specific renderer ──────────────────────────────────────────


def _render_cube_view_svg(view: ViewGraph, spec: dict) -> str:
    """Render cube view with topology validation detail.

    Shows: 6×6 router grid, PE attachments, HBM pseudo channel ports,
    M_CPU/SRAM positions, UCIe connections, BW annotations.
    """
    mesh_data = spec.get("_mesh", {})
    routers = mesh_data.get("routers", {})
    n_rows = mesh_data.get("mesh", {}).get("rows", 6)
    n_cols = mesh_data.get("mesh", {}).get("cols", 6)
    cube = spec.get("cube", {})
    mm = cube.get("memory_map", {})
    clinks = cube.get("links", {})
    cube_w = cube.get("geometry", {}).get("cube_mm", {}).get("w", 17.0)
    cube_h = cube.get("geometry", {}).get("cube_mm", {}).get("h", 14.0)

    channels_per_pe = mm.get("hbm_channels_per_pe", 8)
    channel_bw = mm.get("hbm_channel_bw_gbs", 32.0)
    total_ch = mm.get("hbm_pseudo_channels", 64)
    mode = mm.get("hbm_mapping_mode", "n_to_one")
    agg_bw = channels_per_pe * channel_bw

    scale = 50  # px per mm
    pad = 60
    w_px = int(cube_w * scale + 2 * pad)
    h_px = int(cube_h * scale + 2 * pad + 80)  # extra for legend

    parts: list[str] = []
    parts.append(_svg_header(w_px, h_px, "cube"))

    # Background
    parts.append(f'  <rect width="{w_px}" height="{h_px}" fill="#0f172a"/>')

    # Title
    parts.append(
        f'  <text x="{w_px // 2}" y="22" text-anchor="middle" '
        f'font-family="monospace" font-size="14" font-weight="bold" fill="#94a3b8">'
        f'CUBE TOPOLOGY — {cube_w}×{cube_h}mm | {n_rows}×{n_cols} Router Mesh | '
        f'{mode} mode | {total_ch} pseudo-ch</text>'
    )

    # Subtitle
    parts.append(
        f'  <text x="{w_px // 2}" y="40" text-anchor="middle" '
        f'font-family="monospace" font-size="10" fill="#64748b">'
        f'Per-PE: {channels_per_pe} ch × {channel_bw} GB/s = {agg_bw} GB/s | '
        f'Cube total: {total_ch} × {channel_bw} = {total_ch * channel_bw} GB/s</text>'
    )

    # Cube boundary
    bx, by = pad, pad
    parts.append(
        f'  <rect x="{bx}" y="{by}" width="{cube_w * scale}" height="{cube_h * scale}" '
        f'rx="6" fill="none" stroke="#475569" stroke-width="2" stroke-dasharray="8,4"/>'
    )

    def mm2px(x_mm: float, y_mm: float) -> tuple[float, float]:
        return pad + x_mm * scale, pad + y_mm * scale

    # ── HBM zone background (centered, 9×5mm) ──
    hbm_x, hbm_y = mm2px(4.0, 4.5)
    hbm_w, hbm_h = 9.0 * scale, 5.0 * scale
    parts.append(
        f'  <rect x="{hbm_x:.0f}" y="{hbm_y:.0f}" '
        f'width="{hbm_w:.0f}" height="{hbm_h:.0f}" '
        f'rx="6" fill="#052e16" stroke="#047857" stroke-width="2" opacity="0.6"/>'
    )
    # HBM label
    hcx, hcy = mm2px(8.5, 7.0)
    parts.append(
        f'  <text x="{hcx:.0f}" y="{hcy - 15:.0f}" text-anchor="middle" '
        f'font-family="monospace" font-size="11" font-weight="bold" fill="#047857">'
        f'HBM_CTRL | {total_ch} pseudo channels</text>'
    )
    parts.append(
        f'  <text x="{hcx:.0f}" y="{hcy + 2:.0f}" text-anchor="middle" '
        f'font-family="monospace" font-size="9" fill="#05966988">'
        f'Total BW: {total_ch * channel_bw:.0f} GB/s</text>'
    )

    # ── Pseudo channel ports on HBM top/bottom edges ──
    # Top edge: 32 ports (PE0..PE3, 8 each), Bottom edge: 32 ports (PE4..PE7)
    half_ch = total_ch // 2
    pes_per_half = half_ch // channels_per_pe  # 4 PEs per half
    port_bar_w = hbm_w - 20  # slightly narrower than HBM zone
    port_w = port_bar_w / half_ch
    port_h = 8
    pe_colors = ["#3b82f6", "#60a5fa", "#8b5cf6", "#a78bfa",
                 "#f59e0b", "#fbbf24", "#ef4444", "#f87171"]

    for half_idx, (edge_y, pe_start) in enumerate([
        (hbm_y + 4, 0),                          # top edge, PE0-PE3
        (hbm_y + hbm_h - port_h - 4, pes_per_half),  # bottom edge, PE4-PE7
    ]):
        bar_x = hbm_x + 10
        for i in range(half_ch):
            pe_owner = pe_start + i // channels_per_pe
            c = pe_colors[pe_owner % len(pe_colors)]
            px = bar_x + i * port_w
            parts.append(
                f'  <rect x="{px:.1f}" y="{edge_y:.0f}" '
                f'width="{max(port_w - 0.5, 1):.1f}" height="{port_h}" '
                f'rx="1" fill="{c}" opacity="0.8"/>'
            )
        # Per-PE group labels
        for p in range(pes_per_half):
            gx = bar_x + (p * channels_per_pe + channels_per_pe / 2) * port_w
            label_y = edge_y - 3 if half_idx == 0 else edge_y + port_h + 8
            parts.append(
                f'  <text x="{gx:.0f}" y="{label_y:.0f}" text-anchor="middle" '
                f'font-family="monospace" font-size="6" fill="{pe_colors[(pe_start + p) % len(pe_colors)]}">'
                f'PE{pe_start + p}×{channels_per_pe}ch</text>'
            )

    # Store port group centers for PE→HBM connection lines (used later)
    _pe_hbm_targets: dict[int, tuple[float, float]] = {}
    for half_idx, (edge_y, pe_start) in enumerate([
        (hbm_y + 4, 0),
        (hbm_y + hbm_h - port_h - 4, pes_per_half),
    ]):
        bar_x = hbm_x + 10
        for p in range(pes_per_half):
            pe_id = pe_start + p
            gx = bar_x + (p * channels_per_pe + channels_per_pe / 2) * port_w
            gy = edge_y if half_idx == 0 else edge_y + port_h
            _pe_hbm_targets[pe_id] = (gx, gy)

    # ── Router mesh links ──
    for r in range(n_rows):
        for c in range(n_cols):
            rkey = f"r{r}c{c}"
            if routers.get(rkey) is None:
                continue
            rx, ry = routers[rkey]["pos_mm"]
            sx, sy = mm2px(rx, ry)

            # Horizontal neighbor
            for nc in range(c + 1, n_cols):
                nkey = f"r{r}c{nc}"
                if routers.get(nkey) is None:
                    continue
                nx, ny = routers[nkey]["pos_mm"]
                dx, dy = mm2px(nx, ny)
                parts.append(
                    f'  <line x1="{sx:.0f}" y1="{sy:.0f}" '
                    f'x2="{dx:.0f}" y2="{dy:.0f}" '
                    f'stroke="#475569" stroke-width="1" opacity="0.4"/>'
                )
                break

            # Vertical neighbor
            for nr in range(r + 1, n_rows):
                nkey = f"r{nr}c{c}"
                if routers.get(nkey) is None:
                    continue
                nx, ny = routers[nkey]["pos_mm"]
                dx, dy = mm2px(nx, ny)
                parts.append(
                    f'  <line x1="{sx:.0f}" y1="{sy:.0f}" '
                    f'x2="{dx:.0f}" y2="{dy:.0f}" '
                    f'stroke="#475569" stroke-width="1" opacity="0.4"/>'
                )
                break

    # ── Router nodes + attached component blocks ──
    r_size = 8  # px radius for router circle
    blk_w, blk_h = 32, 16  # px for component blocks

    # Component style definitions
    _COMP_STYLE = {
        "pe":   {"fill": "#2d1f3d", "stroke": "#a855f7", "text": "#a855f7"},
        "mcpu": {"fill": "#451a03", "stroke": "#f59e0b", "text": "#f59e0b"},
        "sram": {"fill": "#1c1917", "stroke": "#d97706", "text": "#d97706"},
        "ucie": {"fill": "#1e1b4b", "stroke": "#8b5cf6", "text": "#8b5cf6"},
    }

    for rkey, rval in routers.items():
        if rval is None:
            continue
        rx, ry = rval["pos_mm"]
        px, py = mm2px(rx, ry)
        attach = rval.get("attach", [])
        is_top = ry < cube_h / 2

        # ── Router circle ──
        has_attach = len(attach) > 0
        r_fill = "#475569" if has_attach else "#334155"
        r_stroke = "#64748b" if has_attach else "#475569"
        parts.append(
            f'  <circle cx="{px:.0f}" cy="{py:.0f}" r="{r_size}" '
            f'fill="{r_fill}" stroke="{r_stroke}" stroke-width="1"/>'
        )
        parts.append(
            f'  <text x="{px:.0f}" y="{py + 3:.0f}" text-anchor="middle" '
            f'font-family="monospace" font-size="6" fill="white">'
            f'{rkey}</text>'
        )

        # ── Router → HBM_CTRL line (deferred, drawn after component blocks) ──

        # ── Attached component blocks ──
        # Collect components to draw, positioned outward from router
        blocks: list[tuple[str, str, dict]] = []  # (label, kind, style)
        pe_items = [a for a in attach if a.endswith(".dma")]
        if pe_items:
            pe_name = pe_items[0].split(".")[0].upper()
            blocks.append((pe_name, "pe", _COMP_STYLE["pe"]))
        if "m_cpu" in attach:
            blocks.append(("M_CPU", "mcpu", _COMP_STYLE["mcpu"]))
        if "sram" in attach:
            blocks.append(("SRAM", "sram", _COMP_STYLE["sram"]))
        # UCIe handled separately below

        # Position blocks outward from router (away from cube center)
        for bi, (label, kind, style) in enumerate(blocks):
            # Determine placement direction: PE/components go outward
            # Use left/right offset for multiple blocks on same router
            offset_x = (bi - (len(blocks) - 1) / 2) * (blk_w + 4)

            gap = 30  # px gap between router and component (room for 2 × 45° stubs)
            if kind == "mcpu":
                # M_CPU: place above (north of) router
                bx = px - blk_w / 2
                by = py - r_size - blk_h - gap
            elif kind == "sram":
                # SRAM: place below (south of) router
                bx = px - blk_w / 2
                by = py + r_size + gap
            else:
                # PE: place above (top half) or below (bottom half)
                bx = px + offset_x - blk_w / 2
                if is_top:
                    by = py - r_size - blk_h - gap - bi * (blk_h + 2)
                else:
                    by = py + r_size + gap + bi * (blk_h + 2)

            # Block rect
            parts.append(
                f'  <rect x="{bx:.0f}" y="{by:.0f}" '
                f'width="{blk_w}" height="{blk_h}" '
                f'rx="3" fill="{style["fill"]}" stroke="{style["stroke"]}" stroke-width="1"/>'
            )
            # Label
            font_sz = 6 if len(label) > 6 else 7
            parts.append(
                f'  <text x="{bx + blk_w / 2:.0f}" y="{by + blk_h / 2 + 3:.0f}" '
                f'text-anchor="middle" font-family="monospace" font-size="{font_sz}" '
                f'font-weight="bold" fill="{style["text"]}">{_escape(label)}</text>'
            )
            # Connector: rule-based (short → 45° line, long → 45°-straight-45°)
            sc = style["stroke"]

            # Determine start (router edge) and end (component edge) points
            bxc = bx + blk_w / 2  # component center x
            if kind == "mcpu":
                rx0, ry0 = px, py - r_size  # router top
                cx0, cy0 = bxc, by + blk_h  # component bottom
            elif kind == "sram":
                rx0, ry0 = px, py + r_size  # router bottom
                cx0, cy0 = bxc, by  # component top
            elif is_top:
                rx0, ry0 = px, py - r_size  # router top
                cx0, cy0 = bx + blk_w / 2 + offset_x, by + blk_h  # component bottom
            else:
                rx0, ry0 = px, py + r_size  # router bottom
                cx0, cy0 = bx + blk_w / 2 + offset_x, by  # component top

            # PE/M_CPU/SRAM directly above/below router (same X):
            # single diagonal line from router center to component right edge
            if abs(cx0 - rx0) < 2 and abs(cy0 - ry0) > 4:
                cx0 = bx + blk_w - 2
                parts.append(
                    f'  <line x1="{rx0:.0f}" y1="{ry0:.0f}" '
                    f'x2="{cx0:.0f}" y2="{cy0:.0f}" '
                    f'stroke="{sc}" stroke-width="1" opacity="0.6"/>'
                )
            else:
                pts = _connector_points(rx0, ry0, cx0, cy0)
                parts.append(
                    f'  <polyline points="{pts}" '
                    f'fill="none" stroke="{sc}" stroke-width="1" opacity="0.6"/>'
                )

        # (PE→HBM BW annotation drawn in the PE→HBM port group section above)

    # ── PE Router → HBM pseudo channel port group lines ──
    # Each PE router connects to its port group center on the HBM edge
    for rkey, rval in routers.items():
        if rval is None:
            continue
        attach = rval.get("attach", [])
        pe_dma_items = [a for a in attach if a.endswith(".dma")]
        if not pe_dma_items:
            continue
        pe_id = int(pe_dma_items[0].split(".")[0].replace("pe", ""))
        if pe_id not in _pe_hbm_targets:
            continue
        rx, ry = rval["pos_mm"]
        rpx, rpy = mm2px(rx, ry)
        tgx, tgy = _pe_hbm_targets[pe_id]
        r_edge_y = rpy + r_size if rpy < hbm_y else rpy - r_size
        # Rule-based connector: router → HBM port group
        pts = _connector_points(rpx, r_edge_y, tgx, tgy)
        parts.append(
            f'  <polyline points="{pts}" '
            f'fill="none" stroke="#10b981" stroke-width="1.5" opacity="0.6" '
            f'stroke-dasharray="4,3"/>'
        )
        # BW annotation at midpoint
        mx = (rpx + tgx) / 2 + 10
        my = (r_edge_y + tgy) / 2
        parts.append(
            f'  <text x="{mx:.0f}" y="{my:.0f}" '
            f'font-family="monospace" font-size="6" fill="#10b98188">'
            f'{agg_bw:.0f}GB/s</text>'
        )

    # ── UCIe port components (position/size from topology.yaml) ──
    # ucie_mm.size = 2.0mm, positions at cube edges (flush)
    ucie_size_mm = cube.get("geometry", {}).get("ucie_mm", {}).get("size", 2.0)
    uh_half = ucie_size_mm * 0.3  # half-height for edge placement
    uw_half = ucie_size_mm * 0.5
    ucie_positions = {
        "N": (cube_w / 2, uh_half),                # flush top edge
        "S": (cube_w / 2, cube_h - uh_half),       # flush bottom edge
        "W": (uh_half, cube_h / 2),                 # flush left edge
        "E": (cube_w - uh_half, cube_h / 2),        # flush right edge
    }

    # Collect UCIe connections per direction
    ucie_by_dir: dict[str, list[tuple[str, str, float, float]]] = {}
    for rkey, rval in routers.items():
        if rval is None:
            continue
        rx, ry = rval["pos_mm"]
        for a in rval.get("attach", []):
            if not a.startswith("ucie_"):
                continue
            parts_a = a.split(".")
            direction = parts_a[0].replace("ucie_", "").upper()
            conn = parts_a[1] if len(parts_a) > 1 else "c0"
            ucie_by_dir.setdefault(direction, []).append((conn, rkey, rx, ry))

    ucie_colors = ["#818cf8", "#a78bfa", "#c084fc", "#e879f9"]

    for direction, conns in ucie_by_dir.items():
        conns.sort(key=lambda x: x[0])
        n_conn = len(conns)
        ucx_mm, ucy_mm = ucie_positions.get(direction, (cube_w / 2, cube_h / 2))
        ucx, ucy = mm2px(ucx_mm, ucy_mm)

        # UCIe box: size from topology, N/S horizontal, E/W vertical
        us = ucie_size_mm * scale
        if direction in ("N", "S"):
            uw, uh = us, us * 0.5
        else:
            uw, uh = us * 0.5, us

        ux = ucx - uw / 2
        uy = ucy - uh / 2

        # UCIe component background
        parts.append(
            f'  <rect x="{ux:.0f}" y="{uy:.0f}" '
            f'width="{uw:.0f}" height="{uh:.0f}" '
            f'rx="3" fill="#1e1b4b" stroke="#8b5cf6" stroke-width="1.5" opacity="0.9"/>'
        )
        # UCIe direction label
        parts.append(
            f'  <text x="{ucx:.0f}" y="{uy - 3:.0f}" text-anchor="middle" '
            f'font-family="monospace" font-size="7" font-weight="bold" fill="#8b5cf6">'
            f'UCIe-{direction}</text>'
        )

        # Connection port boxes inside UCIe component
        for ci, (conn, rkey, crx, cry) in enumerate(conns):
            c_color = ucie_colors[ci % len(ucie_colors)]
            if direction in ("N", "S"):
                cw = max((uw - 4) / n_conn - 1, 6)
                ch = uh - 4
                cx = ux + 2 + ci * (cw + 1)
                cy_box = uy + 2
            else:
                cw = uw - 4
                ch = max((uh - 4) / n_conn - 1, 6)
                cx = ux + 2
                cy_box = uy + 2 + ci * (ch + 1)

            parts.append(
                f'  <rect x="{cx:.0f}" y="{cy_box:.0f}" '
                f'width="{cw:.0f}" height="{ch:.0f}" '
                f'rx="2" fill="{c_color}" opacity="0.7"/>'
            )
            lx = cx + cw / 2
            ly_t = cy_box + ch / 2 + 3
            parts.append(
                f'  <text x="{lx:.0f}" y="{ly_t:.0f}" text-anchor="middle" '
                f'font-family="monospace" font-size="5" fill="white">'
                f'{conn}</text>'
            )

            # Connector: rule-based router → UCIe port
            rpx, rpy = mm2px(crx, cry)
            if direction == "N":
                rx, ry = rpx, rpy - r_size
                tx, ty = lx, cy_box + ch
            elif direction == "S":
                rx, ry = rpx, rpy + r_size
                tx, ty = lx, cy_box
            elif direction == "W":
                rx, ry = rpx - r_size, rpy
                tx, ty = cx + cw, cy_box + ch / 2
            elif direction == "E":
                rx, ry = rpx + r_size, rpy
                tx, ty = cx, cy_box + ch / 2
            else:
                continue
            pts = _connector_points(rx, ry, tx, ty)
            parts.append(
                f'  <polyline points="{pts}" '
                f'fill="none" stroke="{c_color}" stroke-width="1" opacity="0.5"/>'
            )

    # ── Legend ──
    ly = h_px - 35
    legend_items = [
        ("#3b82f6", "PE Router"),
        ("#f59e0b", "M_CPU / SRAM"),
        ("#8b5cf6", "UCIe"),
        ("#334155", "Relay"),
        ("#10b981", "HBM Link"),
        ("#475569", "Mesh Link"),
    ]
    lx = pad
    for color, label in legend_items:
        parts.append(
            f'  <rect x="{lx}" y="{ly}" width="10" height="10" rx="2" '
            f'fill="{color}" stroke="#475569" stroke-width="0.5"/>'
        )
        parts.append(
            f'  <text x="{lx + 14}" y="{ly + 9}" '
            f'font-family="monospace" font-size="8" fill="#94a3b8">'
            f'{label}</text>'
        )
        lx += len(label) * 7 + 24

    parts.append("</svg>")
    return "\n".join(parts)