#!/usr/bin/env python3
"""Close a small, explicit set of residual PCB connections.

This is deliberately a *finisher*, not a full autorouter.  It operates on the
real KiCad board after the global Specctra route has been imported.  For each
requested net it finds the disconnected copper groups, chooses anchors from
pads, existing track endpoints and vias, and runs a fine 3-D A* search through
the board's actual copper obstacle model.  Every committed segment and via is
locked; KiCad DRC remains the acceptance test.

Unlike the older residual helpers, an endpoint does not need room for a new
via.  A route may leave a fine-pitch pad on its own copper layer, travel to a
legal via site, and only then change layers.  Existing same-net vias and track
ends are also valid start/goal anchors.
"""

from __future__ import annotations

import argparse
import heapq
import math
import shutil
import time
from dataclasses import dataclass
from pathlib import Path

import pcbnew

import route_local_power as rlp


# A 0.20-mm grid lands on both halves of the Ti375's 0.8-mm dogbone lattice
# while keeping the Python obstacle search tractable.  Endpoint stubs retain
# their exact coordinates, so off-grid package pads are not displaced.
STEP = 0.20
VIA_DIAMETER = 0.45
VIA_DRILL = 0.225

CLASS_RULES = {
    "Default": (0.12, 0.10),
    "BGA_ESCAPE": (0.10, 0.10),
    "CLOCK": (0.15, 0.15),
    # Only low-current local rails may be passed to this finisher.  Main 5-V
    # and converter paths need deterministic wide copper instead.
    "POWER": (0.25, 0.10),
}


def mm(value: int) -> float:
    return round(pcbnew.ToMM(value), 4)


def MM(value: float) -> int:
    return pcbnew.FromMM(value)


@dataclass(frozen=True)
class Anchor:
    point: tuple[float, float]
    layers: tuple[int, ...]
    label: str


