#!/usr/bin/env python3
"""
Render topo_full.json (produced by export_topology.py) as a draw.io
(mxGraph) diagram: Field-Switch-Topology.drawio.

Edges reference shapes by id (source=/target=), which is how draw.io's
native glue works -- no hand-rolled formulas needed: dragging a shape in
the app keeps every attached line attached.

Usage:
    python3 generate_drawio.py [topo_full.json] [Field-Switch-Topology.drawio]
"""
import json, sys, uuid
from xml.sax.saxutils import escape
import xml.dom.minidom as minidom

IN_PATH = sys.argv[1] if len(sys.argv) > 1 else "topo_full.json"
OUT_PATH = sys.argv[2] if len(sys.argv) > 2 else "Field-Switch-Topology.drawio"

data = json.load(open(IN_PATH))
devs = {d["id"]: d for d in data["devices"]}
edges = data["edges"]  # [{a, a_port, a_vlan, b, b_port, b_vlan}, ...]
vlans = data.get("vlans", [])  # [{vid, name, description, prefixes, site}, ...]

VLAN_PALETTE = ["#D6603F", "#3C9D6B", "#C2547E", "#5B7FD9", "#E8A33D", "#3FAAB0", "#9A6BC9", "#7A8C3F"]
# assigned by position among the VLANs actually present (sorted by vid, same
# order they're listed/rendered in), not a hash -- guarantees no two VLANs
# share a color as long as there are <= len(VLAN_PALETTE) of them, which a
# hash-based assignment can't promise (two vids can hash to the same slot)
_vlan_color_by_vid = {v["vid"]: VLAN_PALETTE[i % len(VLAN_PALETTE)] for i, v in enumerate(vlans)}
def vlan_color(vid):
    return _vlan_color_by_vid.get(vid, "#9AA0AA")

def vlan_tag(pv):
    """Short connector-label suffix for a port's VLAN config, e.g. 'V10' or
    'V10,20T'. Empty string when the port has no 802.1Q config -- which, as
    of this export, is every port (see export_topology.py docstring)."""
    if not pv:
        return ""
    parts = []
    if pv.get("untagged"):
        parts.append(str(pv["untagged"]["vid"]))
    tagged = pv.get("tagged") or []
    if tagged:
        parts.append(",".join(str(v["vid"]) for v in tagged) + "T")
    return f' · V{"+".join(parts)}' if parts else ""

# ---------------- components ----------------
adj = {d: [] for d in devs}
for e in edges:
    adj[e["a"]].append(e["b"])
    adj[e["b"]].append(e["a"])

parent = {d: d for d in devs}
def find(x):
    while parent[x] != x:
        parent[x] = parent[parent[x]]
        x = parent[x]
    return x
def union(a, b):
    ra, rb = find(a), find(b)
    if ra != rb:
        parent[ra] = rb
for e in edges:
    union(e["a"], e["b"])
from collections import defaultdict
comps = defaultdict(list)
for d in devs:
    comps[find(d)].append(d)
comps = sorted(comps.values(), key=len, reverse=True)
main_comp = comps[0]   # fche fabric, rooted at id 1
dmo_comp = comps[1]    # dmo-lob fabric, rooted at id 7
iso = [c[0] for c in comps[2:]]

# ---------------- tidy tree layout ----------------
def layout_tree(root, node_ids):
    node_set = set(node_ids)
    depth, yslot, visited = {}, {}, set()
    leaf_counter = [0]
    def dfs(node, par, d):
        visited.add(node)
        depth[node] = d
        # dedupe: two devices can be joined by more than one real cable
        # (redundant links), which must not create more than one tree child
        seen_here = set()
        children = []
        for n in sorted(adj[node]):
            if n in node_set and n != par and n not in visited and n not in seen_here:
                seen_here.add(n)
                children.append(n)
        if not children:
            y = leaf_counter[0]; leaf_counter[0] += 1
            yslot[node] = y
            return y
        ys = [dfs(c, node, d + 1) for c in children]
        y = sum(ys) / len(ys)
        yslot[node] = y
        return y
    dfs(root, None, 0)
    return depth, yslot

