#!/Applications/KiCad/KiCad.app/Contents/Frameworks/Python.framework/Versions/3.9/bin/python3
from pathlib import Path
import json
import math
import shutil

import wx
app = wx.App(False)
import pcbnew

root = Path(__file__).resolve().parent
board_path = root / "lamp_buck_lmr51450_a1.kicad_pcb"
project_path = root / "lamp_buck_lmr51450_a1.kicad_pro"
backup_path = root / "lamp_buck_lmr51450_a1_pre_refine.kicad_pcb"
if not backup_path.exists():
    shutil.copy2(board_path, backup_path)
else:
    shutil.copy2(backup_path, board_path)
board = pcbnew.LoadBoard(str(board_path))
MM = pcbnew.FromMM


def vec(x, y):
    return pcbnew.VECTOR2I(MM(x), MM(y))


def mmpt(point):
    return (pcbnew.ToMM(point.x), pcbnew.ToMM(point.y))


def close_point(a, b, tol=0.025):
    return math.dist(a, b) <= tol


def same_segment(track, a, b):
    start, end = mmpt(track.GetStart()), mmpt(track.GetEnd())
    return (close_point(start, a) and close_point(end, b)) or (close_point(start, b) and close_point(end, a))


def find_pad(ref, number):
    for footprint in board.GetFootprints():
        footprint = pcbnew.Cast_to_FOOTPRINT(footprint)
        if footprint.GetReference() == ref:
            for pad in footprint.Pads():
                if pad.GetNumber() == str(number):
                    return pad
    raise KeyError(f"missing {ref}.{number}")


# Widen the high-current 5-V path from the buck inductor through the output
# capacitor bank. Keep the existing narrow load-switch and FB branches, which
# match the baseline board topology and must be requalified with load testing.
power_segments = [
    ((60.825, 32.0), (61.83, 33.0)),
    ((61.83, 33.0), (66.92, 33.0)),
    ((66.92, 33.0), (72.025, 33.0)),
    ((72.025, 18.0), (72.025, 23.0)),
    ((72.025, 23.0), (72.025, 28.0)),
    ((72.025, 28.0), (72.025, 33.0)),
    ((77.99, 34.85), (77.99, 38.88)),
    ((77.99, 38.88), (82.61, 43.5)),
]
widened = 0
for track in list(board.GetTracks()):
    if track.GetNetname() == "SYS_5V0_ALWAYS_ON" and not isinstance(track, pcbnew.PCB_VIA):
        if any(same_segment(track, a, b) for a, b in power_segments):
            track.SetWidth(MM(1.2))
            widened += 1

# This low-current enable signal uses 0.15-mm traces in the congested QFN escape.
# The board is explicitly set to a 0.15-mm minimum, while power routes above are
# locally widened to 1.2 mm.
project = json.loads(project_path.read_text(encoding="utf-8"))
project["board"]["design_settings"]["rules"]["min_track_width"] = 0.15
project_path.write_text(json.dumps(project, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")

# Complete the one unrouted LED_QOD link by routing below the 0603 resistor body.
qod_a = find_pad("U3", 5)
qod_b = find_pad("R13", 2)
gnd_pad = find_pad("U3", 4)
qod_points = [mmpt(qod_a.GetPosition()), (85.0, 44.15), (85.0, 46.0),
              (88.0, 46.0), (88.0, 44.0), mmpt(qod_b.GetPosition())]
for a, b in zip(qod_points, qod_points[1:]):
    track = pcbnew.PCB_TRACK(board)
    track.SetStart(vec(*a))
    track.SetEnd(vec(*b))
    track.SetWidth(MM(0.2))
    track.SetLayer(pcbnew.F_Cu)
    track.SetNetCode(qod_a.GetNetCode())
    board.Add(track)

# Delete three copper stubs left by removed DNP/relocated parts.
dangling_anchors = ((58.775, 40.0), (83.8017, 44.175), (33.3147, 33.2103),
                    (57.6563, 41.1187), (35.7063, 41.1187),
                    (84.0, 44.0146), (85.8826, 45.8972),
                    (88.8778, 45.8972), (84.0, 43.55))
removed_dangling = 0
for track in list(board.GetTracks()):
    if track.GetNetname() != "GND" or isinstance(track, pcbnew.PCB_VIA):
        continue
    endpoints = (mmpt(track.GetStart()), mmpt(track.GetEnd()))
    if any(close_point(endpoint, anchor, 0.03) for endpoint in endpoints for anchor in dangling_anchors):
        board.Remove(track)
        removed_dangling += 1

# Replace the ground branch below U3 with a short local ground connection and
# a stitch via into the bottom plane, leaving a clear escape for the QOD signal.
gnd_via = pcbnew.PCB_VIA(board)
gnd_via.SetPosition(vec(84.0, 44.4))
gnd_via.SetViaType(pcbnew.VIATYPE_THROUGH)
gnd_via.SetWidth(MM(0.6))
gnd_via.SetDrill(MM(0.3))
gnd_via.SetLayerPair(pcbnew.F_Cu, pcbnew.B_Cu)
gnd_via.SetFrontTentingMode(pcbnew.TENTING_MODE_TENTED)
gnd_via.SetBackTentingMode(pcbnew.TENTING_MODE_TENTED)
gnd_via.SetNetCode(gnd_pad.GetNetCode())
board.Add(gnd_via)
gnd_link = pcbnew.PCB_TRACK(board)
gnd_link.SetStart(gnd_pad.GetPosition())
gnd_link.SetEnd(gnd_via.GetPosition())
gnd_link.SetWidth(MM(0.2))
gnd_link.SetLayer(pcbnew.F_Cu)
gnd_link.SetNetCode(gnd_pad.GetNetCode())
board.Add(gnd_link)

pcbnew.ZONE_FILLER(board).Fill(board.Zones())
if not pcbnew.SaveBoard(str(board_path), board):
    raise RuntimeError("failed to save refined board")
print(json.dumps({
    "status": "ROUTED_BOARD_REFINEMENT_SAVED_DRC_PENDING",
    "high_current_segments_widened_to_1p2mm": widened,
    "manual_qod_segments_added": len(qod_points) - 1,
    "dangling_ground_tracks_removed": removed_dangling,
    "u3_ground_stitch_via_added": True,
    "board": str(board_path),
}, indent=2))