class FineMaze:
    def __init__(self, board: pcbnew.BOARD) -> None:
        self.board = board
        self.layers = tuple(
            layer for layer in board.GetEnabledLayers().CuStack()
            if board.GetLayerType(layer) == pcbnew.LT_SIGNAL
        )
        self.layer_rank = {layer: index for index, layer in enumerate(self.layers)}
        self.obstacles = rlp.Obstacles(board)

    def track_clear(self, net: str, a: tuple[float, float], b: tuple[float, float],
                    layer: int, width: float, clearance: float) -> bool:
        return self.obstacles.track_clear(
            net, a, b, width, layer, clearance_floor=clearance
        )

    def via_clear(self, net: str, p: tuple[float, float], clearance: float) -> bool:
        return self.obstacles.via_clear(net, p, clearance_floor=clearance)

    @staticmethod
    def _collapse(points: list[tuple[float, float, int]]) -> list[tuple[float, float, int]]:
        if len(points) < 3:
            return points
        out = [points[0]]
        for index in range(1, len(points) - 1):
            a = out[-1]
            b = points[index]
            c = points[index + 1]
            if a[2] == b[2] == c[2]:
                cross = (b[0] - a[0]) * (c[1] - b[1]) - (b[1] - a[1]) * (c[0] - b[0])
                if abs(cross) < 1e-8:
                    continue
            out.append(b)
        out.append(points[-1])
        return out

    def search(self, net: str, start: Anchor, goal: Anchor, width: float,
               clearance: float, margin: float, deadline_s: float = 90.0,
               node_cap: int = 700_000) -> list[tuple[float, float, int]] | None:
        sx, sy = start.point
        gx, gy = goal.point
        x0 = min(sx, gx) - margin
        y0 = min(sy, gy) - margin
        x1 = max(sx, gx) + margin
        y1 = max(sy, gy) + margin

        def snap(value: float) -> int:
            return round(value / STEP)

        ix0, iy0 = snap(x0), snap(y0)
        ix1, iy1 = snap(x1), snap(y1)

        def point(ix: int, iy: int) -> tuple[float, float]:
            return round(ix * STEP, 4), round(iy * STEP, 4)

        def heuristic(ix: int, iy: int) -> float:
            px, py = point(ix, iy)
            return math.hypot(px - gx, py - gy)

        start_layers = tuple(layer for layer in start.layers if layer in self.layer_rank)
        goal_layers = frozenset(layer for layer in goal.layers if layer in self.layer_rank)
        if not start_layers or not goal_layers:
            return None

        queue: list[tuple[float, float, tuple[int, int, int]]] = []
        cost: dict[tuple[int, int, int], float] = {}
        parent: dict[tuple[int, int, int], tuple[int, int, int] | None] = {}

        # Seed several nearby grid points.  The exact pad/via/track endpoint to
        # grid stub is checked just like every other segment; no unchecked snap
        # segment can leak into the output.
        base_x, base_y = snap(sx), snap(sy)
        for layer in start_layers:
            for dx in (-1, 0, 1):
                for dy in (-1, 0, 1):
                    ix, iy = base_x + dx, base_y + dy
                    if not (ix0 <= ix <= ix1 and iy0 <= iy <= iy1):
                        continue
                    p = point(ix, iy)
                    if p != start.point and not self.track_clear(
                            net, start.point, p, layer, width, clearance):
                        continue
                    state = (ix, iy, layer)
                    initial = math.hypot(p[0] - sx, p[1] - sy)
                    if initial < cost.get(state, math.inf):
                        cost[state] = initial
                        parent[state] = None
                        heapq.heappush(queue, (initial + heuristic(ix, iy), initial, state))

        if not queue:
            return None

        stop = time.monotonic() + deadline_s
        hit: tuple[int, int, int] | None = None
        visited = 0
        directions = ((-1, -1), (-1, 0), (-1, 1), (0, -1),
                      (0, 1), (1, -1), (1, 0), (1, 1))

        while queue and visited < node_cap and time.monotonic() < stop:
            _score, current_cost, state = heapq.heappop(queue)
            if current_cost > cost.get(state, math.inf) + 1e-9:
                continue
            visited += 1
            ix, iy, layer = state
            here = point(ix, iy)

            if layer in goal_layers and math.hypot(here[0] - gx, here[1] - gy) <= 0.30:
                if here == goal.point or self.track_clear(
                        net, here, goal.point, layer, width, clearance):
                    hit = state
                    break

            for dx, dy in directions:
                nx, ny = ix + dx, iy + dy
                if not (ix0 <= nx <= ix1 and iy0 <= ny <= iy1):
                    continue
                nxt_point = point(nx, ny)
                if not self.track_clear(net, here, nxt_point, layer, width, clearance):
                    continue
                nxt = (nx, ny, layer)
                move = STEP * (math.sqrt(2.0) if dx and dy else 1.0)
                # Prefer inner signal layers slightly, but do not ban F/B.
                outer_penalty = 0.015 if layer in (pcbnew.F_Cu, pcbnew.B_Cu) else 0.0
                new_cost = current_cost + move + outer_penalty
                if new_cost + 1e-9 < cost.get(nxt, math.inf):
                    cost[nxt] = new_cost
                    parent[nxt] = state
                    heapq.heappush(queue, (new_cost + heuristic(nx, ny), new_cost, nxt))

            if self.via_clear(net, here, clearance):
                for next_layer in self.layers:
                    if next_layer == layer:
                        continue
                    nxt = (ix, iy, next_layer)
                    new_cost = current_cost + 2.0
                    if new_cost + 1e-9 < cost.get(nxt, math.inf):
                        cost[nxt] = new_cost
                        parent[nxt] = state
                        heapq.heappush(queue, (new_cost + heuristic(ix, iy), new_cost, nxt))

        if hit is None:
            return None

        states = []
        cursor: tuple[int, int, int] | None = hit
        while cursor is not None:
            states.append(cursor)
            cursor = parent[cursor]
        states.reverse()
        result = [(sx, sy, states[0][2])]
        result.extend((*point(ix, iy), layer) for ix, iy, layer in states)
        result.append((gx, gy, states[-1][2]))
        return self._collapse(result)

    def _same_net_via_at(self, net: str, p: tuple[float, float]) -> bool:
        for item in self.board.GetTracks():
            if not isinstance(item, pcbnew.PCB_VIA) or item.GetNetname() != net:
                continue
            q = item.GetPosition()
            if math.hypot(mm(q.x) - p[0], mm(q.y) - p[1]) < 0.02:
                return True
        return False

    def add_track(self, net: str, a: tuple[float, float], b: tuple[float, float],
                  layer: int, width: float) -> None:
        if math.hypot(a[0] - b[0], a[1] - b[1]) < 1e-6:
            return
        track = pcbnew.PCB_TRACK(self.board)
        track.SetStart(pcbnew.VECTOR2I_MM(*a))
        track.SetEnd(pcbnew.VECTOR2I_MM(*b))
        track.SetWidth(MM(width))
        track.SetLayer(layer)
        track.SetNet(self.board.FindNet(net))
        track.SetLocked(True)
        self.board.Add(track)
        track.thisown = False
        self.obstacles.add_track(net, a, b, width, layer)

    def add_via(self, net: str, p: tuple[float, float]) -> None:
        if self._same_net_via_at(net, p):
            return
        via = pcbnew.PCB_VIA(self.board)
        via.SetPosition(pcbnew.VECTOR2I_MM(*p))
        via.SetWidth(MM(VIA_DIAMETER))
        via.SetDrill(MM(VIA_DRILL))
        via.SetLayerPair(pcbnew.F_Cu, pcbnew.B_Cu)
        via.SetNet(self.board.FindNet(net))
        via.SetLocked(True)
        self.board.Add(via)
        via.thisown = False
        self.obstacles.add_via(net, p, VIA_DIAMETER / 2, VIA_DRILL / 2)

    def commit(self, net: str, path: list[tuple[float, float, int]], width: float) -> None:
        for a, b in zip(path, path[1:]):
            if a[2] == b[2]:
                self.add_track(net, a[:2], b[:2], a[2], width)
            else:
                if math.hypot(a[0] - b[0], a[1] - b[1]) > 0.02:
                    raise RuntimeError(f"Non-vertical layer change in {net}: {a} -> {b}")
                self.add_via(net, a[:2])