depth_main, yslot_main = layout_tree(1, main_comp)
depth_dmo, yslot_dmo = layout_tree(7, dmo_comp)

# ---------------- geometry constants (px) ----------------
# Realistic rack-equipment silhouettes (draw.io's built-in Rack shape
# library) instead of the old abstract pictogram / hand-drawn box -- both
# switches and patch panels are real 19" rack-mount gear, so both get the
# same width and a below-icon caption; only the icon height (and which
# stencil) differs, matching each one's actual flat 1U proportions.
NODE_W = 108
SWITCH_H = 20    # ~108:20, close to the stencil's native ~5.3:1
PANEL_H = 12     # ~108:12, close to the stencil's native ~10.1:1
CAPTION_H = 46   # hostname + model, wrapped, below the icon
COL_W = 320  # NODE_W + clearance for the connector's own port/VLAN label text
ROW_H = 78
PAD = 40
LABEL_STRIP = 34
SECTION_GAP = 70
MARGIN = 40

def icon_h(dev):
    return PANEL_H if dev["role"] == "patch-panel" else SWITCH_H

pos = {}  # device_id -> (x, y, w, h) of the node's own ICON (caption sits below it)

def place_section(depth_map, yslot_map, y_offset):
    ys = list(yslot_map.values())
    min_y = min(ys)
    max_x = max(depth_map.values())
    for node, d in depth_map.items():
        y = yslot_map[node]
        h = icon_h(devs[node])
        cx = MARGIN + PAD + d * COL_W + NODE_W / 2
        top = y_offset + PAD + LABEL_STRIP + (y - min_y) * ROW_H
        pos[node] = (cx - NODE_W / 2, top, NODE_W, h)
    width = MARGIN * 2 + PAD * 2 + max_x * COL_W + NODE_W
    height = PAD * 2 + LABEL_STRIP + (max(ys) - min_y) * ROW_H + SWITCH_H + CAPTION_H
    return width, height

TITLE_H = 60
LEGEND_ROW_H = 34
VLAN_ROW_H = 26
VLAN_HEADER_H = 30
VLAN_TABLE_W = 900
vlan_rows = max(1, len(vlans))  # the placeholder counts as one row
VLAN_TABLE_H = VLAN_HEADER_H + vlan_rows * VLAN_ROW_H + 16

# a dedicated colour-coded legend (VLAN id block + name block per row),
# stacked to the left of the detail table -- this is the part that reads
# as "legend" at a glance; the table next to it carries the detail
VLAN_LEGEND_W = 260
VLAN_LEGEND_ROW_H = 42
VLAN_LEGEND_GAP = 6
VLAN_LEGEND_H = vlan_rows * (VLAN_LEGEND_ROW_H + VLAN_LEGEND_GAP) - VLAN_LEGEND_GAP

VLAN_SECTION_H = max(VLAN_TABLE_H, VLAN_LEGEND_H)
cursor_y = MARGIN + TITLE_H + LEGEND_ROW_H + VLAN_SECTION_H + 20

w_main, h_main = place_section(depth_main, yslot_main, cursor_y)
main_top = cursor_y
cursor_y += h_main + SECTION_GAP

w_dmo, h_dmo = place_section(depth_dmo, yslot_dmo, cursor_y)
dmo_top = cursor_y
cursor_y += h_dmo + SECTION_GAP

ISO_COLS = 3
iso_top = cursor_y
iso_sorted = sorted(iso, key=lambda d: devs[d]["name"])
for i, node in enumerate(iso_sorted):
    r, c = divmod(i, ISO_COLS)
    h = icon_h(devs[node])
    cx = MARGIN + PAD + c * COL_W + NODE_W / 2
    top = iso_top + PAD + LABEL_STRIP + r * ROW_H
    pos[node] = (cx - NODE_W / 2, top, NODE_W, h)
