#!/usr/bin/env python3
"""
Star Map Sculptor — The Settled Reach
Reads docs/design/star-map.json, adds/removes edges to correct topology
distribution, reclassifies all nodes, writes corrected JSON back.

Edges in star-map.json are stored as arrays: ["S-XXX", "S-YYY"]

Target distribution:
  dead_end      20%  ~60 systems  (degree 1)
  spur_end      15%  ~45 systems  (degree 2, not on cycle, one high-deg neighbour)
  through_route 25%  ~75 systems  (degree 2, not on cycle, both neighbours low-deg)
  loop_member   15%  ~45 systems  (degree 2, on a cycle)
  junction      18%  ~54 systems  (degree 3-4)
  hub            7%  ~21 systems  (degree 5+)

Strategy:
  After generation produces a sparse graph (too many dead_ends, too few loops),
  the sculptor:
    Pass 1 — Add loop-forming edges to convert some dead_ends/through_routes
              into loop_members. Prefer same-sector pairs at BFS distance 4-10.
    Pass 2 — Strip excess edges from over-connected core sector.
    Pass 3 — Strip excess edges from deep_frontier / outer band.
    Pass 4 — General cleanup to a target edge count (~330 edges).
    Fine-tune — Iterate if loop_member or dead_end counts are still off.

Standard library only. No external dependencies.
"""

from __future__ import annotations

import json
from collections import defaultdict, deque
from pathlib import Path
from typing import Optional

REPO_ROOT   = Path(__file__).parent.parent
INPUT_JSON  = REPO_ROOT / "docs" / "design" / "star-map.json"
OUTPUT_JSON = INPUT_JSON

TARGETS = {
    "dead_end":      0.20,
    "spur_end":      0.15,
    "through_route": 0.25,
    "loop_member":   0.15,
    "junction":      0.18,
    "hub":           0.07,
}

# ── Graph helpers ──────────────────────────────────────────────────────────────

def build_adjacency(edges: list[list]) -> dict[str, set[str]]:
    adj: dict[str, set[str]] = defaultdict(set)
    for e in edges:
        a, b = e[0], e[1]
        adj[a].add(b)
        adj[b].add(a)
    return adj


def is_connected(all_nodes: list[str], adj: dict[str, set[str]]) -> bool:
    if not all_nodes:
        return True
    start = all_nodes[0]
    visited = {start}
    queue = deque([start])
    while queue:
        cur = queue.popleft()
        for nb in adj.get(cur, set()):
            if nb not in visited:
                visited.add(nb)
                queue.append(nb)
    return len(visited) == len(all_nodes)


def find_bridges(all_nodes: list[str], adj: dict[str, set[str]]) -> set[frozenset]:
    """
    Iterative Tarjan bridge-finding on undirected graph.
    Returns set of frozenset({u, v}) for each bridge edge.
    """
    n = len(all_nodes)
    idx_map = {v: i for i, v in enumerate(all_nodes)}
    disc = [-1] * n
    low  = [-1] * n
    timer = [0]
    bridges: set[frozenset] = set()

    for start_node in all_nodes:
        si = idx_map[start_node]
        if disc[si] != -1:
            continue

        nbrs_start = sorted(adj.get(start_node, set()))
        # Stack: (node, parent_or_None, sorted_neighbor_list, current_index)
        stack: list[tuple[str, Optional[str], list[str], int]] = [
            (start_node, None, nbrs_start, 0)
        ]
        disc[si] = low[si] = timer[0]
        timer[0] += 1

        while stack:
            node, par, nbr_list, ni = stack[-1]
            node_i = idx_map[node]

            if ni < len(nbr_list):
                nb = nbr_list[ni]
                stack[-1] = (node, par, nbr_list, ni + 1)
                nb_i = idx_map[nb]
                if disc[nb_i] == -1:
                    disc[nb_i] = low[nb_i] = timer[0]
                    timer[0] += 1
                    nb_nbrs = sorted(adj.get(nb, set()))
                    stack.append((nb, node, nb_nbrs, 0))
                elif nb != par:
                    low[node_i] = min(low[node_i], disc[nb_i])
            else:
                stack.pop()
                if par is not None:
                    par_i = idx_map[par]
                    low[par_i] = min(low[par_i], low[node_i])
                    if low[node_i] > disc[par_i]:
                        bridges.add(frozenset([par, node]))

    return bridges


