"""A grid-based maze router (MVP): multi-layer, vias, and rip-up-and-retry.

Still an *MVP* — not a port of FreeRouting's expansion-room maze / ripper — but
it routes across the board's signal layers, changes layers through vias, and
recovers from bad greedy net orderings with rip-up-and-retry:

* the occupancy grid has a layer axis; A* nodes are ``(col, row, layer)``;
* in-plane moves (8-connected, 45-degree) cost distance; a **via move** changes
  layer at a cell for a configurable ``via_cost``; a through via must be clear on
  every signal layer;
* occupancy is tracked *per connection* so a specific routed connection can be
  removed. When a connection cannot reach its target through free space, the
  router runs a **rip-up** search that may pass through other nets' traces (at a
  penalty), rips up the connections it crosses, routes the failing connection,
  and re-queues the ripped connections for a later pass;
* passes iterate up to ``max_passes``, keeping the best (fewest-unrouted) result
  and stopping on full success or when a pass changes nothing;
* **thrash prevention:** each connection may be ripped at most ``rip_cap`` times,
  and the penalty to rip a connection escalates with its rip count, so
  repeatedly-ripped connections harden into walls. Everything is deterministic
  (no RNG).

Deferred (still): FreeRouting's free-space expansion rooms, exact shove /
DRC-clean geometry, and blind/buried via spans.
"""

from __future__ import annotations

import heapq
import math

from freeroute.board.board import BasicBoard
from freeroute.board.items import ObstacleArea, Pin
from freeroute.geometry import IntPoint

__all__ = ["GridRouter", "RouteResult", "route_board"]

# 8-connected in-plane neighbourhood (orthogonal + diagonal) for 45-degree routing.
_NEIGHBOURS = [
    (1, 0),
    (-1, 0),
    (0, 1),
    (0, -1),
    (1, 1),
    (1, -1),
    (-1, 1),
    (-1, -1),
]

Cell = tuple[int, int]
Node = tuple[int, int, int]  # (col, row, layer)


class RouteResult:
    """Per-net routing outcome in *board* units.

    ``wires`` maps a net number to a list of ``(layer_index, [IntPoint, ...])``
    trace segments; ``vias`` maps a net number to a list of via locations.
    ``routed_net_numbers`` are the nets that got any wire or via.
    """

    __slots__ = ("wires", "vias", "half_width", "unrouted")

    def __init__(self, half_width: int) -> None:
        self.wires: dict[int, list[tuple[int, list[IntPoint]]]] = {}
        self.vias: dict[int, list[IntPoint]] = {}
        self.half_width = half_width
        #: number of ratsnest connections left unrouted
        self.unrouted = 0

    def add_wire(self, net_no: int, layer: int, path: list[IntPoint]) -> None:
        self.wires.setdefault(net_no, []).append((layer, path))

    def add_via(self, net_no: int, location: IntPoint) -> None:
        self.vias.setdefault(net_no, []).append(location)

    @property
    def routed_net_numbers(self) -> set[int]:
        nets = {n for n, w in self.wires.items() if w}
        nets |= {n for n, v in self.vias.items() if v}
        return nets

    def via_count(self) -> int:
        return sum(len(v) for v in self.vias.values())


class _Connection:
    """One ratsnest connection (a pin pair) and its routed state."""

    __slots__ = (
        "index",
        "net_no",
        "a",
        "b",
        "routed",
        "rip_count",
        "wires",
        "vias",
        "trace_cells",
        "via_cells",
    )

    def __init__(self, index: int, net_no: int, a: Pin, b: Pin) -> None:
        self.index = index
        self.net_no = net_no
        self.a = a
        self.b = b
        self.routed = False
        self.rip_count = 0
        self.wires: list[tuple[int, list[IntPoint]]] = []
        self.vias: list[IntPoint] = []
        self.trace_cells: list[tuple[Cell, int]] = []
        self.via_cells: list[Cell] = []


