"""Exact-geometry routing track (orthogonal), built on the search tree.

The grid MVP (:mod:`freeroute.route.grid_router`) approximates clearance with
cell occupancy and emits centrelines with no width-aware DRC check. This track
keeps the grid's rip-up/multi-layer/via *topology* search (in orthogonal mode,
so segments are axis-aligned) but replaces the occupancy check with an **exact**
one: every item's copper is an exact :class:`~freeroute.geometry.IntBox` tile
indexed in a :class:`~freeroute.board.search_tree.ShapeSearchTree`, and the
routed board is verified to have **no clearance violation** by exact tile
intersection. A net whose copper would violate is dropped, so the output is
DRC-clean by construction.

**Shove (opt-in, ``shove=True``).** Before a dropped net is abandoned, the
router tries to make room by *moving an existing trace aside* geometrically
(:func:`~freeroute.geometry.shove_segment`) rather than dropping the incoming
net. A shove is accepted only if the moved trace still connects its pads, clears
every other item exactly (via the tree), and stays in the board outline — so
cleanliness is preserved. Shoves are bounded (``max_shove_depth`` and a
per-trace cap) and deterministic.

This is an MVP shove for orthogonal traces: it recovers a straight-line net by
shoving the traces that cross it perpendicular. On this grid-based track its
extra reach over rip-up is limited (rip-up already reroutes); the larger density
win needs a continuous expansion-room router (the next gap). Also deferred:
45-degree / ``IntOctagon`` trace caps.
"""

from __future__ import annotations

from dataclasses import dataclass, field
import math

from freeroute.board.board import BasicBoard
from freeroute.board.items import Pin
from freeroute.board.search_tree import ShapeSearchTree, TreeShape
from freeroute.geometry import (
    IntBox,
    IntPoint,
    Polyline,
    PolylineShape,
    segment_box,
    segment_octagon,
    shove_segment,
)
from freeroute.geometry.polyline import corner_box

from .grid_router import RouteResult, route_board

__all__ = ["ExactRouteResult", "route_board_exact"]


@dataclass
class _Accepted:
    """A committed net's per-net geometry and its tile owners in the tree."""

    wires: list[tuple[int, list[IntPoint]]] = field(default_factory=list)
    vias: list[IntPoint] = field(default_factory=list)
    owners: list[int] = field(default_factory=list)
    shove_count: int = 0


@dataclass
class ExactRouteResult:
    """The routed geometry plus the exact spatial index it was verified against."""

    result: RouteResult
    tree: ShapeSearchTree
    #: ``True`` if the routed board has no different-net clearance violation
    drc_clean: bool
    #: number of successful shoves performed
    shoves: int = 0

    @property
    def routed_net_numbers(self):
        return self.result.routed_net_numbers

    def via_count(self) -> int:
        return self.result.via_count()


