MakerFLOSS/scripts/gen_rack.py

#!/usr/bin/env python3
"""Generate per-rack elevation SVG + page from hardware frontmatter.

Reads `docs/hardware/*.md`, selects files that declare a `rack` field,
validates rack placement, and writes for each rack:
  docs/infrastructure/racks/<rack>-elevation.svg
  docs/infrastructure/racks/<rack>.md

Deterministic, offline, stdlib + PyYAML. Non-zero exit on schema violation.
The physical rack is labeled U1 at the top; the SVG renders U1 at the top.
"""

from __future__ import annotations

import re
import sys
from pathlib import Path

import yaml

REPO_ROOT = Path(__file__).resolve().parent.parent
HARDWARE_DIR = REPO_ROOT / "docs" / "hardware"
OUTPUT_DIR = REPO_ROOT / "docs" / "infrastructure" / "racks"
RACK_UNITS = 48

FRONTMATTER_RE = re.compile(r"^---\s*\n(.*?)\n---\s*\n", re.DOTALL)

FACES = {"front", "rear", "both", "left", "right"}
ZERO_U_FACES = {"left", "right"}
SHELF_FACES = {"front", "rear"}

KIND_COLORS = {
    "server": "#4c78a8",
    "switch": "#59a14f",
    "patch-panel": "#9c755f",
    "pdu": "#e15759",
    "ups": "#edc948",
    "shelf": "#bab0ac",
    "kvm": "#b07aa1",
    "blank": "#d4d4d4",
}
DEFAULT_COLOR = "#888888"


class SchemaError(Exception):
    pass


def parse_frontmatter(path: Path) -> dict | None:
    text = path.read_text(encoding="utf-8")
    m = FRONTMATTER_RE.match(text)
    if not m:
        return None
    data = yaml.safe_load(m.group(1))
    if not isinstance(data, dict):
        raise SchemaError(f"{path}: frontmatter is not a mapping")
    return data


def load_rack_items(hardware_dir: Path) -> list[dict]:
    """Return frontmatter dicts for hardware files that declare a rack."""
    items: list[dict] = []
    for path in sorted(hardware_dir.glob("*.md")):
        if path.name == "index.md":
            continue
        fm = parse_frontmatter(path)
        if fm is None or "rack" not in fm:
            continue
        fm = dict(fm)
        fm["_path"] = str(path)
        items.append(fm)
    return items


def validate_item(fm: dict) -> None:
    name = fm.get("hostname") or fm.get("_path", "?")
    rack = fm.get("rack")
    if not isinstance(rack, str) or not rack:
        raise SchemaError(f"{name}: rack must be a non-empty string")
    if "mounted_on" in fm:
        mounted_on = fm.get("mounted_on")
        if not isinstance(mounted_on, str) or not mounted_on:
            raise SchemaError(f"{name}: mounted_on must be a non-empty string")
        for forbidden in ("rack_u", "u_height", "rack_face"):
            if forbidden in fm:
                raise SchemaError(
                    f"{name}: mounted item must omit {forbidden}"
                )
        sface = fm.get("shelf_face")
        if sface not in SHELF_FACES:
            raise SchemaError(
                f"{name}: shelf_face={sface!r} not in {sorted(SHELF_FACES)}"
            )
        slot = fm.get("shelf_slot")
        if not isinstance(slot, int) or slot < 1:
            raise SchemaError(f"{name}: shelf_slot must be an integer >= 1")
        return
    face = fm.get("rack_face")
    if face not in FACES:
        raise SchemaError(f"{name}: rack_face={face!r} not in {sorted(FACES)}")
    if face in ZERO_U_FACES:
        if "rack_u" in fm or "u_height" in fm:
            raise SchemaError(
                f"{name}: 0U item (face={face}) must omit rack_u/u_height"
            )
        return
    u = fm.get("rack_u")
    h = fm.get("u_height")
    if not isinstance(u, int) or not isinstance(h, int):
        raise SchemaError(f"{name}: rack_u and u_height must be integers")
    if u < 1 or u > RACK_UNITS:
        raise SchemaError(f"{name}: rack_u={u} out of range 1..{RACK_UNITS}")
    if h < 1:
        raise SchemaError(f"{name}: u_height={h} must be >= 1")
    if u + h - 1 > RACK_UNITS:
        raise SchemaError(
            f"{name}: occupies U{u}..U{u + h - 1}, exceeds {RACK_UNITS}U"
        )