iso_rows = (len(iso_sorted) + ISO_COLS - 1) // ISO_COLS
w_iso = MARGIN * 2 + PAD * 2 + (ISO_COLS - 1) * COL_W + NODE_W
h_iso = PAD * 2 + LABEL_STRIP + (iso_rows - 1) * ROW_H + SWITCH_H + CAPTION_H

PAGE_W = max(w_main, w_dmo, w_iso, VLAN_LEGEND_W + 30 + VLAN_TABLE_W + MARGIN * 2)
PAGE_H = cursor_y + h_iso + MARGIN

# ---------------- style palette ----------------
STYLE = {
    "dist":    dict(fill="#2F5FAD", stroke="#173D6E", font="#FFFFFF", dashed=0, bold=1),
    "field":   dict(fill="#F3F5FA", stroke="#7C8AA6", font="#1A2233", dashed=0, bold=0),
    "spare":   dict(fill="#F1F1F2", stroke="#9AA0AA", font="#5B6478", dashed=1, bold=0),
    "decom":   dict(fill="#FBE6E2", stroke="#C04B34", font="#8A3423", dashed=1, bold=0),
    "planned": dict(fill="#FDF1CF", stroke="#B8862C", font="#7A5C00", dashed=1, bold=0),
    "panel":   dict(fill="#F1EDE6", stroke="#A69C89", font="#6B6154", dashed=0, bold=0),
}

def style_for(dev):
    if dev["role"] == "patch-panel":
        return STYLE["panel"]
    if dev["role"] == "dist-switch":
        return STYLE["dist"]
    if dev["status"] == "inventory":
        return STYLE["spare"]
    if dev["status"] == "decommissioning":
        return STYLE["decom"]
    if dev["status"] == "planned":
        return STYLE["planned"]
    return STYLE["field"]

# ---------------- mxGraph XML ----------------
cells = []
def add_cell(xml):
    cells.append(xml)

def vertex(cell_id, value, x, y, w, h, style, parent="1"):
    add_cell(
        f'<mxCell id="{cell_id}" value="{escape(value)}" style="{style}" '
        f'vertex="1" parent="{parent}">'
        f'<mxGeometry x="{x:.1f}" y="{y:.1f}" width="{w:.1f}" height="{h:.1f}" as="geometry"/>'
        f'</mxCell>'
    )

def edge(cell_id, source, target, style, value=""):
    add_cell(
        f'<mxCell id="{cell_id}" value="{escape(value)}" style="{style}" edge="1" parent="1" '
        f'source="{source}" target="{target}"><mxGeometry relative="1" as="geometry"/></mxCell>'
    )

def container(cell_id, x, y, w, h, label):
    add_cell(
        f'<mxCell id="{cell_id}" value="" style="rounded=1;arcSize=3;fillColor=#F7F8FB;'
        f'strokeColor=#DDE1EA;verticalAlign=top;" vertex="1" parent="1">'
        f'<mxGeometry x="{x:.1f}" y="{y:.1f}" width="{w:.1f}" height="{h:.1f}" as="geometry"/></mxCell>'
    )
    add_cell(
        f'<mxCell id="{cell_id}-label" value="{escape(label)}" '
        f'style="text;html=1;align=left;verticalAlign=top;fontStyle=1;fontSize=13;fontColor=#33415C;" '
        f'vertex="1" parent="1"><mxGeometry x="{x+14:.1f}" y="{y+8:.1f}" width="{w-28:.1f}" height="22" as="geometry"/></mxCell>'
    )

container("bg-main", MARGIN, main_top, w_main - MARGIN * 2, h_main,
           f"FCHE Cable Headend — distribution fabric  ({len(main_comp)} devices)")
container("bg-dmo", MARGIN, dmo_top, w_dmo - MARGIN * 2, h_dmo,
           f"DMO-LOB — distribution fabric  ({len(dmo_comp)} devices)")