def find_cycle_members(all_nodes: list[str], adj: dict[str, set[str]]) -> set[str]:
    """
    A node is on a cycle iff at least one of its incident edges is not a bridge.
    We only need this for degree-2 nodes (loop_member vs through_route/spur_end).
    """
    bridge_set = find_bridges(all_nodes, adj)
    on_cycle: set[str] = set()
    for node in all_nodes:
        neighbors = adj.get(node, set())
        if len(neighbors) < 2:
            continue
        non_bridge_count = sum(
            1 for nb in neighbors
            if frozenset([node, nb]) not in bridge_set
        )
        if non_bridge_count > 0:
            on_cycle.add(node)
    return on_cycle


def classify_node(
    node_id: str,
    adj: dict[str, set[str]],
    cycle_members: set[str],
    is_gateway: bool,
) -> str:
    if is_gateway:
        return "hub"
    deg = len(adj.get(node_id, set()))
    if deg <= 1:
        return "dead_end"
    if deg == 2:
        if node_id in cycle_members:
            return "loop_member"
        neighbors = list(adj[node_id])
        nb_degs = [len(adj.get(nb, set())) for nb in neighbors]
        # spur_end: NOT on a cycle, and at least one neighbour is branching
        # (degree >= 3). A spur hangs off a busier node — it doesn't require
        # BOTH neighbours to be high-degree.
        if any(d >= 3 for d in nb_degs):
            return "spur_end"
        return "through_route"
    if deg <= 4:
        return "junction"
    return "hub"


def reclassify_all(nodes: list[dict], adj: dict[str, set[str]]) -> None:
    all_ids = [n["system_id"] for n in nodes]
    cycle_members = find_cycle_members(all_ids, adj)
    for n in nodes:
        nid = n["system_id"]
        is_gw = n.get("is_gateway", False)
        topo = classify_node(nid, adj, cycle_members, is_gw)
        n["gate_topology"] = topo
        deg = len(adj.get(nid, set()))
        n["gate_connections"] = deg
        if is_gw:
            # Gateway: keep aperture_count=5 (4 active Reach + 1 dormant Sol)
            n["aperture_count"] = max(5, deg)
        elif topo == "hub":
            # Hub: one spare aperture (max 8 per setting rules)
            n["aperture_count"] = min(8, deg + 1)
        else:
            # All others: aperture == connections (no spare)
            n["aperture_count"] = deg


def count_topology(nodes: list[dict], adj: dict[str, set[str]]) -> dict[str, int]:
    all_ids = [n["system_id"] for n in nodes]
    cycle_members = find_cycle_members(all_ids, adj)
    counts: dict[str, int] = defaultdict(int)
    for n in nodes:
        t = classify_node(n["system_id"], adj, cycle_members, n.get("is_gateway", False))
        counts[t] += 1
    return dict(counts)


# ── Edge removal ──────────────────────────────────────────────────────────────

def remove_edge(a: str, b: str, adj: dict[str, set[str]], edges: list) -> None:
    adj[a].discard(b)
    adj[b].discard(a)
    edges[:] = [
        e for e in edges
        if not ((e[0] == a and e[1] == b) or (e[0] == b and e[1] == a))
    ]


def remove_edges_to_target(
    nodes: list[dict],
    edges: list,
    adj: dict[str, set[str]],
    target: int,
    mode: str = "default",
) -> int:
    """
    Remove up to `target` non-bridge edges using the given priority mode.
    mode: "core"     — only remove edges where both endpoints are in core sector
          "frontier" — prefer deep_frontier and outer band
          "default"  — combined: core > frontier > outer > rest, weighted by degree
    Returns: number of edges actually removed.
    """
    node_map = {n["system_id"]: n for n in nodes}
    all_ids = [n["system_id"] for n in nodes]
    removed = 0

    while removed < target:
        bridge_set = find_bridges(all_ids, adj)

        best_score: Optional[float] = None
        best_a: Optional[str] = None
        best_b: Optional[str] = None

        for e in edges:
            a, b = e[0], e[1]
            key = frozenset([a, b])

            if key in bridge_set:
                continue

            da = len(adj.get(a, set()))
            db = len(adj.get(b, set()))
            if da <= 1 or db <= 1:
                continue

            na = node_map.get(a, {})
            nb_d = node_map.get(b, {})
            if na.get("is_gateway") or nb_d.get("is_gateway"):
                continue

            sector_a = na.get("geographic_sector", "")
            sector_b = nb_d.get("geographic_sector", "")
            band_a   = na.get("geographic_band", "")
            band_b   = nb_d.get("geographic_band", "")

            if mode == "core":
                if sector_a != "core" or sector_b != "core":
                    continue
                score = -(da + db)

            elif mode == "frontier":
                is_frontier = sector_a == "deep_frontier" or sector_b == "deep_frontier"
                is_outer    = band_a == "outer" or band_b == "outer"
                if not (is_frontier or is_outer):
                    continue
                score = -(da + db)
                if is_frontier:
                    score -= 100

            else:  # default
                if sector_a == "core" and sector_b == "core":
                    base = -400
                elif sector_a == "core" or sector_b == "core":
                    base = -200
                elif sector_a == "deep_frontier" and sector_b == "deep_frontier":
                    base = -350
                elif sector_a == "deep_frontier" or sector_b == "deep_frontier":
                    base = -175
                elif band_a == "outer" and band_b == "outer":
                    base = -100
                elif band_a == "outer" or band_b == "outer":
                    base = -50
                else:
                    base = 0
                score = base - (da + db)

            if best_score is None or score < best_score:
                best_score = score
                best_a, best_b = a, b

        if best_a is None:
            print(f"  [sculptor] No removable edges at {removed}/{target} (mode={mode})")
            break

        remove_edge(best_a, best_b, adj, edges)
        removed += 1

    return removed