def check_overlaps(items: list[dict]) -> None:
    """Raise if two items share a U on the same face within one rack."""
    occupied: dict[tuple[str, int], str] = {}
    for fm in items:
        if "mounted_on" in fm:
            continue
        face = fm.get("rack_face")
        if face in ZERO_U_FACES:
            continue
        faces = ("front", "rear") if face == "both" else (face,)
        u = fm["rack_u"]
        h = fm["u_height"]
        name = fm.get("hostname", "?")
        for f in faces:
            for uu in range(u, u + h):
                key = (f, uu)
                if key in occupied:
                    raise SchemaError(
                        f"U{uu} {f}: {name} overlaps {occupied[key]}"
                    )
                occupied[key] = name


def check_shelves(items: list[dict]) -> None:
    """Validate shelf-mounted devices within one rack.

    Every mounted_on resolves to a placed kind:shelf item in the same rack;
    no two devices share (shelf, face, slot).
    """
    by_host = {fm.get("hostname"): fm for fm in items}
    occupied: dict[tuple[str, str, int], str] = {}
    for fm in items:
        if "mounted_on" not in fm:
            continue
        name = fm.get("hostname", "?")
        shelf_name = fm["mounted_on"]
        target = by_host.get(shelf_name)
        if target is None:
            raise SchemaError(
                f"{name}: mounted_on={shelf_name!r} is not in this rack"
            )
        if target.get("kind") != "shelf":
            raise SchemaError(
                f"{name}: mounted_on={shelf_name!r} is not a kind:shelf item"
            )
        if not isinstance(target.get("rack_u"), int) or not isinstance(
            target.get("u_height"), int
        ):
            raise SchemaError(
                f"{name}: shelf {shelf_name!r} is not placed (needs rack_u/u_height)"
            )
        key = (shelf_name, fm["shelf_face"], fm["shelf_slot"])
        if key in occupied:
            raise SchemaError(
                f"{shelf_name} {fm['shelf_face']} slot {fm['shelf_slot']}: "
                f"{name} overlaps {occupied[key]}"
            )
        occupied[key] = name


def load_hardware_index(hardware_dir: Path) -> dict[str, dict]:
    """Map hostname -> frontmatter for every hardware file (global peer lookup)."""
    index: dict[str, dict] = {}
    for path in sorted(hardware_dir.glob("*.md")):
        if path.name == "index.md":
            continue
        fm = parse_frontmatter(path)
        if fm is None:
            continue
        name = fm.get("hostname")
        if isinstance(name, str) and name:
            index[name] = fm
    return index


def validate_links(items: list[dict], hw_index: dict[str, dict]) -> None:
    """Validate `links` cable declarations (rule 4).

    Every links[].peer must resolve to a real hardware file (global lookup via
    hw_index); peer_port must fall within the peer's declared `ports` when it
    declares an integer count.
    """
    for fm in items:
        links = fm.get("links")
        if links is None:
            continue
        name = fm.get("hostname", "?")
        if not isinstance(links, list):
            raise SchemaError(f"{name}: links must be a list")
        for link in links:
            if not isinstance(link, dict):
                raise SchemaError(f"{name}: links entry must be a mapping")
            local = link.get("local")
            peer = link.get("peer")
            peer_port = link.get("peer_port")
            if not isinstance(local, str) or not local:
                raise SchemaError(f"{name}: links entry needs a non-empty 'local'")
            if not isinstance(peer, str) or not peer:
                raise SchemaError(f"{name}: links entry needs a non-empty 'peer'")
            if not isinstance(peer_port, int):
                raise SchemaError(
                    f"{name}: links entry for {peer} needs an integer 'peer_port'"
                )
            target = hw_index.get(peer)
            if target is None:
                raise SchemaError(
                    f"{name}: links peer={peer!r} is not a known hardware file"
                )
            ports = target.get("ports")
            if isinstance(ports, int) and (peer_port < 1 or peer_port > ports):
                raise SchemaError(
                    f"{name}: peer_port {peer_port} out of range 1..{ports} on {peer}"
                )


