#!/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. Patch panel port ticks/accent
strips are parented to the panel's body cell, so they move with it too.

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", "#E8A33D", "#5B7FD9", "#9A6BC9", "#3FAAB0", "#C2547E", "#7A8C3F"]
def vlan_color(vid):
    return VLAN_PALETTE[hash(vid) % len(VLAN_PALETTE)]

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) ----------------
SWITCH_ICON = 46          # icon box; label renders below it
SWITCH_LABEL_H = 30
PANEL_W, PANEL_H = 132, 42
COL_W = 200
ROW_H = 86
PAD = 40
LABEL_STRIP = 34
SECTION_GAP = 70
MARGIN = 40

def footprint(dev):
    if dev["role"] == "patch-panel":
        return PANEL_W, PANEL_H
    return SWITCH_ICON, SWITCH_ICON + SWITCH_LABEL_H

pos = {}  # device_id -> (x, y) TOP-LEFT of the node's own bounding box (not center)

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]
        w, h = footprint(devs[node])
        cx = MARGIN + PAD + d * COL_W + SWITCH_ICON / 2
        cy = y_offset + PAD + LABEL_STRIP + (y - min_y) * ROW_H + SWITCH_ICON / 2
        pos[node] = (cx - w / 2, cy - h / 2, w, h)
    width = MARGIN * 2 + PAD * 2 + max_x * COL_W + SWITCH_ICON
    height = PAD * 2 + LABEL_STRIP + (max(ys) - min_y) * ROW_H + SWITCH_ICON + SWITCH_LABEL_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
cursor_y = MARGIN + TITLE_H + LEGEND_ROW_H + VLAN_TABLE_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)
    w, h = footprint(devs[node])
    cx = MARGIN + PAD + c * COL_W + SWITCH_ICON / 2
    cy = iso_top + PAD + LABEL_STRIP + r * ROW_H + SWITCH_ICON / 2
    pos[node] = (cx - w / 2, cy - h / 2, w, h)
iso_rows = (len(iso_sorted) + ISO_COLS - 1) // ISO_COLS
w_iso = MARGIN * 2 + PAD * 2 + (ISO_COLS - 1) * COL_W + SWITCH_ICON
h_iso = PAD * 2 + LABEL_STRIP + (iso_rows - 1) * ROW_H + SWITCH_ICON + SWITCH_LABEL_H

PAGE_W = max(w_main, w_dmo, w_iso, 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="#3A4457", stroke="#242C38", font="#FFFFFF", 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 ----------------
# Switches: a real network-switch icon, hostname caption below it.
# Patch panels: a small rack-style unit -- dark body, a lighter header strip,
# and a row of port ticks along the bottom -- built from cells parented to
# the panel's own body cell, so the whole unit moves together when dragged.
def panel_shape(did, x, y, w, h, name, st):
    body_id = f"n{did}"
    vertex(body_id, name, x, y, w, h,
           f'rounded=1;arcSize=10;whiteSpace=wrap;html=1;fillColor={st["fill"]};strokeColor={st["stroke"]};'
           f'fontColor={st["font"]};fontSize=9;verticalAlign=bottom;spacingBottom=5;')
    # header strip (lighter steel), relative to the body
    vertex(f"{body_id}-strip", "", 5, 4, w - 10, 5,
           'rounded=1;arcSize=40;fillColor=#8C97A6;strokeColor=none;', parent=body_id)
    # port ticks along the top-middle, evenly spaced
    n_ticks = 6
    tick_w, tick_h = 10, 7
    usable = w - 16
    gap = usable / n_ticks
    for k in range(n_ticks):
        tx = 8 + gap * k + (gap - tick_w) / 2
        vertex(f"{body_id}-t{k}", "", tx, 15, tick_w, tick_h,
               'rounded=0;fillColor=#2A313D;strokeColor=#5B6577;strokeWidth=0.75;', parent=body_id)

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 ""
    if dev["role"] == "patch-panel":
        panel_shape(did, x, y, w, h, dev["name"], st)
    else:
        style = (
            'shape=mxgraph.networks.switch;html=1;verticalLabelPosition=bottom;verticalAlign=top;'
            f'labelBackgroundColor=none;fontSize=10;fontColor={st["font"] if st["font"] != "#FFFFFF" else "#1A2233"};'
            f'fillColor={st["fill"]};strokeColor={st["stroke"]};fontStyle={st["bold"]};{dashed}'
        )
        vertex(f"n{did}", dev["name"], x, y, SWITCH_ICON, SWITCH_ICON, 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"),
    ("#3A4457", "#242C38", "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 / 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 table renders as an explicit "not configured" placeholder, not as
# blank space or made-up VLAN numbers.
vt_x, vt_y = MARGIN, 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")