def item_anchor(item, routing_layers: tuple[int, ...], label: str) -> list[Anchor]:
    if isinstance(item, pcbnew.PAD):
        p = item.GetPosition()
        layers = tuple(layer for layer in item.GetLayerSet().Seq() if layer in routing_layers)
        return [Anchor((mm(p.x), mm(p.y)), layers, label)] if layers else []
    if isinstance(item, pcbnew.PCB_VIA):
        p = item.GetPosition()
        return [Anchor((mm(p.x), mm(p.y)), routing_layers, label)]
    if isinstance(item, pcbnew.PCB_TRACK):
        layer = item.GetLayer()
        if layer not in routing_layers:
            return []
        a, b = item.GetStart(), item.GetEnd()
        return [Anchor((mm(a.x), mm(a.y)), (layer,), label + ":start"),
                Anchor((mm(b.x), mm(b.y)), (layer,), label + ":end")]
    return []


def copper_groups(board: pcbnew.BOARD, net: str,
                  routing_layers: tuple[int, ...]) -> list[tuple[list[pcbnew.PAD], list[Anchor]]]:
    connectivity = board.GetConnectivity()
    pads = [pad for fp in board.GetFootprints() for pad in fp.Pads()
            if pad.GetNetname() == net]
    seen: set[str] = set()
    groups = []
    for pad in pads:
        uid = pad.m_Uuid.AsString()
        if uid in seen:
            continue
        connected = list(connectivity.GetConnectedItems(pad, 0))
        if all(item.m_Uuid.AsString() != uid for item in connected):
            connected.append(pad)
        group_pads = [item for item in connected if isinstance(item, pcbnew.PAD)]
        if not group_pads:
            group_pads = [pad]
        seen.update(item.m_Uuid.AsString() for item in group_pads)
        anchors: list[Anchor] = []
        anchor_seen = set()
        for index, item in enumerate(connected):
            for anchor in item_anchor(item, routing_layers, f"{item.GetClass()}#{index}"):
                key = (anchor.point, anchor.layers)
                if key not in anchor_seen:
                    anchors.append(anchor)
                    anchor_seen.add(key)
        if not anchors:
            anchors.extend(item_anchor(pad, routing_layers, "pad"))
        groups.append((group_pads, anchors))
    return groups