def _pdu_index(items: list[dict]) -> dict[str, dict]:
    """Map hostname -> frontmatter for every kind:pdu item."""
    return {
        fm.get("hostname"): fm
        for fm in items
        if fm.get("kind") == "pdu"
    }


def validate_power(items: list[dict]) -> None:
    """Validate PDU outlet declarations and `power` feeds within one rack.

    Rule 3: every power[].pdu resolves to a kind:pdu file, and outlet is
    within that PDU's `outlets` count.
    """
    pdus = _pdu_index(items)
    for name, fm in pdus.items():
        outlets = fm.get("outlets")
        if not isinstance(outlets, int) or outlets < 1:
            raise SchemaError(
                f"{name}: kind:pdu must declare a positive integer 'outlets'"
            )
    for fm in items:
        feeds = fm.get("power")
        if feeds is None:
            continue
        name = fm.get("hostname", "?")
        if not isinstance(feeds, list):
            raise SchemaError(f"{name}: power must be a list")
        for feed in feeds:
            if not isinstance(feed, dict):
                raise SchemaError(f"{name}: power entry must be a mapping")
            pdu = feed.get("pdu")
            outlet = feed.get("outlet")
            if not isinstance(pdu, str) or not pdu:
                raise SchemaError(f"{name}: power entry needs a non-empty 'pdu'")
            if not isinstance(outlet, int):
                raise SchemaError(
                    f"{name}: power entry for {pdu} needs an integer 'outlet'"
                )
            target = pdus.get(pdu)
            if target is None:
                raise SchemaError(
                    f"{name}: power pdu={pdu!r} is not a known kind:pdu file"
                )
            count = target["outlets"]
            if outlet < 1 or outlet > count:
                raise SchemaError(
                    f"{name}: outlet {outlet} out of range 1..{count} on {pdu}"
                )


def _esc(s: object) -> str:
    return str(s).replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")


def _sorted_items(items: list[dict]) -> list[dict]:
    """Deterministic order: faced items by U then hostname, 0U items last."""
    return sorted(
        items,
        key=lambda i: (
            0 if i.get("rack_face") not in ZERO_U_FACES else 1,
            i.get("rack_u", 0) if isinstance(i.get("rack_u"), int) else 0,
            i.get("hostname", ""),
        ),
    )