container("bg-iso", MARGIN, iso_top, w_iso - MARGIN * 2, h_iso,
           "Not in a resolved cable path")

# title
add_cell(
    f'<mxCell id="title" value="Field Switch Topology" style="text;html=1;align=left;'
    f'verticalAlign=top;fontStyle=1;fontSize=22;fontColor=#161B26;" vertex="1" parent="1">'
    f'<mxGeometry x="{MARGIN}" y="10" width="600" height="30" as="geometry"/></mxCell>'
)

# ---------------- node shapes ----------------
# Realistic 19" rack-equipment silhouettes from draw.io's built-in Rack
# shape library (verified against the actual stencil source, not guessed):
# a real network-switch front panel for switches, a real 24-port patch
# panel front panel for patch panels. Neither is FS.com/AFL/Primex/etc.
# branded art -- no such stencils exist anywhere, this vendor mix is far
# too niche -- so "by model" is carried as text: every node's caption is
# its hostname plus its actual NetBox device_type.model, wrapped below the
# icon, instead of pretending the icon itself is model-specific.
SWITCH_STENCIL = "mxgraph.rack.general.switches_1"
PANEL_STENCIL = "mxgraph.rack.general.cat5e_rack_mount_patch_panel_24_ports"

def caption(name, model):
    # single-quote the inner HTML attribute -- this whole string is itself
    # going into an XML value="..." attribute (vertex() escapes it), and a
    # literal " from a double-quoted style="..." here would prematurely
    # close that outer XML attribute
    return f"{escape(name)}<br><font style='font-size:9px'>{escape(model)}</font>"

for did, dev in devs.items():
    if did not in pos:
        continue
    x, y, w, h = pos[did]
    st = style_for(dev)
    dashed = "dashed=1;dashPattern=4 3;" if st["dashed"] else ""
    stencil = PANEL_STENCIL if dev["role"] == "patch-panel" else SWITCH_STENCIL
    font = st["font"] if st["font"] != "#FFFFFF" else "#1A2233"
    style = (
        f'shape={stencil};html=1;whiteSpace=wrap;verticalLabelPosition=bottom;verticalAlign=top;'
        f'labelBackgroundColor=none;fontSize=11;fontColor={font};'
        f'fillColor={st["fill"]};strokeColor={st["stroke"]};fontStyle={st["bold"]};{dashed}'
    )
    vertex(f"n{did}", caption(dev["name"], dev.get("device_type", "")), x, y, w, h, style)

# ---------------- connectors, labeled with the port at each end ----------------
for idx, e in enumerate(edges):
    a, b = e["a"], e["b"]
    if a in pos and b in pos:
        label = f'{e["a_port"]}{vlan_tag(e.get("a_vlan"))} ↔ {e["b_port"]}{vlan_tag(e.get("b_vlan"))}'
        edge(f"e{idx}", f"n{a}", f"n{b}",
             "edgeStyle=orthogonalEdgeStyle;rounded=1;html=1;strokeColor=#9AA6B8;strokeWidth=1.2;"
             "startArrow=none;endArrow=none;"
             "jettySize=auto;orthogonalLoop=1;fontSize=13;fontStyle=1;fontColor=#3A4457;"
             "labelBackgroundColor=#F7F8FBee;align=center;verticalAlign=middle;",
             value=label)

# legend
legend_items = [
    ("#2F5FAD", "#173D6E", "Distribution switch"),
    ("#F3F5FA", "#7C8AA6", "Field switch"),
    ("#F1EDE6", "#A69C89", "Patch panel / splice module"),
    ("#F1F1F2", "#9AA0AA", "Spare / decommissioned / planned"),
]
lx = MARGIN
ly = 44
for i, (fill, stroke, label) in enumerate(legend_items):
    add_cell(
        f'<mxCell id="lg{i}" value="" style="rounded=0;fillColor={fill};strokeColor={stroke};" '
        f'vertex="1" parent="1"><mxGeometry x="{lx}" y="{ly}" width="14" height="14" as="geometry"/></mxCell>'
    )
    add_cell(
        f'<mxCell id="lg{i}-t" value="{escape(label)}" style="text;html=1;align=left;'
        f'verticalAlign=middle;fontSize=11;fontColor=#3A4457;" vertex="1" parent="1">'
        f'<mxGeometry x="{lx+20}" y="{ly-3}" width="220" height="20" as="geometry"/></mxCell>'
    )
    lx += 260

