#!/usr/bin/env python3
"""Install deterministic 0.8-mm-pitch dogbone escapes for connected Ti375 balls."""

from __future__ import annotations

import argparse
import shutil
from pathlib import Path

import pcbnew


ROOT = Path(__file__).resolve().parents[1]
DEFAULT_INPUT = ROOT / "hardware" / "board.kicad_pcb"
DEFAULT_OUTPUT = ROOT / "hardware" / "board_prerouted.kicad_pcb"
POWER_NETS = {
    "GND", "VCC0V95", "VCC1V8", "VCC3V3", "VCC3V0_DIG",
}


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


def centre_mm(position: pcbnew.VECTOR2I) -> tuple[float, float]:
    return pcbnew.ToMM(position.x), pcbnew.ToMM(position.y)


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


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


def ball_grid(footprint: pcbnew.FOOTPRINT) -> tuple[dict, float, float, float, float]:
    """Every ball position, connected or not, plus the array extent."""
    positions = {}
    for pad in footprint.Pads():
        p = pad.GetPosition()
        positions[(round(pcbnew.ToMM(p.x), 2), round(pcbnew.ToMM(p.y), 2))] = pad
    xs = [x for x, _ in positions]
    ys = [y for _, y in positions]
    return positions, min(xs), min(ys), max(xs), max(ys)


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)
    parser.add_argument("--surface-rings", type=int, default=1,
                        help="How many outer ball rings escape on F.Cu without "
                             "a via. Each via removed here frees an inner-layer "
                             "channel and one anti-pad per plane.")
    args = parser.parse_args()

    board = pcbnew.LoadBoard(str(args.input))
    fpga = board.FindFootprintByReference("U1")
    if fpga is None:
        raise RuntimeError("U1 Ti375 footprint is missing")
    centre = fpga.GetPosition()
    occupied: set[tuple[int, int]] = set()
    counts = {"power": 0, "signal": 0, "surface": 0}
    _grid, gx0, gy0, gx1, gy1 = ball_grid(fpga)
    pitch = 0.8

    for pad in fpga.Pads():
        if not pad.GetNetname():
            continue
        position = pad.GetPosition()
        x, y = pcbnew.ToMM(position.x), pcbnew.ToMM(position.y)
        ring = min(round((x - gx0) / pitch), round((gx1 - x) / pitch),
                   round((y - gy0) / pitch), round((gy1 - y) / pitch))
        is_power = pad.GetNetname() in POWER_NETS
        # Signal balls on the outer rings leave straight out on F.Cu.  They
        # need no via at all, and the interstitial site they would have taken
        # stays open as a channel for the balls further in.  Power and ground
        # balls always drop, because their whole purpose is to reach a plane.
        if not is_power and ring < args.surface_rings:
            dx = 0.0 if abs(x - centre_mm(centre)[0]) < 0.4 else (
                1.0 if x > centre_mm(centre)[0] else -1.0)
            dy = 0.0 if abs(y - centre_mm(centre)[1]) < 0.4 else (
                1.0 if y > centre_mm(centre)[1] else -1.0)
            if abs(x - centre_mm(centre)[0]) > abs(y - centre_mm(centre)[1]):
                dy = 0.0
            else:
                dx = 0.0
            reach = 1.6
            end = pcbnew.VECTOR2I(position.x + mm(dx * reach),
                                  position.y + mm(dy * reach))
            add_segment(board, pad, end, 0.10)
            counts["surface"] += 1
            continue
        # Move each ball toward its package quadrant.  Centre-row/column balls
        # use the positive direction, leaving a full 0.8-mm via grid without
        # duplicated interstitial locations.
        dx = 0.4 if position.x >= centre.x else -0.4
        dy = 0.4 if position.y >= centre.y else -0.4
        via_position = pcbnew.VECTOR2I(position.x + mm(dx), position.y + mm(dy))
        key = (via_position.x, via_position.y)
        if key in occupied:
            raise RuntimeError(f"Duplicate dogbone via at {key} for U1.{pad.GetNumber()}")
        occupied.add(key)
        add_segment(board, pad, via_position, 0.15 if is_power else 0.10)
        add_via(board, pad, via_position)
        counts["power" if is_power else "signal"] += 1

    if not pcbnew.ZONE_FILLER(board).Fill(board.Zones()):
        raise RuntimeError("Zone fill failed after BGA fanout")
    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(f"Installed {counts['power']} power/GND and {counts['signal']} signal "
          f"dogbones, {counts['surface']} via-less surface escapes")
    print(f"Saved {args.output}")


if __name__ == "__main__":
    main()