def render_svg(rack: str, items: list[dict]) -> str:
    U_H = 20
    COL_W = 240
    LABEL_W = 30
    RAIL_W = 16
    PAD = 12
    GAP = 50
    TITLE_H = 28

    items = _sorted_items(items)
    left_items = [i for i in items if i.get("rack_face") == "left"]
    right_items = [i for i in items if i.get("rack_face") == "right"]

    body_h = RACK_UNITS * U_H
    height = PAD + TITLE_H + body_h + PAD
    front_x = PAD + len(left_items) * RAIL_W + LABEL_W
    rear_x = front_x + COL_W + GAP
    width = rear_x + COL_W + len(right_items) * RAIL_W + PAD
    top = PAD + TITLE_H

    def u_y(u: int) -> int:
        # U1 at the top; U numbers increase downward.
        return top + (u - 1) * U_H

    p: list[str] = []
    p.append(
        f'<svg xmlns="http://www.w3.org/2000/svg" width="{width}" '
        f'height="{height}" viewBox="0 0 {width} {height}" '
        f'font-family="sans-serif" font-size="11">'
    )
    p.append(f'<rect width="{width}" height="{height}" fill="#ffffff"/>')
    p.append(
        f'<text x="{PAD}" y="{PAD + 16}" font-size="16" '
        f'font-weight="bold">Rack {_esc(rack)}</text>'
    )

    for col_x, col_label in ((front_x, "front"), (rear_x, "rear")):
        p.append(
            f'<text x="{col_x + COL_W // 2}" y="{top - 6}" '
            f'text-anchor="middle" font-weight="bold">{col_label}</text>'
        )
        for u in range(1, RACK_UNITS + 1):
            y = u_y(u)
            p.append(
                f'<rect x="{col_x}" y="{y}" width="{COL_W}" height="{U_H}" '
                f'fill="#f5f5f5" stroke="#e0e0e0"/>'
            )

    # U numbers in the gutter left of the front column.
    for u in range(1, RACK_UNITS + 1):
        y = u_y(u)
        p.append(
            f'<text x="{front_x - 4}" y="{y + 14}" text-anchor="end" '
            f'fill="#999">{u}</text>'
        )

    def draw_device(fm: dict, col_x: int) -> None:
        u = fm["rack_u"]
        h = fm["u_height"]
        y = u_y(u)
        box_h = h * U_H
        color = KIND_COLORS.get(fm.get("kind", ""), DEFAULT_COLOR)
        name = fm.get("hostname", "?")
        urange = f"U{u}" if h == 1 else f"U{u}–U{u + h - 1}"
        p.append(
            f'<rect x="{col_x + 1}" y="{y + 1}" width="{COL_W - 2}" '
            f'height="{box_h - 2}" rx="3" fill="{color}" stroke="#333"/>'
        )
        p.append(
            f'<text x="{col_x + COL_W // 2}" y="{y + box_h // 2 + 4}" '
            f'text-anchor="middle" fill="#ffffff">'
            f'{_esc(name)} ({urange})</text>'
        )

    for fm in items:
        if fm.get("kind") == "shelf" or "mounted_on" in fm:
            continue
        face = fm.get("rack_face")
        if face in ("front", "both"):
            draw_device(fm, front_x)
        if face in ("rear", "both"):
            draw_device(fm, rear_x)

    def draw_rail(fm: dict, x: int) -> None:
        color = KIND_COLORS.get(fm.get("kind", ""), DEFAULT_COLOR)
        name = fm.get("hostname", "?")
        cx = x + RAIL_W // 2
        cy = top + body_h // 2
        p.append(
            f'<rect x="{x}" y="{top}" width="{RAIL_W}" height="{body_h}" '
            f'fill="{color}" stroke="#333"/>'
        )
        p.append(
            f'<text x="{cx}" y="{cy}" text-anchor="middle" fill="#ffffff" '
            f'transform="rotate(-90 {cx} {cy})">{_esc(name)}</text>'
        )

    for idx, fm in enumerate(left_items):
        draw_rail(fm, PAD + idx * RAIL_W)
    for idx, fm in enumerate(right_items):
        draw_rail(fm, rear_x + COL_W + idx * RAIL_W)

    SHELF_STRIP_H = 6
    shelves = [i for i in items if i.get("kind") == "shelf"]
    mounted = [i for i in items if "mounted_on" in i]

    def draw_shelf(fm: dict) -> None:
        u = fm["rack_u"]
        h = fm["u_height"]
        y = u_y(u)
        block_h = h * U_H
        strip_y = y + block_h - SHELF_STRIP_H
        avail_h = block_h - SHELF_STRIP_H
        shelf_color = KIND_COLORS.get("shelf", DEFAULT_COLOR)
        sname = fm.get("hostname", "?")
        for col_x, sface in ((front_x, "front"), (rear_x, "rear")):
            occ = sorted(
                (m for m in mounted
                 if m.get("mounted_on") == sname
                 and m.get("shelf_face") == sface),
                key=lambda m: (m.get("shelf_slot", 0), m.get("hostname", "")),
            )
            n = len(occ)
            for idx, m in enumerate(occ):
                sub_w = COL_W // n
                bx = col_x + idx * sub_w
                bw = (COL_W - idx * sub_w) if idx == n - 1 else sub_w
                mcolor = KIND_COLORS.get(m.get("kind", ""), DEFAULT_COLOR)
                mname = m.get("hostname", "?")
                p.append(
                    f'<rect x="{bx + 1}" y="{y + 1}" width="{bw - 2}" '
                    f'height="{avail_h - 2}" rx="3" fill="{mcolor}" stroke="#333"/>'
                )
                p.append(
                    f'<text x="{bx + bw // 2}" y="{y + avail_h // 2 + 4}" '
                    f'text-anchor="middle" fill="#ffffff">{_esc(mname)}</text>'
                )
            p.append(
                f'<rect x="{col_x}" y="{strip_y}" width="{COL_W}" '
                f'height="{SHELF_STRIP_H}" fill="{shelf_color}" stroke="#333"/>'
            )
        p.append(
            f'<text x="{front_x + COL_W // 2}" y="{strip_y + SHELF_STRIP_H - 1}" '
            f'text-anchor="middle" fill="#333" font-size="9">{_esc(sname)}</text>'
        )

    for fm in sorted(shelves, key=lambda s: s.get("hostname", "")):
        draw_shelf(fm)

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


