#!/usr/bin/env python3
import json
import subprocess
import uuid
from datetime import datetime, timezone
from pathlib import Path
import xml.etree.ElementTree as ET
from xml.sax.saxutils import escape as xml_escape

SNAPSHOT = Path('/home/nikola/codex-cli/projects/incus-topology-map/data/incus-snapshot-20260409-132237.json')
OUT = Path('/home/nikola/codex-cli/projects/incus-topology-map/incus-topology-corporate.drawio')


def safe_id(prefix: str) -> str:
    return f"{prefix}_{uuid.uuid4().hex[:8]}"


def is_private_bridge_ip(ip: str) -> bool:
    if not ip:
        return False
    octets = ip.split('.')
    if len(octets) != 4:
        return False
    try:
        a, b = int(octets[0]), int(octets[1])
    except ValueError:
        return False
    return a == 172 and 16 <= b <= 31


def extract_ips(instance: dict) -> list[str]:
    out = []
    net = ((instance.get('state') or {}).get('network') or {})
    for iface in net.values():
        for addr in iface.get('addresses', []):
            if addr.get('family') == 'inet' and addr.get('scope') == 'global':
                ip = addr.get('address')
                if ip and ip not in out:
                    out.append(ip)
    return out


def primary_ip(ips: list[str]) -> str:
    for ip in ips:
        if not is_private_bridge_ip(ip):
            return ip
    return ips[0] if ips else 'n/a'


def add_vertex(root, _id, value, style, x, y, w, h, parent='1'):
    cell = ET.SubElement(root, 'mxCell', {
        'id': _id,
        'value': value,
        'style': style,
        'vertex': '1',
        'parent': parent,
    })
    ET.SubElement(cell, 'mxGeometry', {
        'x': str(x), 'y': str(y), 'width': str(w), 'height': str(h), 'as': 'geometry'
    })
    return cell


def add_edge(root, _id, source, target, style, parent='1'):
    cell = ET.SubElement(root, 'mxCell', {
        'id': _id,
        'edge': '1',
        'parent': parent,
        'source': source,
        'target': target,
        'style': style,
    })
    ET.SubElement(cell, 'mxGeometry', {'relative': '1', 'as': 'geometry'})
    return cell


def html(text: str) -> str:
    return xml_escape(str(text))


def compact_endpoint(addr: str) -> str:
    if not addr or addr == 'n/a':
        return 'n/a'
    return addr.replace('https://', '').replace('http://', '')


def remote_palette(name: str) -> tuple[str, str, str]:
    if name.startswith('hetzner'):
        return ('#FFF4E8', '#E38B1A', '#7A3E00')
    if name == 'local':
        return ('#F5F7FA', '#8B97A8', '#273244')
    return ('#ECF3FF', '#3E7BDA', '#123765')


def remote_label(name: str, server_name: str, clustered: bool, addr: str, vm_count: int, node_count: int) -> str:
    endpoint = compact_endpoint(addr)
    endpoint_line = f'🌍 {html(endpoint)}<br>' if endpoint != 'n/a' else ''
    return (
        f'<div style="line-height:1.35;">'
        f'<span style="font-size:17px;"><b>{html(name)}</b></span><br>'
        f'🖥 <b>{html(server_name)}</b> | cluster: <b>{str(clustered).lower()}</b><br>'
        f'{endpoint_line}'
        f'📦 nodes: <b>{node_count}</b> | vms: <b>{vm_count}</b>'
        f'</div>'
    )


def classify_vm(vm_name: str) -> tuple[str, str, str, str]:
    n = (vm_name or '').lower()
    if any(k in n for k in ('gateway', 'proxy', 'ingress', 'route', 'vpn', 'wacli')):
        return ('network', 'NET', '🌐', '#1D4ED8')
    if any(k in n for k in ('postgres', 'mysql', 'mongo', 'redis', 'cassandra', 'supabase', 'db', 'pg-')):
        return ('database', 'DB', '🗄', '#7C3AED')
    if any(k in n for k in ('jenkins', 'runner', 'ci', 'build', 'deploy', 'harness')):
        return ('cicd', 'CI', '⚙', '#B45309')
    if any(k in n for k in ('grafana', 'prometheus', 'loki', 'elk', 'monitor', 'uptime', 'alert')):
        return ('observability', 'OBS', '📈', '#0F766E')
    if any(k in n for k in ('auth', 'vault', 'keycloak', 'infisical', 'secret')):
        return ('security', 'SEC', '🔐', '#BE123C')
    if any(k in n for k in ('shell', 'console', 'test', 'ubuntu', 'stage', 'airstrip', 'fileserver')):
        return ('utility', 'UTIL', '🧰', '#475569')
    return ('application', 'APP', '🧩', '#166534')