class GridRouter:
    """Routes a :class:`BasicBoard` with a grid A*, vias, and rip-up-and-retry."""

    def __init__(
        self,
        board: BasicBoard,
        *,
        trace_width: int,
        clearance: int,
        layers: list[int] | None = None,
        via_cost: float = 10.0,
        rip_up: bool = True,
        max_passes: int = 10,
        rip_cap: int = 4,
        rip_penalty: float = 30.0,
        orthogonal: bool = False,
    ) -> None:
        self.board = board
        self.trace_width = max(trace_width, 1)
        self.clearance = max(clearance, 0)
        self.half_width = self.trace_width // 2
        self.via_cost = via_cost
        self.rip_up = rip_up
        self.max_passes = max(max_passes, 1)
        self.rip_cap = max(rip_cap, 0)
        self.rip_penalty = rip_penalty
        # orthogonal routing (axis-aligned segments only) yields exact IntBox
        # trace copper; 8-connected also produces 45-degree segments.
        self._neighbours = _NEIGHBOURS[:4] if orthogonal else _NEIGHBOURS
        # cell size must fit a trace plus its clearance to a neighbouring trace
        self.step = max(self.trace_width + self.clearance, 1)

        if layers is None:
            layers = [i for i, layer in enumerate(board.layer_structure.arr) if layer.is_signal]
        self.layers = list(layers)

        box = board.bounding_box
        self.origin_x = box.ll.x
        self.origin_y = box.ll.y
        self.cols = max(1, math.ceil((box.ur.x - box.ll.x) / self.step) + 1)
        self.rows = max(1, math.ceil((box.ur.y - box.ll.y) / self.step) + 1)

        # static obstacles
        self._pad_block: dict[Node, set[int]] = {}
        self._keepout: set[Node] = set()
        # per-connection occupancy (owner = connection index)
        self._cell_owner: dict[Node, int] = {}
        self._via_owner: dict[Cell, int] = {}
        self._rasterize_obstacles()

    # --- grid helpers -------------------------------------------------------

    def _cell_of(self, point: IntPoint) -> Cell:
        gx = round((point.x - self.origin_x) / self.step)
        gy = round((point.y - self.origin_y) / self.step)
        gx = min(max(gx, 0), self.cols - 1)
        gy = min(max(gy, 0), self.rows - 1)
        return gx, gy

    def _cell_center(self, cell: Cell) -> IntPoint:
        return IntPoint(self.origin_x + cell[0] * self.step, self.origin_y + cell[1] * self.step)

    def _cells_in_box(self, box, margin: int):
        lo_x = round((box.ll.x - margin - self.origin_x) / self.step)
        hi_x = round((box.ur.x + margin - self.origin_x) / self.step)
        lo_y = round((box.ll.y - margin - self.origin_y) / self.step)
        hi_y = round((box.ur.y + margin - self.origin_y) / self.step)
        for gx in range(max(lo_x, 0), min(hi_x, self.cols - 1) + 1):
            for gy in range(max(lo_y, 0), min(hi_y, self.rows - 1) + 1):
                yield gx, gy

    def _rasterize_obstacles(self) -> None:
        margin = self.clearance + self.half_width
        for item in self.board.get_items():
            if isinstance(item, Pin):
                box = item.shape.bounding_box() if item.shape is not None else None
                if box is None or box.is_empty():
                    continue
                net = item.net_nos[0] if item.net_nos else 0
                pin_layers = [layer for layer in item.layers if layer in self.layers]
                for cell in self._cells_in_box(box, margin):
                    for layer in pin_layers:
                        self._pad_block.setdefault((*cell, layer), set()).add(net)
            elif isinstance(item, ObstacleArea):
                if item.layer not in self.layers:
                    continue
                box = item.bounding_box()
                if box.is_empty():
                    continue
                for cell in self._cells_in_box(box, self.clearance):
                    self._keepout.add((*cell, item.layer))

    # --- occupancy cost -----------------------------------------------------

    def _enter_cost(
        self, cell: Cell, layer: int, net_no: int, allow_ripup: bool, conns: list[_Connection]
    ) -> float | None:
        """Cost to enter ``(cell, layer)`` for ``net_no``; ``None`` if blocked.

        Static pads/keepouts always block. A cell occupied by another net's trace
        or via is impassable in free-space mode; in rip-up mode it is passable at
        a penalty that escalates with the owner's rip count, unless the owner has
        already hit ``rip_cap`` (then it hardens into a wall).
        """
        node = (*cell, layer)
        if node in self._keepout:
            return None
        pads = self._pad_block.get(node)
        if pads and any(n != net_no for n in pads):
            return None
        blocking = None
        owner = self._cell_owner.get(node)
        if owner is not None and conns[owner].net_no != net_no:
            blocking = owner
        else:
            via_owner = self._via_owner.get(cell)
            if via_owner is not None and conns[via_owner].net_no != net_no:
                blocking = via_owner
        if blocking is None:
            return 0.0
        if not allow_ripup:
            return None
        rip_count = conns[blocking].rip_count
        if rip_count >= self.rip_cap:
            return None
        return self.rip_penalty * (1 + rip_count)

    def _via_placeable(self, cell: Cell, net_no: int, conns: list[_Connection]) -> bool:
        """A through via needs every signal layer strictly free of other nets."""
        return all(
            self._enter_cost(cell, layer, net_no, False, conns) == 0.0 for layer in self.layers
        )

    # --- A* search ----------------------------------------------------------

    def _search(
        self, conn: _Connection, allow_ripup: bool, conns: list[_Connection]
    ) -> list[Node] | None:
        start_cell = self._cell_of(conn.a.location)
        goal_cell = self._cell_of(conn.b.location)
        a_layers = [layer for layer in conn.a.layers if layer in self.layers]
        b_layers = [layer for layer in conn.b.layers if layer in self.layers]
        if not a_layers or not b_layers:
            return None
        starts = {(*start_cell, layer) for layer in a_layers}
        goals = {(*goal_cell, layer) for layer in b_layers}
        if starts & goals:
            return [min(starts & goals)]
        net_no = conn.net_no

        open_heap: list[tuple[float, Node]] = []
        g_score: dict[Node, float] = {}
        came_from: dict[Node, Node] = {}
        for s in sorted(starts):
            g_score[s] = 0.0
            heapq.heappush(open_heap, (self._h(s, goal_cell), s))

        while open_heap:
            _, current = heapq.heappop(open_heap)
            if current in goals:
                return _reconstruct(came_from, current)
            cx, cy, cl = current
            base = g_score[current]
            for dx, dy in self._neighbours:
                cell = (cx + dx, cy + dy)
                if not (0 <= cell[0] < self.cols and 0 <= cell[1] < self.rows):
                    continue
                nxt = (*cell, cl)
                if nxt in goals:
                    enter = 0.0
                else:
                    enter = self._enter_cost(cell, cl, net_no, allow_ripup, conns)
                    if enter is None:
                        continue
                step = 1.0 if dx == 0 or dy == 0 else math.sqrt(2)
                self._relax(
                    current, nxt, base + step + enter, goal_cell, g_score, came_from, open_heap
                )
            if len(self.layers) > 1 and self._via_placeable((cx, cy), net_no, conns):
                for layer in self.layers:
                    if layer == cl:
                        continue
                    nxt = (cx, cy, layer)
                    self._relax(
                        current, nxt, base + self.via_cost, goal_cell, g_score, came_from, open_heap
                    )
        return None

    def _relax(self, current, nxt, tentative, goal_cell, g_score, came_from, open_heap):
        if tentative < g_score.get(nxt, math.inf):
            came_from[nxt] = current
            g_score[nxt] = tentative
            heapq.heappush(open_heap, (tentative + self._h(nxt, goal_cell), nxt))

    def _h(self, node: Node, goal_cell: Cell) -> float:
        return math.hypot(node[0] - goal_cell[0], node[1] - goal_cell[1])

    # --- commit / rip -------------------------------------------------------

    def _commit(self, conn: _Connection, path: list[Node]) -> None:
        wire_segs, via_cells = _split_path(path)
        conn.trace_cells = [(cell, layer) for layer, cells in wire_segs for cell in cells]
        conn.via_cells = list(via_cells)
        for cell, layer in conn.trace_cells:
            self._cell_owner[(*cell, layer)] = conn.index
        for cell in conn.via_cells:
            self._via_owner[cell] = conn.index
        conn.wires = []
        for idx, (layer, cells) in enumerate(wire_segs):
            points = [self._cell_center(c) for c in cells]
            if idx == 0:
                points.insert(0, conn.a.location)
            if idx == len(wire_segs) - 1:
                points.append(conn.b.location)
            points = _simplify(points)
            if len(points) >= 2:
                conn.wires.append((layer, points))
        conn.vias = [self._cell_center(c) for c in conn.via_cells]
        conn.routed = True

    def _rip(self, conn: _Connection) -> None:
        for cell, layer in conn.trace_cells:
            if self._cell_owner.get((*cell, layer)) == conn.index:
                del self._cell_owner[(*cell, layer)]
        for cell in conn.via_cells:
            if self._via_owner.get(cell) == conn.index:
                del self._via_owner[cell]
        conn.trace_cells = []
        conn.via_cells = []
        conn.wires = []
        conn.vias = []
        conn.routed = False
        conn.rip_count += 1

    def _crossed_owners(self, path: list[Node], net_no: int, conns: list[_Connection]) -> set[int]:
        owners: set[int] = set()
        for cx, cy, cl in path:
            owner = self._cell_owner.get((cx, cy, cl))
            if owner is not None and conns[owner].net_no != net_no:
                owners.add(owner)
            via_owner = self._via_owner.get((cx, cy))
            if via_owner is not None and conns[via_owner].net_no != net_no:
                owners.add(via_owner)
        return owners

    # --- passes -------------------------------------------------------------

    def route(self) -> RouteResult:
        """Route all ratsnest connections, using rip-up-and-retry across passes."""
        conns = self._build_connections()
        if not conns:
            return RouteResult(self.half_width)

        passes = self.max_passes if self.rip_up else 1
        best_count = -1
        best_snapshot: list[tuple[int, list, list]] | None = None
        best_unrouted = len(conns)
        for _ in range(passes):
            changed = False
            for conn in conns:  # deterministic: connection creation order
                if conn.routed:
                    continue
                if self._try_route(conn, conns):
                    changed = True
            routed_count = sum(1 for c in conns if c.routed)
            if routed_count > best_count:
                best_count = routed_count
                best_unrouted = len(conns) - routed_count
                best_snapshot = self._snapshot(conns)
            if routed_count == len(conns):
                break
            if not changed:
                break
        return self._result_from(best_snapshot or [], best_unrouted)

    def _try_route(self, conn: _Connection, conns: list[_Connection]) -> bool:
        path = self._search(conn, allow_ripup=False, conns=conns)
        if path is None:
            if not self.rip_up:
                return False
            path = self._search(conn, allow_ripup=True, conns=conns)
            if path is None:
                return False
            for idx in sorted(self._crossed_owners(path, conn.net_no, conns)):
                self._rip(conns[idx])
        self._commit(conn, path)
        return True

    def _build_connections(self) -> list[_Connection]:
        pins_by_net = self._pins_by_net()
        conns: list[_Connection] = []
        # deterministic: fewer-pin nets first, then by net number
        for net_no in sorted(pins_by_net, key=lambda n: (len(pins_by_net[n]), n)):
            pins = pins_by_net[net_no]
            for a, b in zip(pins, pins[1:], strict=False):
                conns.append(_Connection(len(conns), net_no, a, b))
        return conns

    def _pins_by_net(self) -> dict[int, list[Pin]]:
        by_net: dict[int, list[Pin]] = {}
        for pin in self.board.get_pins():
            if not any(layer in self.layers for layer in pin.layers):
                continue
            for net_no in pin.net_nos:
                by_net.setdefault(net_no, []).append(pin)
        return by_net

    def _snapshot(self, conns: list[_Connection]) -> list[tuple[int, list, list]]:
        return [
            (c.net_no, [(layer, list(pts)) for layer, pts in c.wires], list(c.vias))
            for c in conns
            if c.routed
        ]

    def _result_from(self, snapshot, unrouted: int) -> RouteResult:
        result = RouteResult(self.half_width)
        result.unrouted = unrouted
        for net_no, wires, vias in snapshot:
            for layer, points in wires:
                result.add_wire(net_no, layer, points)
            for location in vias:
                result.add_via(net_no, location)
        return result