def _node_id(name: str) -> str:
    """A mermaid-safe node id derived from a hostname."""
    return re.sub(r"[^0-9A-Za-z]", "_", str(name))


def render_power(rack: str, items: list[dict]) -> str:
    """Return a mermaid power-distribution flowchart, or '' if no feeds.

    Assumes `validate_power(items)` has already passed: every referenced PDU
    resolves to a kind:pdu item with a positive-int `outlets`. `generate`
    guarantees this by validating before any render call.
    """
    powered = [fm for fm in items if fm.get("power")]
    if not powered:
        return ""
    pdus = _pdu_index(items)

    edges: list[tuple[str, int, str]] = []
    for fm in powered:
        device = fm.get("hostname", "?")
        for feed in fm["power"]:
            edges.append((feed["pdu"], feed["outlet"], device))
    edges.sort()

    lines: list[str] = ["```mermaid", "flowchart LR"]
    for pdu in sorted(pdus):
        outlets = pdus[pdu].get("outlets")
        lines.append(f'    {_node_id(pdu)}["{pdu}<br/>{outlets} outlets"]')
    devices = sorted(
        powered,
        key=lambda i: (
            i.get("rack_u", 0) if isinstance(i.get("rack_u"), int) else 0,
            i.get("hostname", ""),
        ),
    )
    for fm in devices:
        device = fm.get("hostname", "?")
        lines.append(f'    {_node_id(device)}["{device}"]')
    for pdu, outlet, device in edges:
        lines.append(
            f"    {_node_id(pdu)} -->|outlet {outlet}| {_node_id(device)}"
        )
    lines.append("```")
    return "\n".join(lines) + "\n"


def render_network(rack: str, items: list[dict]) -> str:
    """Return a mermaid network-cabling flowchart, or '' if no links.

    Assumes `validate_links` has already passed: every link has a non-empty
    `local`/`peer` and an integer `peer_port`, and `peer` resolves to a real
    hardware file. `generate` validates before any render call.
    """
    linked = [fm for fm in items if fm.get("links")]
    if not linked:
        return ""

    by_host = {fm.get("hostname"): fm for fm in items}

    edges: list[tuple[str, str, str, int, object]] = []
    nodes: set[str] = set()
    for fm in linked:
        source = fm.get("hostname", "?")
        nodes.add(source)
        for link in fm["links"]:
            peer = link["peer"]
            nodes.add(peer)
            edges.append(
                (source, link["local"], peer, link["peer_port"],
                 link.get("speed_gbps"))
            )
    edges.sort(key=lambda e: (e[0], e[1], e[2], e[3]))

    def node_label(name: str) -> str:
        fm = by_host.get(name)
        kind = fm.get("kind") if fm else None
        if kind in ("switch", "patch-panel"):
            return f"{name}<br/>{kind}"
        return name

    lines: list[str] = ["```mermaid", "flowchart LR"]
    for name in sorted(nodes):
        lines.append(f'    {_node_id(name)}["{node_label(name)}"]')
    for source, local, peer, peer_port, speed in edges:
        label = f"{local} → p{peer_port}"
        if speed is not None:
            label += f" · {speed}G"
        lines.append(f"    {_node_id(source)} -->|{label}| {_node_id(peer)}")
    lines.append("```")
    return "\n".join(lines) + "\n"