def node_category_summary(vms: list[dict]) -> str:
    counts: dict[str, int] = {}
    labels: dict[str, str] = {}
    for vm in vms:
        key, short, _, _ = classify_vm(vm.get('name', ''))
        counts[key] = counts.get(key, 0) + 1
        labels[key] = short
    ordered = sorted(counts.items(), key=lambda item: (-item[1], item[0]))
    parts = [f'{labels[key]} {value}' for key, value in ordered[:4]]
    return ' | '.join(parts) if parts else 'n/a'


def node_label(node_name: str, vm_count: int, running: int, stopped: int, categories: str) -> str:
    return (
        f'<div style="line-height:1.35;">'
        f'<b>Node: {html(node_name)}</b><br>'
        f'VMs: <b>{vm_count}</b> | 🟢 {running} | 🔴 {stopped}<br>'
        f'roles: {html(categories)}'
        f'</div>'
    )


def vm_label(project: str, vm_name: str, status: str, pip: str) -> str:
    is_running = str(status).lower() == 'running'
    state_text = 'RUNNING' if is_running else 'STOPPED'
    state_dot = '🟢' if is_running else '🔴'
    display_name = f'{project}/{vm_name}' if project != 'default' else vm_name
    if len(display_name) > 26:
        display_name = display_name[:23] + '...'
    return (
        f'<div style="line-height:1.3;">'
        f'<span style="font-size:13px;"><b>{html(display_name)}</b></span><br>'
        f'{state_dot} <b>{state_text}</b> | {html(pip)}'
        f'</div>'
    )


