#!/usr/bin/env python3
"""Import a Specctra session and install audited RF/USB copper.

Freerouting is used only for the dense non-critical interconnect.  The GPS
front end and USB2 pair are removed from that result and rebuilt here with
deterministic geometry derived from the selected JLC08161H-3313 stackup.
"""

from __future__ import annotations

import argparse
import math
import shutil
from pathlib import Path

import pcbnew

ROOT = Path(__file__).resolve().parents[1]
PLACED = ROOT / "hardware" / "board.kicad_pcb"
ROUTED = ROOT / "hardware" / "board_routed.kicad_pcb"

RF_NETS = {"RF_SMA", "RF_PRE_SAW", "RF_POST_SAW", "MAX_MIXIN_HI"}
USB_NETS = {"USB_D_P", "USB_D_N", "FTDI_D_P", "FTDI_D_N"}
CRITICAL_NETS = RF_NETS | USB_NETS


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


def point(x: float, y: float) -> pcbnew.VECTOR2I:
    return pcbnew.VECTOR2I_MM(x, y)


def footprint(board: pcbnew.BOARD, ref: str) -> pcbnew.FOOTPRINT:
    item = board.FindFootprintByReference(ref)
    if item is None:
        raise RuntimeError(f"Missing footprint {ref}")
    return item


def pads(board: pcbnew.BOARD, ref: str, number: str) -> list[pcbnew.PAD]:
    result = [pad for pad in footprint(board, ref).Pads() if pad.GetNumber() == str(number)]
    if not result:
        raise RuntimeError(f"Missing pad {ref}.{number}")
    return result


def pad_xy(board: pcbnew.BOARD, ref: str, number: str, index: int = 0) -> tuple[float, float]:
    position = pads(board, ref, number)[index].GetPosition()
    return pcbnew.ToMM(position.x), pcbnew.ToMM(position.y)


def assert_ref_nets(board: pcbnew.BOARD, ref: str, expected: set[str]) -> None:
    actual = {pad.GetNetname() for pad in footprint(board, ref).Pads() if pad.GetNetname()}
    if actual != expected:
        raise RuntimeError(f"{ref} net mismatch: expected {sorted(expected)}, got {sorted(actual)}")


def add_segment(board: pcbnew.BOARD, net_name: str,
                start: tuple[float, float], end: tuple[float, float],
                width: float, layer: int = pcbnew.F_Cu, locked: bool = True) -> None:
    if start == end:
        return
    track = pcbnew.PCB_TRACK(board)
    track.SetStart(point(*start))
    track.SetEnd(point(*end))
    track.SetWidth(mm(width))
    track.SetLayer(layer)
    track.SetNet(board.FindNet(net_name))
    track.SetLocked(locked)
    board.Add(track)


def add_polyline(board: pcbnew.BOARD, net_name: str,
                 vertices: list[tuple[float, float]], width: float,
                 layer: int = pcbnew.F_Cu) -> None:
    for start, end in zip(vertices, vertices[1:]):
        add_segment(board, net_name, start, end, width, layer)


def add_via(board: pcbnew.BOARD, net_name: str, position: tuple[float, float],
            diameter: float = 0.45, drill: float = 0.225) -> None:
    via = pcbnew.PCB_VIA(board)
    via.SetPosition(point(*position))
    via.SetWidth(mm(diameter))
    via.SetDrill(mm(drill))
    via.SetLayerPair(pcbnew.F_Cu, pcbnew.B_Cu)
    via.SetNet(board.FindNet(net_name))
    via.SetLocked(True)
    board.Add(via)


def remove_critical_routes(board: pcbnew.BOARD) -> int:
    doomed = [item for item in board.GetTracks() if item.GetNetname() in CRITICAL_NETS]
    for item in doomed:
        board.Remove(item)
    return len(doomed)


def remove_temporary_keepouts(board: pcbnew.BOARD) -> int:
    doomed = [zone for zone in board.Zones()
              if zone.GetIsRuleArea() and zone.GetZoneName().startswith("TEMP_")]
    for zone in doomed:
        board.Remove(zone)
    return len(doomed)


def route_rf(board: pcbnew.BOARD) -> None:
    # 0.164-mm L1 CPWG with 0.200-mm ground gap: 49.65 ohms on the exact
    # JLC08161H-3313 L1/L2 dielectric construction.
    for ref, expected in {
        "C83": {"RF_SMA", "RF_PRE_SAW"},
        "C84": {"RF_POST_SAW", "GND"},
        "R39": {"RF_POST_SAW", "MAX_MIXIN_HI"},
        "C85": {"MAX_MIXIN_HI", "GND"},
    }.items():
        assert_ref_nets(board, ref, expected)

    w = 0.164
    j2 = pad_xy(board, "J2", "1")
    dc_block_1 = pad_xy(board, "C83", "1")
    l4_2 = pad_xy(board, "L4", "2")
    add_polyline(board, "RF_SMA", [j2, l4_2, dc_block_1], w)
    d1 = pad_xy(board, "D1", "1")
    add_polyline(board, "RF_SMA", [d1, (d1[0], 45.07)], w)

    add_polyline(board, "RF_PRE_SAW",
                 [pad_xy(board, "C83", "2"), pad_xy(board, "FL1", "1")], w)

    fl_out = pad_xy(board, "FL1", "4")
    tune1 = pad_xy(board, "C84", "1")
    series_1 = pad_xy(board, "R39", "1")
    add_polyline(board, "RF_POST_SAW", [fl_out, tune1, series_1], w)

    series_2 = pad_xy(board, "R39", "2")
    mixin = pad_xy(board, "U21", "25")
    add_polyline(board, "MAX_MIXIN_HI",
                 [series_2, (159.69, 41.80), (mixin[0], 41.80), mixin], w)
    tune2 = pad_xy(board, "C85", "1")
    add_polyline(board, "MAX_MIXIN_HI", [tune2, (tune2[0], 41.80)], w)