def render_page(rack: str, items: list[dict]) -> str:
    items = _sorted_items(items)
    lines: list[str] = []
    lines.append(f"# Rack {rack}")
    lines.append("")
    lines.append(
        f"_Auto-generated from `docs/hardware/*.md` (items with `rack: {rack}`) "
        f"— do not edit by hand. Run `make docs-index` after changing a "
        f"source file._"
    )
    lines.append("")
    lines.append("## Elevation")
    lines.append("")
    lines.append(f"![Rack {rack} elevation]({rack}-elevation.svg)")
    lines.append("")
    power = render_power(rack, items)
    if power:
        lines.append("## Power")
        lines.append("")
        lines.append(power.rstrip())
        lines.append("")
    network = render_network(rack, items)
    if network:
        lines.append("## Network")
        lines.append("")
        lines.append(network.rstrip())
        lines.append("")
    lines.append("## Occupancy")
    lines.append("")
    lines.append("| U | Device | Kind | Face | Status |")
    lines.append("|---|---|---|---|---|")
    by_host = {fm.get("hostname"): fm for fm in items}
    mounted_by_shelf: dict[str, list[dict]] = {}
    for fm in items:
        if "mounted_on" in fm:
            mounted_by_shelf.setdefault(fm["mounted_on"], []).append(fm)

    def occ_row(fm: dict) -> str:
        name = fm.get("hostname", "?")
        link = f"[{name}](../../hardware/{name}.md)"
        if "mounted_on" in fm:
            target = by_host.get(fm["mounted_on"])
            if target and isinstance(target.get("rack_u"), int):
                su = target["rack_u"]
                sh = target["u_height"]
                urange = f"U{su}" if sh == 1 else f"U{su}–U{su + sh - 1}"
            else:
                urange = "—"
            face = (
                f"{fm.get('shelf_face', '')} · "
                f"{fm['mounted_on']}/{fm.get('shelf_slot', '')}"
            )
        else:
            face = fm.get("rack_face", "")
            if face in ZERO_U_FACES:
                urange = "0U"
            else:
                u = fm["rack_u"]
                h = fm["u_height"]
                urange = f"U{u}" if h == 1 else f"U{u}–U{u + h - 1}"
        return (
            f"| {urange} | {link} | {fm.get('kind', '')} | {face} "
            f"| {fm.get('status', '')} |"
        )

    for fm in _sorted_items([i for i in items if "mounted_on" not in i]):
        lines.append(occ_row(fm))
        if fm.get("kind") == "shelf":
            occ = sorted(
                mounted_by_shelf.get(fm.get("hostname"), []),
                key=lambda m: (m.get("shelf_face", ""), m.get("shelf_slot", 0)),
            )
            for m in occ:
                lines.append(occ_row(m))
    lines.append("")
    return "\n".join(lines).rstrip() + "\n"


def generate(hardware_dir: Path, output_dir: Path) -> int:
    items = load_rack_items(hardware_dir)
    hw_index = load_hardware_index(hardware_dir)

    errors: list[str] = []
    for fm in items:
        try:
            validate_item(fm)
        except SchemaError as e:
            errors.append(str(e))

    racks: dict[str, list[dict]] = {}
    for fm in items:
        racks.setdefault(fm["rack"], []).append(fm)

    if not errors:  # only check overlaps once placements are individually valid
        for rack, ritems in racks.items():
            try:
                check_overlaps(ritems)
                validate_power(ritems)
                validate_links(ritems, hw_index)
                check_shelves(ritems)
            except SchemaError as e:
                errors.append(f"{rack}: {e}")

    if errors:
        for err in errors:
            print(f"ERROR: {err}", file=sys.stderr)
        return 1

    output_dir.mkdir(parents=True, exist_ok=True)
    for rack in sorted(racks):
        ritems = racks[rack]
        (output_dir / f"{rack}-elevation.svg").write_text(
            render_svg(rack, ritems), encoding="utf-8"
        )
        (output_dir / f"{rack}.md").write_text(
            render_page(rack, ritems), encoding="utf-8"
        )
        print(f"Wrote {rack}.md + {rack}-elevation.svg ({len(ritems)} item(s))")
    return 0


def main() -> int:
    return generate(HARDWARE_DIR, OUTPUT_DIR)


if __name__ == "__main__":
    sys.exit(main())