def main():
    data = json.loads(SNAPSHOT.read_text())
    remotes = [r for r in data['remotes'] if r.get('name') != 'local']
    try:
        remotes_meta = json.loads(
            subprocess.check_output(["incus", "remote", "list", "--format", "json"], text=True)
        )
    except Exception:
        remotes_meta = {}

    mxfile = ET.Element('mxfile', {
        'host': 'app.diagrams.net',
        'modified': datetime.now(timezone.utc).isoformat(),
        'agent': 'codex-gpt-5',
        'version': '24.7.17',
        'type': 'device',
    })
    diagram = ET.SubElement(mxfile, 'diagram', {'id': uuid.uuid4().hex[:12], 'name': 'Incus Corporate Topology'})
    model = ET.SubElement(diagram, 'mxGraphModel', {
        'dx': '2000', 'dy': '1200', 'grid': '1', 'gridSize': '10', 'guides': '1',
        'tooltips': '1', 'connect': '1', 'arrows': '1', 'fold': '1', 'page': '1',
        'pageScale': '1', 'pageWidth': '2200', 'pageHeight': '1300', 'math': '0', 'shadow': '0'
    })
    root = ET.SubElement(model, 'root')
    ET.SubElement(root, 'mxCell', {'id': '0'})
    ET.SubElement(root, 'mxCell', {'id': '1', 'parent': '0'})

    title_style = (
        'rounded=1;whiteSpace=wrap;html=1;align=left;verticalAlign=middle;'
        'fontSize=22;fontStyle=1;fontFamily=Inter;fillColor=#F8FAFC;fontColor=#0F172A;strokeColor=#CBD5E1;'
        'spacingLeft=20;arcSize=10;'
    )
    generated_at = data.get('generated_at', 'n/a')
    add_vertex(
        root,
        safe_id('title'),
        f'<div><b>Incus Infrastructure Topology</b><br>Generated: {html(generated_at)}</div>',
        title_style,
        20,
        20,
        2160,
        78,
    )

    remote_x = 20
    remote_y = 165
    remote_w = 525
    remote_gap = 20
    max_cols = 4
    # Keep a fixed generous header area so metadata never gets covered by node cards.
    remote_header_h = 106

    remote_style_base = (
        'rounded=1;whiteSpace=wrap;html=1;align=left;verticalAlign=top;'
        'fontSize=13;fontStyle=1;fontFamily=Inter;spacingTop=12;spacingLeft=14;arcSize=10;'
    )
    node_style = (
        'rounded=1;whiteSpace=wrap;html=1;align=left;verticalAlign=top;'
        'fontSize=12;fontStyle=1;fontFamily=Inter;fillColor=#F8FAFC;strokeColor=#CBD5E1;fontColor=#0F172A;'
        'spacingTop=8;spacingLeft=10;arcSize=8;'
    )
    vm_running_style = (
        'rounded=1;whiteSpace=wrap;html=1;align=left;verticalAlign=top;'
        'fontSize=11;fontFamily=Inter;fontColor=#0B1F14;fillColor=#ECFDF3;strokeColor=#86EFAC;'
        'spacingTop=10;spacingLeft=10;arcSize=6;'
    )
    vm_stopped_style = (
        'rounded=1;whiteSpace=wrap;html=1;align=left;verticalAlign=top;'
        'fontSize=11;fontFamily=Inter;fontColor=#3A1414;fillColor=#FEF2F2;strokeColor=#FCA5A5;'
        'spacingTop=10;spacingLeft=10;arcSize=6;'
    )
    vm_badge_style_base = (
        'shape=ellipse;whiteSpace=wrap;html=1;align=center;verticalAlign=middle;'
        'fontSize=9;fontStyle=1;fontFamily=Inter;strokeWidth=0;'
    )
    edge_style = (
        'endArrow=none;rounded=1;strokeColor=#94A3B8;strokeWidth=1;opacity=20;dashed=1;'
        'orthogonalLoop=1;jettySize=auto;orthogonal=1;'
    )

    remotes_sorted = sorted(
        remotes,
        key=lambda r: (
            0 if r['name'] == 'local' else 1 if not r['name'].startswith('hetzner') else 2,
            r['name'],
        ),
    )

    total_nodes = 0
    total_vms = 0
    total_running = 0
    total_stopped = 0

    for remote in remotes_sorted:
        instances = remote.get('instances', [])
        total_vms += len(instances)
        total_running += sum(1 for vm in instances if str(vm.get('status', '')).lower() == 'running')
        total_stopped += sum(1 for vm in instances if str(vm.get('status', '')).lower() != 'running')

        server = remote.get('server', {})
        env = server.get('environment', {})
        server_name = env.get('server_name', 'unknown')
        nodes = {}
        for inst in instances:
            loc = inst.get('location') or 'none'
            node_name = server_name if loc in ('none', '') else loc
            nodes.setdefault(node_name, []).append(inst)
        if not nodes:
            nodes = {server_name: []}
        total_nodes += len(nodes)

    summary_style = (
        'rounded=1;whiteSpace=wrap;html=1;align=left;verticalAlign=middle;'
        'fontSize=13;fontStyle=1;fontFamily=Inter;fillColor=#FFFFFF;strokeColor=#CBD5E1;arcSize=8;spacingLeft=12;'
    )
    chip_w = 210
    chip_h = 36
    chip_y = 115
    add_vertex(root, safe_id('s1'), f'Remotes: <b>{len(remotes_sorted)}</b>', summary_style, 20, chip_y, chip_w, chip_h)
    add_vertex(root, safe_id('s2'), f'Nodes: <b>{total_nodes}</b>', summary_style, 240, chip_y, chip_w, chip_h)
    add_vertex(root, safe_id('s3'), f'VMs: <b>{total_vms}</b>', summary_style, 460, chip_y, chip_w, chip_h)
    add_vertex(root, safe_id('s4'), f'Running: <b>{total_running}</b>', vm_running_style, 680, chip_y, chip_w, chip_h)
    add_vertex(root, safe_id('s5'), f'Stopped: <b>{total_stopped}</b>', vm_stopped_style, 900, chip_y, chip_w, chip_h)

    max_remote_bottom = 0
    row_bottom = remote_y

    for idx_remote, remote in enumerate(remotes_sorted):
        name = remote['name']
        server = remote.get('server', {})
        env = server.get('environment', {})
        server_name = env.get('server_name', 'unknown')
        clustered = env.get('server_clustered', False)
        addr = (remotes_meta.get(name) or {}).get('Addr') or 'n/a'

        instances = remote.get('instances', [])
        nodes = {}
        for inst in instances:
            loc = inst.get('location') or 'none'
            node_name = server_name if loc in ('none', '') else loc
            nodes.setdefault(node_name, []).append(inst)

        if not nodes:
            nodes = {server_name: []}

        node_heights = []
        for node_name, vms in nodes.items():
            rows = max(1, (len(vms) + 1) // 2)
            node_h = 50 + rows * 74 + 16
            node_heights.append(node_h)

        remote_h = remote_header_h + sum(node_heights) + (len(node_heights) - 1) * 16 + 20

        fill, stroke, font = remote_palette(name)
        remote_style = remote_style_base + f'fillColor={fill};strokeColor={stroke};fontColor={font};'

        remote_id = safe_id('remote')
        remote_label_value = remote_label(name, server_name, clustered, addr if addr else 'n/a', len(instances), len(nodes))
        add_vertex(root, remote_id, remote_label_value, remote_style, remote_x, remote_y, remote_w, remote_h)

        current_y = remote_y + remote_header_h
        for node_name, vms in nodes.items():
            rows = max(1, (len(vms) + 1) // 2)
            node_h = 50 + rows * 74 + 16
            node_id = safe_id('node')
            running = sum(1 for vm in vms if str(vm.get('status', '')).lower() == 'running')
            stopped = len(vms) - running
            categories = node_category_summary(vms)
            add_vertex(
                root,
                node_id,
                node_label(node_name, len(vms), running, stopped, categories),
                node_style,
                remote_x + 12,
                current_y,
                remote_w - 24,
                node_h,
            )

            vm_w = (remote_w - 24 - 18) / 2
            for idx, vm in enumerate(vms):
                col = idx % 2
                row = idx // 2
                vm_x = remote_x + 20 + col * (vm_w + 10)
                vm_y = current_y + 46 + row * 74

                status = vm.get('status', 'Unknown')
                vm_name = vm.get('name', 'unknown')
                project = vm.get('project', 'default')
                ips = extract_ips(vm)
                pip = primary_ip(ips)
                _, short_role, role_icon, role_color = classify_vm(vm_name)

                vm_style = vm_running_style if str(status).lower() == 'running' else vm_stopped_style
                vm_style += 'strokeWidth=1.2;'
                vm_id = safe_id('vm')
                add_vertex(root, vm_id, vm_label(project, vm_name, status, pip), vm_style, vm_x, vm_y, vm_w, 66)
                add_vertex(
                    root,
                    safe_id('role'),
                    role_icon,
                    vm_badge_style_base + f'fillColor={role_color};strokeColor={role_color};fontColor=#FFFFFF;',
                    vm_x + vm_w - 26,
                    vm_y + 5,
                    18,
                    18,
                )

            current_y += node_h + 16

        max_remote_bottom = max(max_remote_bottom, remote_y + remote_h)
        row_bottom = max(row_bottom, remote_y + remote_h)
        if (idx_remote + 1) % max_cols == 0:
            remote_x = 20
            remote_y = row_bottom + 20
            row_bottom = remote_y
        else:
            remote_x += remote_w + remote_gap

    legend_y = max_remote_bottom + 20
    legend_style = (
        'rounded=1;whiteSpace=wrap;html=1;align=left;verticalAlign=middle;'
        'fontSize=12;fontStyle=1;fontFamily=Inter;fillColor=#FFFFFF;strokeColor=#CBD5E1;fontColor=#0F172A;arcSize=8;'
    )
    add_vertex(root, safe_id('legend_title'), 'Legend', legend_style, 20, legend_y, 140, 36)
    add_vertex(root, safe_id('lg1'), 'Running VM', vm_running_style, 170, legend_y, 180, 36)
    add_vertex(root, safe_id('lg2'), 'Stopped VM', vm_stopped_style, 360, legend_y, 180, 36)
    add_vertex(root, safe_id('lg5'), 'Role badges: 🌐 🗄 ⚙ 📈 🔐 🧰 🧩', legend_style, 550, legend_y, 340, 36)

    ET.indent(mxfile, space='  ')
    OUT.write_text(ET.tostring(mxfile, encoding='unicode'))
    print(str(OUT))


if __name__ == '__main__':
    main()