def route_usb(board: pcbnew.BOARD) -> None:
    # Main device-side pair uses the upper receptacle contact pair.  The two
    # mirrored USB-C contacts join through a sub-1.5-mm L3 crossover so neither
    # polarity creates a same-layer crossing at the receptacle.
    w = 0.135
    dm_main = pad_xy(board, "J1", "B7")
    dp_main = pad_xy(board, "J1", "A6")
    dm_mirror = pad_xy(board, "J1", "A7")
    dp_mirror = pad_xy(board, "J1", "B6")
    dm_join, dp_join = (28.30, 47.25), (28.80, 47.75)
    dm_aux, dp_aux = (27.30, 48.25), (27.80, 48.75)

    add_polyline(board, "USB_D_N", [dm_main, dm_join], w)
    add_polyline(board, "USB_D_P", [dp_main, dp_join], w)
    add_polyline(board, "USB_D_N", [dm_mirror, dm_aux], w)
    add_polyline(board, "USB_D_P", [dp_mirror, dp_aux], w)
    for net, location in (("USB_D_N", dm_join), ("USB_D_P", dp_join),
                          ("USB_D_N", dm_aux), ("USB_D_P", dp_aux)):
        add_via(board, net, location)
    add_polyline(board, "USB_D_N", [dm_aux, dm_join], w, pcbnew.In2_Cu)
    add_polyline(board, "USB_D_P", [dp_aux, dp_join], w, pcbnew.In2_Cu)

    dm_esd = pad_xy(board, "U6", "3")
    dp_esd = pad_xy(board, "U6", "1")
    add_polyline(board, "USB_D_N", [dm_join, (30.00, 49.00), (30.00, 50.70), dm_esd], w)
    add_polyline(board, "USB_D_P", [dp_join, (30.285, 49.285), (30.285, 51.64), dp_esd], w)

    dm_r = pad_xy(board, "R64", "1")
    dp_r = pad_xy(board, "R63", "1")
    add_polyline(board, "USB_D_N",
                 [dm_esd, (31.00, 51.40), (34.00, 54.40), (35.80, 54.40), dm_r], w)
    add_polyline(board, "USB_D_P",
                 [dp_esd, (30.715, 51.685), (33.715, 54.685), (35.80, 54.685), dp_r], w)

    # Equal 1.511-mm post-termination legs into the FT4232H USB PHY.
    add_polyline(board, "FTDI_D_N", [pad_xy(board, "R64", "2"), pad_xy(board, "U2", "7")], w)
    add_polyline(board, "FTDI_D_P", [pad_xy(board, "R63", "2"), pad_xy(board, "U2", "8")], w)


def add_rf_ground(board: pcbnew.BOARD) -> None:
    zone = pcbnew.ZONE(board)
    zone.SetLayer(pcbnew.F_Cu)
    zone.SetNet(board.FindNet("GND"))
    zone.SetLocalClearance(mm(0.20))
    zone.SetMinThickness(mm(0.10))
    zone.SetPadConnection(pcbnew.ZONE_CONNECTION_FULL)
    zone.SetIslandRemovalMode(pcbnew.ISLAND_REMOVAL_MODE_NEVER)
    zone.SetAssignedPriority(100)
    outline = zone.Outline()
    outline.NewOutline()
    for x, y in ((152.0, 36.0), (174.5, 36.0), (174.5, 49.5), (152.0, 49.5)):
        outline.Append(point(x, y))
    board.Add(zone)

    # Lambda/20 is about 9.5 mm at GPS L1 in FR-4.  This 2-mm fence is much
    # tighter and also gives every local shunt/EP a short plane connection.
    fence = [(x, 38.7) for x in range(153, 174, 2)]
    fence += [(x, 45.0) for x in range(161, 174, 2)]
    for location in fence:
        add_via(board, "GND", location, 0.45, 0.225)


def routed_lengths(board: pcbnew.BOARD) -> dict[str, float]:
    totals = {net: 0.0 for net in CRITICAL_NETS}
    for item in board.GetTracks():
        net = item.GetNetname()
        if net not in totals or isinstance(item, pcbnew.PCB_VIA):
            continue
        start, end = item.GetStart(), item.GetEnd()
        totals[net] += math.hypot(pcbnew.ToMM(end.x - start.x), pcbnew.ToMM(end.y - start.y))
    return totals


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("session", type=Path, help="Freerouting .ses file")
    parser.add_argument("--output", type=Path, default=ROUTED)
    args = parser.parse_args()

    board = pcbnew.LoadBoard(str(PLACED))
    if not pcbnew.ImportSpecctraSES(board, str(args.session)):
        raise RuntimeError(f"Could not import {args.session}")
    removed = remove_critical_routes(board)
    removed_keepouts = remove_temporary_keepouts(board)
    route_rf(board)
    route_usb(board)
    add_rf_ground(board)
    if not pcbnew.ZONE_FILLER(board).Fill(board.Zones()):
        raise RuntimeError("Final zone fill failed")
    pcbnew.SaveBoard(str(args.output), board)
    for suffix in (".kicad_pro", ".kicad_dru"):
        source = PLACED.with_suffix(suffix)
        if source.exists():
            shutil.copyfile(source, args.output.with_suffix(suffix))
    print(f"Removed {removed} autorouted critical-net items")
    print(f"Removed {removed_keepouts} temporary critical-route keepouts")
    print(f"Saved {args.output}")
    for net, length in sorted(routed_lengths(board).items()):
        print(f"{net:16s} {length:8.3f} mm")


if __name__ == "__main__":
    main()