def class_rule(board: pcbnew.BOARD, net: str) -> tuple[float, float, str]:
    for fp in board.GetFootprints():
        for pad in fp.Pads():
            if pad.GetNetname() == net:
                name = pad.GetNetClassName()
                width, clearance = CLASS_RULES.get(name, CLASS_RULES["Default"])
                return width, clearance, name
    raise RuntimeError(f"Net {net} has no pads")


def finish_net(maze: FineMaze, net: str) -> int:
    board = maze.board
    joined = 0
    while True:
        board.BuildConnectivity()
        groups = copper_groups(board, net, maze.layers)
        if len(groups) < 2:
            return joined
        groups.sort(key=lambda group: len(group[0]), reverse=True)
        main = groups[0]
        branch = min(groups[1:], key=lambda group: len(group[0]))
        width, clearance, class_name = class_rule(board, net)

        pairs = []
        for source in branch[1]:
            for target in main[1]:
                distance = math.hypot(source.point[0] - target.point[0],
                                      source.point[1] - target.point[1])
                pairs.append((distance, source, target))
        pairs.sort(key=lambda entry: entry[0])
        routed = False
        for distance, source, target in pairs[:10]:
            for margin in (5.0, 12.0):
                path = maze.search(net, source, target, width, clearance, margin,
                                   deadline_s=20.0, node_cap=240_000)
                if path is None:
                    continue
                maze.commit(net, path, width)
                board.BuildConnectivity()
                print(f"{net}: joined {source.label} -> {target.label}; "
                      f"class={class_name} width={width:.3f} mm span={distance:.1f} mm")
                joined += 1
                routed = True
                break
            if routed:
                break
        if not routed:
            print(f"{net}: FAILED ({len(groups)} copper groups, {len(pairs)} anchor pairs)")
            return joined


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("board", type=Path)
    parser.add_argument("--output", type=Path)
    parser.add_argument("--nets", required=True,
                        help="Comma- or semicolon-separated residual net names")
    args = parser.parse_args()
    output = args.output or args.board
    nets = [name.strip() for name in args.nets.replace(";", ",").split(",")
            if name.strip()]

    board = pcbnew.LoadBoard(str(args.board))
    if not pcbnew.ZONE_FILLER(board).Fill(board.Zones()):
        raise RuntimeError("Zone refill failed before residual routing")
    board.BuildConnectivity()
    maze = FineMaze(board)

    for net in nets:
        if board.FindNet(net) is None:
            raise RuntimeError(f"Unknown net: {net}")
        finish_net(maze, net)

    if not pcbnew.ZONE_FILLER(board).Fill(board.Zones()):
        raise RuntimeError("Zone refill failed after residual routing")
    board.BuildConnectivity()
    pcbnew.SaveBoard(str(output), board)
    if output != args.board:
        for suffix in (".kicad_pro", ".kicad_dru"):
            source = args.board.with_suffix(suffix)
            if source.exists():
                shutil.copyfile(source, output.with_suffix(suffix))
    print(f"Unconnected endpoints: {board.GetConnectivity().GetUnconnectedCount(True)}")
    print(f"Saved {output}")


if __name__ == "__main__":
    main()