#!/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 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 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))
    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}

    for pad in fpga.Pads():
        if not pad.GetNetname():
            continue
        position = pad.GetPosition()
        # 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)
        is_power = pad.GetNetname() in POWER_NETS
        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 dogbones")
    print(f"Saved {args.output}")


if __name__ == "__main__":
    main()