# ── Edge addition ─────────────────────────────────────────────────────────────

def bfs_path_length(adj: dict[str, set[str]], start: str, end: str) -> int:
    """BFS shortest path length from start to end. Returns 9999 if unreachable."""
    if start == end:
        return 0
    visited = {start}
    queue = deque([(start, 0)])
    while queue:
        node, dist = queue.popleft()
        for nb in adj.get(node, set()):
            if nb == end:
                return dist + 1
            if nb not in visited:
                visited.add(nb)
                queue.append((nb, dist + 1))
    return 9999


def add_edge(a: str, b: str, adj: dict[str, set[str]], edges: list) -> None:
    adj[a].add(b)
    adj[b].add(a)
    edges.append([a, b])


def add_loop_edges(
    nodes: list[dict],
    edges: list,
    adj: dict[str, set[str]],
    target: int,
    min_path: int = 4,
    max_path: int = 10,
) -> int:
    """
    Add up to `target` loop-forming edges.
    Connects pairs of nodes that are already connected by a path of length
    [min_path, max_path] — adding the shortcut edge creates a cycle, converting
    nodes on that path to loop_members.
    Prioritises same-sector pairs on inner/core bands.
    Avoids gateway (S-001) and avoids edges that already exist.
    Returns: number of edges added.
    """
    node_map = {n["system_id"]: n for n in nodes}
    all_ids = [n["system_id"] for n in nodes]
    edges_set = {(e[0], e[1]) for e in edges} | {(e[1], e[0]) for e in edges}
    added = 0

    # Build candidate pairs: same sector, reasonable degree (not already hubs)
    # and not already connected
    candidates: list[tuple[str, str, int]] = []  # (a, b, score)

    for i, na in enumerate(nodes):
        a = na["system_id"]
        if a == "S-001":
            continue
        da = len(adj.get(a, set()))
        if da >= 5:  # don't inflate hubs
            continue
        sector_a = na.get("geographic_sector", "")
        band_a = na.get("geographic_band", "")

        for nb in nodes[i + 1:]:
            b = nb["system_id"]
            if b == "S-001":
                continue
            if (a, b) in edges_set:
                continue
            db = len(adj.get(b, set()))
            if db >= 5:
                continue
            sector_b = nb.get("geographic_sector", "")
            band_b = nb.get("geographic_band", "")

            # Must be same sector for clean loop topology
            if sector_a != sector_b:
                continue

            # Score: prefer inner-band pairs
            score = 0
            if band_a in ("inner", "core"):
                score += 1
            if band_b in ("inner", "core"):
                score += 1
            candidates.append((a, b, score))

    # Sort by score desc, then shuffle within score groups for variety
    import random as _random
    _random.shuffle(candidates)
    candidates.sort(key=lambda x: -x[2])

    for a, b, _ in candidates:
        if added >= target:
            break
        # Check current path length — must be in [min_path, max_path]
        # (if shorter, adding this edge just creates a tiny cycle; if longer,
        # the loop created spans too many systems to classify cleanly)
        if (a, b) in edges_set or (b, a) in edges_set:
            continue
        path_len = bfs_path_length(adj, a, b)
        if min_path <= path_len <= max_path:
            add_edge(a, b, adj, edges)
            edges_set.add((a, b))
            edges_set.add((b, a))
            added += 1

    return added


# ── Reporting ─────────────────────────────────────────────────────────────────