# ---------------- VLAN legend (colour blocks) + mapping table ----------------
# Structure is always rendered; content depends on whether NetBox actually
# has VLANs defined yet (see export_topology.py). Nothing here is invented --
# an empty legend/table renders as an explicit "not configured" placeholder,
# not as blank space or made-up VLAN numbers.
vlg_x, vlg_y = MARGIN, MARGIN + TITLE_H + LEGEND_ROW_H
vlg_id_w, vlg_name_w = 70, VLAN_LEGEND_W - 70
if not vlans:
    add_cell(
        f'<mxCell id="vlg-empty" value="No VLANs" style="rounded=1;arcSize=12;whiteSpace=wrap;html=1;'
        f'fillColor=#F1F1F2;strokeColor=#9AA0AA;dashed=1;dashPattern=4 3;fontColor=#7C8AA6;fontSize=11;" '
        f'vertex="1" parent="1"><mxGeometry x="{vlg_x}" y="{vlg_y}" width="{VLAN_LEGEND_W}" '
        f'height="{VLAN_LEGEND_ROW_H}" as="geometry"/></mxCell>'
    )
else:
    ly2 = vlg_y
    for v in vlans:
        color = vlan_color(v["vid"])
        add_cell(
            f'<mxCell id="vlg-id-{v["vid"]}" value="VLAN&#10;{v["vid"]}" style="rounded=0;whiteSpace=wrap;'
            f'html=1;fillColor={color};strokeColor=none;fontColor=#FFFFFF;fontStyle=1;fontSize=13;'
            f'align=center;verticalAlign=middle;lineHeight=1.1;" vertex="1" parent="1">'
            f'<mxGeometry x="{vlg_x}" y="{ly2}" width="{vlg_id_w}" height="{VLAN_LEGEND_ROW_H}" as="geometry"/></mxCell>'
        )
        add_cell(
            f'<mxCell id="vlg-name-{v["vid"]}" value="{escape(v["name"])}" style="rounded=0;whiteSpace=wrap;'
            f'html=1;fillColor=#FFFFFF;strokeColor={color};strokeWidth=2;fontColor=#1A2233;fontStyle=1;'
            f'fontSize=12;align=center;verticalAlign=middle;" vertex="1" parent="1">'
            f'<mxGeometry x="{vlg_x+vlg_id_w}" y="{ly2}" width="{vlg_name_w}" height="{VLAN_LEGEND_ROW_H}" '
            f'as="geometry"/></mxCell>'
        )
        ly2 += VLAN_LEGEND_ROW_H + VLAN_LEGEND_GAP

vt_x, vt_y = MARGIN + VLAN_LEGEND_W + 30, MARGIN + TITLE_H + LEGEND_ROW_H
vt_cols = [("VLAN", 60), ("Name", 130), ("Description", 300), ("Subnet", 160), ("Site", 170)]
add_cell(
    f'<mxCell id="vt-bg" value="" style="rounded=0;fillColor=#FFFFFF;strokeColor=#DDE1EA;" '
    f'vertex="1" parent="1"><mxGeometry x="{vt_x}" y="{vt_y}" width="{VLAN_TABLE_W}" '
    f'height="{VLAN_TABLE_H}" as="geometry"/></mxCell>'
)
add_cell(
    f'<mxCell id="vt-title" value="VLANs" style="text;html=1;align=left;verticalAlign=middle;'
    f'fontStyle=1;fontSize=15;fontColor=#33415C;" vertex="1" parent="1">'
    f'<mxGeometry x="{vt_x+10}" y="{vt_y}" width="200" height="{VLAN_HEADER_H}" as="geometry"/></mxCell>'
)
cx = vt_x + 90
for label, w in vt_cols:
    add_cell(
        f'<mxCell id="vt-h-{label}" value="{escape(label)}" style="text;html=1;align=left;'
        f'verticalAlign=middle;fontStyle=1;fontSize=12;fontColor=#7C8AA6;" vertex="1" parent="1">'
        f'<mxGeometry x="{cx}" y="{vt_y}" width="{w}" height="{VLAN_HEADER_H}" as="geometry"/></mxCell>'
    )
    cx += w

