#!/usr/bin/env python3
"""Connect SMT power/ground pads locally to their intended solid planes.

This deliberately does not replace power-stage routing.  It only installs
short pad-to-via connections where a filled, same-net plane is physically
present below the candidate via.  High-current converter loops and plane
bridges are handled separately.
"""

from __future__ import annotations

import argparse
import math
import shutil
from pathlib import Path

import pcbnew


ROOT = Path(__file__).resolve().parents[1]
DEFAULT_INPUT = ROOT / "hardware" / "board_prerouted.kicad_pcb"
DEFAULT_OUTPUT = ROOT / "hardware" / "board_plane_stitched.kicad_pcb"
PLANE_NETS = {
    "GND", "SYS5", "VCC0V95", "VCC1V8", "VCC3V3", "VCC3V3_AUX",
    "VCC3V0_CLK", "VCC3V0_DIG", "VCC3V0_RF", "VCC3V3_ANT",
}
VIA_DIAMETER = 0.45
VIA_DRILL = 0.225
CLEARANCE = 0.10


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


def distance_mm(a: pcbnew.VECTOR2I, b: pcbnew.VECTOR2I) -> float:
    return math.hypot(pcbnew.ToMM(a.x - b.x), pcbnew.ToMM(a.y - b.y))


def add_track(board: pcbnew.BOARD, pad: pcbnew.PAD,
              end: pcbnew.VECTOR2I, width: float) -> None:
    track = pcbnew.PCB_TRACK(board)
    track.SetStart(pad.GetPosition())
    track.SetEnd(end)
    track.SetWidth(mm(width))
    track.SetLayer(pad.GetLayer())
    track.SetNet(pad.GetNet())
    track.SetLocked(True)
    board.Add(track)


def add_via(board: pcbnew.BOARD, pad: pcbnew.PAD,
            position: pcbnew.VECTOR2I) -> pcbnew.PCB_VIA:
    via = pcbnew.PCB_VIA(board)
    via.SetPosition(position)
    via.SetWidth(mm(VIA_DIAMETER))
    via.SetDrill(mm(VIA_DRILL))
    via.SetLayerPair(pcbnew.F_Cu, pcbnew.B_Cu)
    via.SetNet(pad.GetNet())
    via.SetLocked(True)
    board.Add(via)
    return via


def plane_at(zones_by_net: dict[str, list[pcbnew.ZONE]], net: str,
             position: pcbnew.VECTOR2I) -> bool:
    return any(zone.HitTestFilledArea(zone.GetLayer(), position)
               for zone in zones_by_net.get(net, ()))


def clear_of_other_copper(position: pcbnew.VECTOR2I, net: str,
                          pads: list[pcbnew.PAD], vias: list[pcbnew.PCB_VIA]) -> bool:
    x, y = pcbnew.ToMM(position.x), pcbnew.ToMM(position.y)
    margin = VIA_DIAMETER / 2 + CLEARANCE + 0.03
    for other in pads:
        if other.GetNetname() == net:
            continue
        box = other.GetBoundingBox()
        x0, y0 = pcbnew.ToMM(box.GetX()) - margin, pcbnew.ToMM(box.GetY()) - margin
        x1, y1 = pcbnew.ToMM(box.GetRight()) + margin, pcbnew.ToMM(box.GetBottom()) + margin
        if x0 <= x <= x1 and y0 <= y <= y1:
            return False
    for other in vias:
        minimum = (VIA_DIAMETER + pcbnew.ToMM(other.GetWidth(pcbnew.F_Cu))) / 2 + CLEARANCE
        if distance_mm(position, other.GetPosition()) < minimum:
            return False
    return True


def clear_track_path(start: pcbnew.VECTOR2I, end: pcbnew.VECTOR2I, net: str,
                     source: pcbnew.PAD, pads: list[pcbnew.PAD], width: float) -> bool:
    """Conservatively reject short dogbones that cross another pad bbox."""
    margin = width / 2 + CLEARANCE + 0.02
    for step in range(1, 12):
        fraction = step / 12
        sample = pcbnew.VECTOR2I(
            round(start.x + (end.x - start.x) * fraction),
            round(start.y + (end.y - start.y) * fraction),
        )
        x, y = pcbnew.ToMM(sample.x), pcbnew.ToMM(sample.y)
        for other in pads:
            if other is source or other.GetNetname() == net:
                continue
            box = other.GetBoundingBox()
            if (pcbnew.ToMM(box.GetX()) - margin <= x <= pcbnew.ToMM(box.GetRight()) + margin and
                    pcbnew.ToMM(box.GetY()) - margin <= y <= pcbnew.ToMM(box.GetBottom()) + margin):
                return False
    return True