def print_distribution(counts: dict[str, int], total: int, label: str) -> None:
    print(f"\n  {label}:")
    order = ["dead_end", "spur_end", "through_route", "loop_member", "junction", "hub"]
    for t in order:
        c = counts.get(t, 0)
        pct = c / total * 100
        tgt = TARGETS.get(t, 0) * 100
        delta = pct - tgt
        flag = " <<<" if abs(delta) > 4 else ""
        print(f"    {t:<14}  {c:>4} ({pct:5.1f}%)  target={tgt:.0f}%  delta={delta:+.1f}%{flag}")


# ── Main ───────────────────────────────────────────────────────────────────────

def main() -> None:
    print(f"Reading {INPUT_JSON} ...")
    with open(INPUT_JSON) as f:
        data = json.load(f)

    nodes: list[dict] = data["nodes"]
    edges: list       = data["edges"]
    total = len(nodes)
    print(f"  {total} nodes, {len(edges)} edges")

    # Verify edge format
    sample = edges[0]
    assert isinstance(sample, list) and len(sample) == 2, \
        f"Unexpected edge format: {sample}"

    adj = build_adjacency(edges)
    all_ids = [n["system_id"] for n in nodes]

    assert is_connected(all_ids, adj), "ERROR: Input graph is not connected!"
    print("  Connectivity: OK")

    print_distribution(count_topology(nodes, adj), total, "BEFORE")

    # ── Diagnostic: core sector ──────────────────────────────────────────────
    node_map = {n["system_id"]: n for n in nodes}
    core_nodes = [n for n in nodes if n.get("geographic_sector") == "core"]
    core_degs = sorted(
        [(n["system_id"], len(adj.get(n["system_id"], set()))) for n in core_nodes],
        key=lambda x: -x[1]
    )
    print(f"\n  Core sector: {len(core_nodes)} nodes")
    print("  Top core degrees:")
    for nid, deg in core_degs[:10]:
        print(f"    {nid}: deg={deg}")

    # ── Pass 1: Add loop-forming edges ───────────────────────────────────────
    # Generation now leaves the graph sparse (too many dead_ends, ~0 loops).
    # Add edges that form clean cycles of length 4-10 within the same sector.
    # Target: push loop_member count toward ~45 (15% of 300).
    # We add more than strictly needed here because Pass 3/4 may remove some.
    target_lm = int(TARGETS["loop_member"] * total)       # 45
    current_lm = count_topology(nodes, adj).get("loop_member", 0)
    loops_to_add = max(0, target_lm - current_lm + 10)   # +10 cushion
    if loops_to_add > 0:
        print(f"\n  Pass 1: add loop-forming edges (target +{loops_to_add})...")
        r1 = add_loop_edges(nodes, edges, adj, loops_to_add, min_path=4, max_path=10)
        print(f"  Added {r1}. Edges: {len(edges)}")
        assert is_connected(all_ids, adj), "ERROR: Disconnected after pass 1!"
        reclassify_all(nodes, adj)
        print_distribution(count_topology(nodes, adj), total, "After pass 1")

    # ── Pass 2: Strip core sector ────────────────────────────────────────────
    # Core hub count may be inflated. Strip non-bridge core edges to bring
    # hub count to ~3 in core (Gateway + 2 others).
    core_hubs = sum(
        1 for n in nodes
        if n.get("geographic_sector") == "core"
        and len(adj.get(n["system_id"], set())) >= 5
    )
    core_strip = max(0, (core_hubs - 3) * 2)  # rough: each hub removal ~2 edges
    if core_strip > 0:
        print(f"\n  Pass 2: core sector strip (target -{core_strip} edges)...")
        r2 = remove_edges_to_target(nodes, edges, adj, core_strip, mode="core")
        print(f"  Removed {r2}. Edges: {len(edges)}")
        assert is_connected(all_ids, adj), "ERROR: Disconnected after pass 2!"
        reclassify_all(nodes, adj)
        print_distribution(count_topology(nodes, adj), total, "After pass 2")

    # ── Pass 3: Frontier and outer strip ─────────────────────────────────────
    # Deep frontier should be sparse. Remove excess connections there.
    frontier_hubs = sum(
        1 for n in nodes
        if n.get("geographic_sector") == "deep_frontier"
        and len(adj.get(n["system_id"], set())) >= 5
    )
    frontier_strip = frontier_hubs * 2
    if frontier_strip > 0:
        print(f"\n  Pass 3: frontier/outer strip (target -{frontier_strip} edges)...")
        r3 = remove_edges_to_target(nodes, edges, adj, frontier_strip, mode="frontier")
        print(f"  Removed {r3}. Edges: {len(edges)}")
        assert is_connected(all_ids, adj), "ERROR: Disconnected after pass 3!"
        reclassify_all(nodes, adj)
        print_distribution(count_topology(nodes, adj), total, "After pass 3")

    # ── Pass 4: General cleanup to target edge count ─────────────────────────
    # Target ~330 edges for 300 nodes gives an average degree of 2.2,
    # consistent with the target distribution (lots of dead_ends + through_routes,
    # moderate loops, fewer junctions/hubs).
    target_edges = 330
    current_edges = len(edges)
    if current_edges > target_edges:
        need = current_edges - target_edges
        print(f"\n  Pass 4: general cleanup (target -{need} edges to reach {target_edges})...")
        r4 = remove_edges_to_target(nodes, edges, adj, need, mode="default")
        print(f"  Removed {r4}. Edges: {len(edges)}")
        assert is_connected(all_ids, adj), "ERROR: Disconnected after pass 4!"
        reclassify_all(nodes, adj)
        c4 = count_topology(nodes, adj)
        print_distribution(c4, total, "After pass 4")
    else:
        reclassify_all(nodes, adj)
        c4 = count_topology(nodes, adj)

    # ── Fine-tune: loop_member deficit? Add more loops ────────────────────────
    lm_now = c4.get("loop_member", 0)
    tolerance = int(total * 0.04)  # 4pp = 12 systems
    if lm_now < target_lm - tolerance:
        deficit = target_lm - lm_now
        print(f"\n  Fine-tune A: {lm_now} loop_members < target {target_lm}, +{deficit} edges...")
        add_loop_edges(nodes, edges, adj, deficit, min_path=4, max_path=12)
        assert is_connected(all_ids, adj), "ERROR: Disconnected after fine-tune A!"
        reclassify_all(nodes, adj)
        cf = count_topology(nodes, adj)
        print_distribution(cf, total, "After fine-tune A")
        c4 = cf

    # ── Fine-tune: too many loop_members? Strip more ──────────────────────────
    lm_now = c4.get("loop_member", 0)
    if lm_now > target_lm + tolerance:
        extra = min((lm_now - target_lm) // 3, 20)
        print(f"\n  Fine-tune B: {lm_now} loop_members > target {target_lm}, -{extra} edges...")
        remove_edges_to_target(nodes, edges, adj, extra, mode="default")
        assert is_connected(all_ids, adj), "ERROR: Disconnected after fine-tune B!"
        reclassify_all(nodes, adj)
        cf = count_topology(nodes, adj)
        print_distribution(cf, total, "After fine-tune B")

    # ── Update metadata ──────────────────────────────────────────────────────
    data["_meta"]["edge_count"] = len(edges)
    data["_meta"]["sculpted"]   = "2026-03-13"
    data["_meta"]["sculpt_note"] = (
        "Sculpted topology: loop edges added, excess core/frontier edges removed. "
        "Reclassified all nodes. Aperture counts updated."
    )

    print(f"\n  Writing {OUTPUT_JSON} ...")
    with open(OUTPUT_JSON, "w") as f:
        json.dump(data, f, indent=2)
    print("  Written.")

    # ── Final report ─────────────────────────────────────────────────────────
    print("\n" + "=" * 60)
    print("FINAL TOPOLOGY SUMMARY")
    print("=" * 60)
    final = count_topology(nodes, adj)
    print_distribution(final, total, "Final distribution")
    print(f"\n  Nodes: {total}   Edges: {len(edges)}")

    print("\n  Breakdown by sector:")
    st: dict[str, dict[str, int]] = defaultdict(lambda: defaultdict(int))
    for n in nodes:
        st[n.get("geographic_sector", "?")][n["gate_topology"]] += 1
    for sec in sorted(st.keys()):
        sd = st[sec]
        n_in = sum(sd.values())
        parts = "  ".join(f"{t}={c}" for t, c in sorted(sd.items()))
        print(f"    {sec:<15}  n={n_in:>3}  {parts}")

    print("\n  Core sector final degrees:")
    core_final = [
        (n["system_id"], len(adj.get(n["system_id"], set())), n["gate_topology"])
        for n in nodes if n.get("geographic_sector") == "core"
    ]
    core_final.sort(key=lambda x: -x[1])
    for nid, deg, topo in core_final:
        print(f"    {nid}: deg={deg}  {topo}")

    # Verify aperture >= connections for all nodes
    violations = [
        n for n in nodes
        if n.get("aperture_count", 0) < n.get("gate_connections", 0)
    ]
    if violations:
        print(f"\n  WARNING: {len(violations)} nodes have aperture < connections!")
        for v in violations:
            print(f"    {v['system_id']}: aperture={v['aperture_count']} < connections={v['gate_connections']}")
    else:
        print("\n  Aperture >= connections: OK (all nodes)")


if __name__ == "__main__":
    main()