row_y = vt_y + VLAN_HEADER_H
if not vlans:
    add_cell(
        f'<mxCell id="vt-empty-sw" value="" style="rounded=0;fillColor=#F1F1F2;strokeColor=#9AA0AA;'
        f'dashed=1;dashPattern=4 3;" vertex="1" parent="1">'
        f'<mxGeometry x="{vt_x+10}" y="{row_y+2}" width="14" height="14" as="geometry"/></mxCell>'
    )
    add_cell(
        f'<mxCell id="vt-empty-txt" value="No VLANs defined in NetBox yet (ipam.VLAN) — this table '
        f'and each connector’s VLAN tag fill in automatically once VLANs and interface assignments '
        f'are added, next time this diagram is regenerated." style="text;html=1;align=left;'
        f'verticalAlign=middle;fontSize=12;fontColor=#7C8AA6;fontStyle=2;whiteSpace=wrap;" '
        f'vertex="1" parent="1"><mxGeometry x="{vt_x+90}" y="{row_y}" width="{VLAN_TABLE_W-100}" '
        f'height="{VLAN_ROW_H}" as="geometry"/></mxCell>'
    )
else:
    for v in vlans:
        add_cell(
            f'<mxCell id="vt-sw-{v["vid"]}" value="" style="rounded=0;fillColor={vlan_color(v["vid"])};'
            f'strokeColor=none;" vertex="1" parent="1">'
            f'<mxGeometry x="{vt_x+10}" y="{row_y+3}" width="14" height="14" as="geometry"/></mxCell>'
        )
        row_vals = [
            str(v["vid"]), v["name"], v.get("description") or "—",
            ", ".join(v.get("prefixes") or []) or "—", v.get("site") or "—",
        ]
        cx = vt_x + 90
        for (label, w), val in zip(vt_cols, row_vals):
            add_cell(
                f'<mxCell id="vt-{v["vid"]}-{label}" value="{escape(val)}" style="text;html=1;align=left;'
                f'verticalAlign=middle;fontSize=12;fontColor=#3A4457;" vertex="1" parent="1">'
                f'<mxGeometry x="{cx}" y="{row_y}" width="{w}" height="{VLAN_ROW_H}" as="geometry"/></mxCell>'
            )
            cx += w
        row_y += VLAN_ROW_H

xml = (
    '<mxfile host="app.diagrams.net" agent="netbox-topology-generator" version="24.0.0">'
    f'<diagram name="Topology" id="{uuid.uuid4()}">'
    f'<mxGraphModel dx="800" dy="600" grid="1" gridSize="10" guides="1" tooltips="1" '
    f'connect="1" arrows="1" fold="1" page="1" pageScale="1" '
    f'pageWidth="{PAGE_W:.0f}" pageHeight="{PAGE_H:.0f}" math="0" shadow="0">'
    '<root><mxCell id="0"/><mxCell id="1" parent="0"/>'
    + "".join(cells) +
    '</root></mxGraphModel></diagram></mxfile>'
)

minidom.parseString(xml)  # fail fast on malformed XML
open(OUT_PATH, "w").write(xml)
print(f"wrote {OUT_PATH}: {len(cells)} cells, page {PAGE_W:.0f}x{PAGE_H:.0f}px")