def _split_path(path: list[Node]) -> tuple[list[tuple[int, list[Cell]]], list[Cell]]:
    """Split a node path into per-layer wire segments and via cells."""
    wires: list[tuple[int, list[Cell]]] = []
    vias: list[Cell] = []
    i = 0
    n = len(path)
    while i < n:
        layer = path[i][2]
        cells: list[Cell] = []
        while i < n and path[i][2] == layer:
            cells.append((path[i][0], path[i][1]))
            i += 1
        wires.append((layer, cells))
        if i < n:  # a via connects this segment's last cell to the next layer
            vias.append(cells[-1])
    return wires, vias


def _reconstruct(came_from: dict[Node, Node], current: Node) -> list[Node]:
    path = [current]
    while current in came_from:
        current = came_from[current]
        path.append(current)
    path.reverse()
    return path


def _simplify(points: list[IntPoint]) -> list[IntPoint]:
    """Drop duplicate and collinear points from a polyline."""
    out: list[IntPoint] = []
    for p in points:
        if out and out[-1].x == p.x and out[-1].y == p.y:
            continue
        if len(out) >= 2:
            a, b = out[-2], out[-1]
            if (b.x - a.x) * (p.y - a.y) - (b.y - a.y) * (p.x - a.x) == 0:
                out[-1] = p
                continue
        out.append(p)
    return out


def route_board(
    board: BasicBoard,
    *,
    trace_width: int,
    clearance: int,
    layers: list[int] | None = None,
    via_cost: float = 10.0,
    rip_up: bool = True,
    max_passes: int = 10,
    orthogonal: bool = False,
) -> RouteResult:
    """Convenience wrapper: build a :class:`GridRouter` and route the board."""
    router = GridRouter(
        board,
        trace_width=trace_width,
        clearance=clearance,
        layers=layers,
        via_cost=via_cost,
        rip_up=rip_up,
        max_passes=max_passes,
        orthogonal=orthogonal,
    )
    return router.route()