def candidates(pad: pcbnew.PAD, footprint: pcbnew.FOOTPRINT) -> list[pcbnew.VECTOR2I]:
    pad_pos = pad.GetPosition()
    fp_pos = footprint.GetPosition()
    away = math.atan2(pad_pos.y - fp_pos.y, pad_pos.x - fp_pos.x)
    if pad_pos == fp_pos:
        away = 0.0
    radius = max(pcbnew.ToMM(pad.GetSizeX()), pcbnew.ToMM(pad.GetSizeY())) / 2
    radius += VIA_DIAMETER / 2 + CLEARANCE + 0.08
    angles = [away]
    for delta_deg in (30, -30, 60, -60, 90, -90, 120, -120, 150, -150, 180):
        angles.append(away + math.radians(delta_deg))
    result = []
    for scale in (1.0, 1.35, 1.75):
        for angle in angles:
            result.append(pcbnew.VECTOR2I(
                pad_pos.x + mm(radius * scale * math.cos(angle)),
                pad_pos.y + mm(radius * scale * math.sin(angle)),
            ))
    return result


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("input", nargs="?", type=Path, default=DEFAULT_INPUT)
    parser.add_argument("output", nargs="?", type=Path, default=DEFAULT_OUTPUT)
    args = parser.parse_args()

    board = pcbnew.LoadBoard(str(args.input))
    if not pcbnew.ZONE_FILLER(board).Fill(board.Zones()):
        raise RuntimeError("Could not fill planes before stitching")

    zones_by_net: dict[str, list[pcbnew.ZONE]] = {}
    for zone in board.Zones():
        if not zone.GetIsRuleArea() and zone.GetNetname() in PLANE_NETS:
            zones_by_net.setdefault(zone.GetNetname(), []).append(zone)

    all_pads = [pad for fp in board.GetFootprints() for pad in fp.Pads()]
    all_vias = [item for item in board.GetTracks() if isinstance(item, pcbnew.PCB_VIA)]
    connected: dict[str, int] = {}
    skipped: dict[str, int] = {}

    for fp in board.GetFootprints():
        if fp.GetReference() == "U1":
            continue
        for pad in fp.Pads():
            net = pad.GetNetname()
            if net not in PLANE_NETS or pad.GetAttribute() != pcbnew.PAD_ATTRIB_SMD:
                continue
            chosen = None
            pad_min = min(pcbnew.ToMM(pad.GetSizeX()), pcbnew.ToMM(pad.GetSizeY()))
            width = max(0.15, min(0.30, pad_min * 0.65))
            for candidate in candidates(pad, fp):
                x, y = pcbnew.ToMM(candidate.x), pcbnew.ToMM(candidate.y)
                if not (20.5 < x < 174.5 and 20.5 < y < 109.5):
                    continue
                if not plane_at(zones_by_net, net, candidate):
                    continue
                if not clear_of_other_copper(candidate, net, all_pads, all_vias):
                    continue
                if not clear_track_path(pad.GetPosition(), candidate, net, pad, all_pads, width):
                    continue
                chosen = candidate
                break
            if chosen is None:
                skipped[net] = skipped.get(net, 0) + 1
                continue
            add_track(board, pad, chosen, width)
            all_vias.append(add_via(board, pad, chosen))
            connected[net] = connected.get(net, 0) + 1

    if not pcbnew.ZONE_FILLER(board).Fill(board.Zones()):
        raise RuntimeError("Could not refill planes after stitching")
    pcbnew.SaveBoard(str(args.output), board)
    for suffix in (".kicad_pro", ".kicad_dru"):
        source = args.input.with_suffix(suffix)
        if source.exists():
            shutil.copyfile(source, args.output.with_suffix(suffix))
    print("Plane connections:", ", ".join(f"{n}={k}" for n, k in sorted(connected.items())))
    print("Pads without a local same-net plane:",
          ", ".join(f"{n}={k}" for n, k in sorted(skipped.items())))
    print(f"Saved {args.output}")


if __name__ == "__main__":
    main()