#!/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"]

# ---------------- components ----------------
adj = {d: [] for d in devs}
for a, b in edges:
    adj[a].append(b)
    adj[b].append(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 a, b in edges:
    union(a, 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
        children = [n for n in sorted(adj[node]) if n in node_set and n != par and n not in visited]
        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 = 130, 34
COL_W = 190
ROW_H = 82
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
cursor_y = MARGIN + TITLE_H

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)
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="#EEF0F2", stroke="#8C97A6", font="#495566", 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):
    add_cell(
        f'<mxCell id="{cell_id}" 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 get a real network-switch icon (label rendered below it by
# draw.io's own verticalLabelPosition, not baked into the geometry), patch panels get
# a plain labeled box.
node_style = {}
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":
        style = (
            f'rounded=0;whiteSpace=wrap;html=1;fillColor={st["fill"]};strokeColor={st["stroke"]};'
            f'fontColor={st["font"]};fontSize=9;{dashed}'
        )
        vertex(f"n{did}", dev["name"], x, y, w, h, style)
    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 -- glued natively via source/target, orthogonal routing
for idx, (a, b) in enumerate(edges):
    if a in pos and b in pos:
        edge(f"e{idx}", f"n{a}", f"n{b}",
             "edgeStyle=orthogonalEdgeStyle;rounded=1;html=1;strokeColor=#9AA6B8;strokeWidth=1.2;"
             "jettySize=auto;orthogonalLoop=1;")

# legend
legend_items = [
    ("#2F5FAD", "#173D6E", "Distribution switch"),
    ("#F3F5FA", "#7C8AA6", "Field switch"),
    ("#EEF0F2", "#8C97A6", "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

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")