def route_board_exact(
    board: BasicBoard,
    *,
    trace_width: int,
    clearance: int,
    layers: list[int] | None = None,
    via_cost: float = 10.0,
    max_passes: int = 10,
    rip_up: bool = True,
    shove: bool = False,
    diagonal: bool = False,
    max_shove_depth: int = 5,
    shove_cap: int = 3,
    via_half: int | None = None,
    edge_clearance: int | None = None,
    resolution: int = 1,
    heuristic_weight: float = 1.0,
    seed: int | None = None,
    route_first: frozenset | None = None,
    access_cost: float = 0.0,
    rollback: bool = True,
) -> ExactRouteResult:
    """Route ``board`` orthogonally and verify the result with exact geometry.

    ``diagonal`` (opt-in) adds a 45-degree recovery pass: a still-dropped 2-pin net
    is retried as an octilinear (0/45/90/135-degree) trace whose diagonal copper is
    covered by an exact integer octagon, so it fits diagonal corridors the
    orthogonal L cannot. With ``diagonal=False`` the pass is skipped and the output
    is byte-for-byte the orthogonal router."""
    signal_layers = layers or [
        i for i, layer in enumerate(board.layer_structure.arr) if layer.is_signal
    ]
    half_width = max(trace_width // 2, 1)
    via_half = via_half or half_width
    step = trace_width + clearance

    result = route_board(
        board,
        trace_width=trace_width,
        clearance=clearance,
        layers=signal_layers,
        via_cost=via_cost,
        max_passes=max_passes,
        rip_up=rip_up,
        orthogonal=True,
        via_half=via_half,
        edge_clearance=edge_clearance,
        resolution=resolution,
        heuristic_weight=heuristic_weight,
        seed=seed,
        route_first=route_first,
        access_cost=access_cost,
        rollback=rollback,
    )
    edge = 0 if edge_clearance is None else edge_clearance
    outline = board.outline

    tree = ShapeSearchTree()
    owner = _counter()

    # static items: pad and keepout copper (routed=False)
    for pin in board.get_pins():
        if pin.shape is None:
            continue
        box = pin.shape.bounding_box()
        if box.is_empty():
            continue
        net = pin.net_nos[0] if pin.net_nos else -1
        pin_layers = frozenset(layer for layer in pin.layers if layer in signal_layers)
        if pin_layers:
            tree.insert(TreeShape(next(owner), net, pin_layers, box))
    for obstacle in board.get_obstacle_areas():
        if obstacle.layer not in signal_layers:
            continue
        blocks = (
            "all" if obstacle.blocks_wires and obstacle.blocks_vias
            else "wires" if obstacle.blocks_wires else "vias"
        )
        for tile in obstacle.tiles:
            box = tile.bounding_box()
            if not box.is_empty():
                tree.insert(
                    TreeShape(
                        next(owner), -1, frozenset({obstacle.layer}), box,
                        blocks=blocks, clearance=obstacle.clearance,
                    )
                )

    via_layers = frozenset(signal_layers)
    accepted: dict[int, _Accepted] = {}

    # Verify connection by connection, so one conflicting segment of a
    # 25-pin supply net costs that one connection, not the whole net.
    rejected: list[tuple[int, str, str]] = []
    ends_list = result.connection_ends or [("?", "?", "")] * len(result.connections)
    for (net_no, wires, vias), ends in sorted(
        zip(result.connections, ends_list, strict=False), key=lambda c: c[0][0]
    ):
        boxes = _copper_boxes(wires, vias, half_width, via_layers, via_half)
        wire_count = len(_copper_boxes(wires, [], half_width, via_layers, via_half))
        if any(
            tree.clearance_conflict(
                box, net_no, layer, clearance, "wire" if i < wire_count else "via"
            )
            for i, (box, box_layers) in enumerate(boxes)
            for layer in box_layers
        ) or _end_in_rule_area(tree, wires, half_width) or (
            outline is not None and not all(outline.box_clear(b, edge) for b, _ in boxes)
        ):
            rejected.append((net_no, *ends))
            continue  # exact-clearance conflict -> drop (may be recovered by shove)
        acc = accepted.get(net_no)
        if acc is None:
            acc = accepted[net_no] = _Accepted(wires=[], vias=[])
        acc.wires.extend(wires)
        acc.vias.extend(vias)
        for box, box_layers in boxes:
            oid = next(owner)
            tree.insert(TreeShape(oid, net_no, box_layers, box, routed=True))
            acc.owners.append(oid)

    accepted_before_recovery = set(accepted)
    # a net is dropped if it has a real ratsnest (>= 2 pins on a signal layer)
    # but was not accepted — whether the grid failed to route it or its exact
    # copper conflicted.
    dropped = sorted(_routable_nets(board, signal_layers) - set(accepted))

    shoves = 0
    if shove and dropped:
        outline = board.bounding_box
        for net_no in dropped:
            done = _shove_recover(
                net_no,
                board,
                tree,
                accepted,
                owner,
                signal_layers,
                half_width,
                clearance,
                step,
                outline,
                max_shove_depth,
                shove_cap,
            )
            if done:
                shoves += done

    if diagonal:
        outline = board.bounding_box
        for net_no in sorted(_routable_nets(board, signal_layers)):
            _diagonalize(
                net_no, board, tree, accepted, owner, signal_layers, half_width, clearance, outline
            )

    clean = RouteResult(half_width)
    # what is still missing: what the grid could not route, plus what the
    # exact check rejected, less any net shove or the diagonal pass recovered
    clean.failed = [
        *result.failed,
        *(r for r in rejected if not (r[0] in accepted and r[0] not in accepted_before_recovery)),
    ]
    for net_no, acc in accepted.items():
        clean.wires[net_no] = acc.wires
        clean.vias[net_no] = acc.vias

    drc_clean = tree.has_violation(clearance) is None
    return ExactRouteResult(result=clean, tree=tree, drc_clean=drc_clean, shoves=shoves)


# --- 45-degree (diagonal) recovery -------------------------------------------


def _octilinear_candidates(a: IntPoint, b: IntPoint) -> list[list[IntPoint]]:
    """Octilinear (0/45/90/135-degree) polylines from ``a`` to ``b``.

    A direct diagonal when the offset is exactly 45-degree, otherwise the two
    two-segment paths that spend ``min(|dx|, |dy|)`` on a 45-degree run and the
    rest on an axis run (diagonal-first and diagonal-last)."""
    dx, dy = b.x - a.x, b.y - a.y
    sx = (dx > 0) - (dx < 0)
    sy = (dy > 0) - (dy < 0)
    adx, ady = abs(dx), abs(dy)
    if adx == ady:
        raw = [[a, b]]
    else:
        m = min(adx, ady)
        raw = [
            [a, IntPoint(a.x + sx * m, a.y + sy * m), b],  # diagonal first
            [a, IntPoint(b.x - sx * m, b.y - sy * m), b],  # diagonal last
        ]
    out: list[list[IntPoint]] = []
    for corners in raw:
        deduped: list[IntPoint] = []
        for p in corners:
            if deduped and deduped[-1].x == p.x and deduped[-1].y == p.y:
                continue
            deduped.append(p)
        if len(deduped) >= 2:
            out.append(deduped)
    return out


def _candidate_tiles(corners, net_no, layer, tree, half_width, clearance, outline):
    """``(length, tiles)`` if this octilinear candidate is exactly clean, else
    ``None``. ``tiles`` are ``(box, exact, seg)`` triples ready to insert: a plain
    box for an orthogonal segment or corner, a bounding box + octagon copper + seg
    for a diagonal segment."""
    tiles: list[tuple[IntBox, object, object]] = []
    length = 0.0
    for p, q in zip(corners, corners[1:], strict=False):
        dx, dy = q.x - p.x, q.y - p.y
        length += math.hypot(dx, dy)
        if dx == 0 or dy == 0:  # orthogonal
            box = segment_box(p, q, half_width)
            if not outline.contains_box(box):
                return None
            if tree.clearance_conflict(box, net_no, layer, clearance):
                return None
            tiles.append((box, None, None))
        elif abs(dx) == abs(dy):  # 45-degree diagonal
            copper = segment_octagon(p, q, half_width)
            if not outline.contains_box(copper.bounding_box()):
                return None
            grown = segment_octagon(p, q, half_width + clearance)
            if tree.clearance_conflict_shape(grown, net_no, layer):
                return None
            tiles.append((copper.bounding_box(), copper, (p, q, half_width)))
        else:  # not octilinear -- reject
            return None
    for corner in corners[1:-1]:
        square = IntBox(
            corner.x - half_width,
            corner.y - half_width,
            corner.x + half_width,
            corner.y + half_width,
        )
        if not outline.contains_box(square):
            return None
        if tree.clearance_conflict(square, net_no, layer, clearance):
            return None
        tiles.append((square, None, None))
    return length, tiles


def _wire_length(wires) -> float:
    total = 0.0
    for _layer, corners in wires:
        for p, q in zip(corners, corners[1:], strict=False):
            total += math.hypot(q.x - p.x, q.y - p.y)
    return total


def _diagonalize(
    net_no, board, tree, accepted, owner, signal_layers, half_width, clearance, outline
) -> bool:
    """Route (or shorten) a 2-pin net with the shortest exactly-clean octilinear
    trace. Recovers a dropped net; replaces an orthogonal route only when the
    diagonal is strictly shorter. Endpoints stay on the pads; diagonal copper is
    stored as an exact octagon so later nets clear it precisely. Same-net copper is
    ignored by the clearance check, so the candidate is verified against every
    *other* net before the old copper is removed. Returns ``True`` on a change."""
    pins = _net_pins(board, net_no, signal_layers)
    if len(pins) != 2:
        return False
    a, b = pins
    layer = next(iter(set(a.layers) & set(b.layers) & set(signal_layers)), None)
    if layer is None:
        return False

    acc = accepted.get(net_no)
    if acc is not None and acc.vias:
        return False  # keep a multi-layer orthogonal route; diagonal MVP is single-layer

    best = None
    for corners in _octilinear_candidates(a.location, b.location):
        found = _candidate_tiles(corners, net_no, layer, tree, half_width, clearance, outline)
        if found is None:
            continue
        length, tiles = found
        if best is None or length < best[0]:
            best = (length, corners, tiles)
    if best is None:
        return False
    length, corners, tiles = best

    if acc is not None:
        if length >= _wire_length(acc.wires):
            return False  # diagonal is not shorter -- keep the orthogonal route
        for oid in acc.owners:
            tree.remove_owner(oid)  # rip the orthogonal copper; diagonal already clean of others

    new = _Accepted(wires=[(layer, corners)], vias=[])
    for box, exact, seg in tiles:
        oid = next(owner)
        tree.insert(
            TreeShape(oid, net_no, frozenset({layer}), box, routed=True, exact=exact, seg=seg)
        )
        new.owners.append(oid)
    accepted[net_no] = new
    return True


# --- shove recovery ----------------------------------------------------------


def _shove_recover(
    net_no,
    board,
    tree,
    accepted,
    owner,
    signal_layers,
    half_width,
    clearance,
    step,
    outline,
    max_shove_depth,
    shove_cap,
) -> int:
    """Try to route dropped net ``net_no`` by shoving crossing traces aside.

    MVP: a straight orthogonal candidate for a 2-pin net, shoving each single-
    segment trace that crosses it perpendicular until it clears. Returns the
    number of successful shoves (0 if the net stays dropped).
    """
    pins = _net_pins(board, net_no, signal_layers)
    if len(pins) != 2:
        return 0
    a, b = pins
    if a.location.x != b.location.x and a.location.y != b.location.y:
        return 0  # MVP only recovers straight (axis-aligned) nets
    layer = next(iter(set(a.layers) & set(b.layers) & set(signal_layers)), None)
    if layer is None:
        return 0
    corners = [a.location, b.location]
    d_boxes = PolylineShape(Polyline(corners), half_width).tiles()

    # which accepted traces cross this net's straight copper?
    blockers: dict[int, list[int]] = {}
    for box in d_boxes:
        for shape in tree.overlapping(box.offset(clearance)):
            if not shape.routed or shape.net_no == net_no or layer not in shape.layers:
                continue
            if shape.tile.overlaps(box.offset(clearance)):
                blockers.setdefault(shape.net_no, []).append(shape.owner)

    if len(blockers) > max_shove_depth:
        return 0

    shoved: list[int] = []
    for blocker_net in sorted(blockers):
        acc = accepted.get(blocker_net)
        if acc is None or acc.shove_count >= shove_cap:
            return _rollback(shoved, accepted, tree, owner)
        if not _try_shove_trace(
            blocker_net,
            corners,
            layer,
            board,
            tree,
            accepted,
            owner,
            signal_layers,
            half_width,
            clearance,
            step,
            outline,
        ):
            return _rollback(shoved, accepted, tree, owner)
        shoved.append(blocker_net)

    # re-verify the dropped net is now clean, then accept it
    if any(tree.clearance_conflict(box, net_no, layer, clearance) for box in d_boxes):
        return _rollback(shoved, accepted, tree, owner)
    acc = _Accepted(wires=[(layer, corners)], vias=[])
    for box in d_boxes:
        oid = next(owner)
        tree.insert(TreeShape(oid, net_no, frozenset({layer}), box, routed=True))
        acc.owners.append(oid)
    accepted[net_no] = acc
    return len(shoved) + 1


def _try_shove_trace(
    blocker_net,
    d_corners,
    layer,
    board,
    tree,
    accepted,
    owner,
    signal_layers,
    half_width,
    clearance,
    step,
    outline,
) -> bool:
    """Shove ``blocker_net``'s trace clear of the straight net at ``d_corners``.

    Only the single-segment straight-crossing case is handled (MVP). Tries
    increasing perpendicular displacements away from the crossing net until the
    moved trace clears everything exactly.
    """
    acc = accepted[blocker_net]
    if len(acc.wires) != 1:
        return False
    seg_layer, corners = acc.wires[0]
    if seg_layer != layer:
        return False

    # find the blocker's axis-aligned segment (run) that conflicts with the
    # crossing net, and the perpendicular displacements to clear it
    d_boxes = PolylineShape(Polyline(d_corners), half_width).tiles()
    seg_index, horizontal = _conflicting_segment(corners, d_boxes, half_width, clearance)
    if seg_index is None:
        return False

    d0, d1 = d_corners[0], d_corners[1]
    margin = 2 * half_width + clearance
    fine = half_width + clearance
    ref = corners[seg_index]
    displacements: list[tuple[int, int]] = []
    if horizontal:  # displace in y; clear the crossing net's y-span
        lo, hi = min(d0.y, d1.y), max(d0.y, d1.y)
        for extra in range(6):
            displacements.append((0, (hi + margin) - ref.y + extra * fine))
            displacements.append((0, (lo - margin) - ref.y - extra * fine))
    else:  # displace in x
        lo, hi = min(d0.x, d1.x), max(d0.x, d1.x)
        for extra in range(6):
            displacements.append(((hi + margin) - ref.x + extra * fine, 0))
            displacements.append(((lo - margin) - ref.x - extra * fine, 0))

    # remove the blocker's current tiles so its shoved copy can be verified alone
    for oid in acc.owners:
        tree.remove_owner(oid)

    for dx, dy in displacements:
        if dx == 0 and dy == 0:
            continue
        new_corners = shove_segment(list(corners), seg_index, dx, dy)
        if _place_trace(
            blocker_net,
            seg_layer,
            new_corners,
            tree,
            acc,
            owner,
            half_width,
            clearance,
            outline,
        ):
            acc.wires = [(seg_layer, new_corners)]
            acc.shove_count += 1
            return True

    # could not shove cleanly: restore the original tiles
    _place_trace(
        blocker_net,
        seg_layer,
        corners,
        tree,
        acc,
        owner,
        half_width,
        clearance,
        outline,
        force=True,
    )
    return False


def _conflicting_segment(corners, d_boxes, half_width, clearance):
    """Index + orientation of the axis-aligned segment of ``corners`` that
    conflicts with the crossing net's copper ``d_boxes``, or ``(None, False)``."""
    for k in range(len(corners) - 1):
        a, b = corners[k], corners[k + 1]
        horizontal = a.y == b.y
        if not horizontal and a.x != b.x:
            continue  # skip diagonal grid stubs
        from freeroute.geometry import segment_box

        seg = segment_box(a, b, half_width).offset(clearance)
        if any(seg.overlaps(box) for box in d_boxes):
            return k, horizontal
    return None, False


def _place_trace(
    net_no, layer, corners, tree, acc, owner, half_width, clearance, outline, force=False
) -> bool:
    """Verify + insert a trace's copper; ``force`` skips the DRC/outline check."""
    boxes = PolylineShape(Polyline(corners), half_width).tiles()
    if not force:
        for box in boxes:
            if not outline.contains_box(box):
                return False
            if tree.clearance_conflict(box, net_no, layer, clearance):
                return False
    acc.owners = []
    for box in boxes:
        oid = next(owner)
        tree.insert(TreeShape(oid, net_no, frozenset({layer}), box, routed=True))
        acc.owners.append(oid)
    return True


def _rollback(shoved, accepted, tree, owner) -> int:  # noqa: ARG001
    # shoved traces were verified individually; nothing else to undo for the MVP
    return 0


# --- helpers -----------------------------------------------------------------


def _net_boxes(result, net_no, half_width, via_layers, via_half=None):
    return _copper_boxes(
        result.wires.get(net_no, []), result.vias.get(net_no, []), half_width, via_layers, via_half
    )


def _end_in_rule_area(tree, wires, half_width) -> bool:
    """A wire's tiles stop at its end points, but a track ends in a round cap
    half its width beyond them. Against other copper the pad or via there
    covers it; a rule area has no clearance to absorb it, so check it."""
    for layer, points in wires:
        for end in (points[0], points[-1]):
            if tree.rule_area_overlap(corner_box(end, half_width), layer):
                return True
    return False


def _copper_boxes(wires, vias, half_width, via_layers, via_half=None):
    via_half = via_half or half_width
    boxes: list[tuple[IntBox, frozenset[int]]] = []
    for layer, points in wires:
        for box in PolylineShape(Polyline(points), half_width).tiles():
            boxes.append((box, frozenset({layer})))
    for loc in vias:
        vbox = IntBox(loc.x - via_half, loc.y - via_half, loc.x + via_half, loc.y + via_half)
        boxes.append((vbox, via_layers))
    return boxes


def _net_pins(board: BasicBoard, net_no: int, signal_layers) -> list[Pin]:
    return [
        pin
        for pin in board.get_pins()
        if net_no in pin.net_nos and any(layer in signal_layers for layer in pin.layers)
    ]


def _routable_nets(board: BasicBoard, signal_layers) -> set[int]:
    """Net numbers with at least two pins on a signal layer (a real ratsnest)."""
    counts: dict[int, int] = {}
    for pin in board.get_pins():
        if not any(layer in signal_layers for layer in pin.layers):
            continue
        for net_no in pin.net_nos:
            counts[net_no] = counts.get(net_no, 0) + 1
    return {net_no for net_no, c in counts.items() if c >= 2}


def _counter():
    i = 0
    while True:
        yield i
        i += 1