"""Autoroute the PSU carrier board with Freerouting.

Run with KiCad's bundled Python (KiCad must not have the board open):
    "C:/Users/patri/AppData/Local/Programs/KiCad/10.0/bin/python.exe" hardware/kicad/route_pcb.py

Exports a Specctra DSN, adds class-to-class clearance (HV <-> everything else 3 mm, matching
plasma_dmd_psu.kicad_dru, which Freerouting does not read), routes, and imports the session.
Existing tracks are removed first, so the result always comes from the current placement.
"""
import re
import subprocess
import sys
from pathlib import Path

import pcbnew

HERE = Path(__file__).parent
ROOT = HERE.parents[1]
BOARD = HERE / "plasma_dmd_psu.kicad_pcb"
WORK = HERE / "route"
JAVA = Path("C:/Program Files/Eclipse Adoptium/jre-25.0.4.101-hotspot/bin/java.exe")
JAR = ROOT / "tools/freerouting/freerouting-2.5.0.jar"
PASSES = int(sys.argv[1]) if len(sys.argv) > 1 else 100

# Route with oversized tracks, then shrink them to the real width: every track then keeps
# (routing clearance + half the width difference) from its neighbours. This gets 3 mm HV <-> LV
# and 1.3 mm HV <-> HV while tracks can still squeeze out between transistor pins.
#            class: (routing width um, routing clearance um, final width mm)
ROUTE = {"kicad_default": (1800, 250, 0.4), "Power": (2400, 300, 1.0), "Power3A": (2900, 300, 1.5),
         "HV": (2100, 650, 0.6)}
VIA_ROUTE_UM, VIA_FINAL_MM, VIA_DRILL_MM = 2300, 1.0, 0.5   # same trick for vias
HV_TO_LV_UM = 2450     # + (2.1 - 0.6) / 2 after shrinking = 3.2 mm


def add_gnd_plane(board):
    """Bottom-layer GND pour over the whole board (custom rules keep it 3 mm from HV)."""
    for z in list(board.Zones()):
        board.Delete(z)
    bb = board.GetBoardEdgesBoundingBox()
    inset = pcbnew.FromMM(0.5)
    x0, y0 = bb.GetX() + inset, bb.GetY() + inset
    x1, y1 = bb.GetRight() - inset, bb.GetBottom() - inset
    z = pcbnew.ZONE(board)
    z.SetLayer(pcbnew.B_Cu)
    z.SetNet(board.FindNet("/GND"))
    z.SetLocalClearance(pcbnew.FromMM(0.5))
    z.SetMinThickness(pcbnew.FromMM(0.5))
    z.SetPadConnection(pcbnew.ZONE_CONNECTION_THERMAL)
    z.SetThermalReliefGap(pcbnew.FromMM(0.5))
    z.SetThermalReliefSpokeWidth(pcbnew.FromMM(0.8))
    ol = z.Outline()
    ol.NewOutline()
    for x, y in ((x0, y0), (x1, y0), (x1, y1), (x0, y1)):
        ol.Append(x, y)
    board.Add(z)


def add_connector_stubs(board, ref="J2", length=5.0):
    """Fixed escape stubs on the panel connector: LV pins exit left, HV pins exit right.

    Its 3.96 mm pitch is tighter than the router's HV <-> LV clearance, so the router could not
    leave those pins by itself. GND pins connect to the bottom plane and need no stub.
    """
    fp = board.FindFootprintByReference(ref)
    # final widths: a router track must end on the stub copper itself, which still holds after shrinking
    widths = {k: v[2] for k, v in ROUTE.items()}
    for pad in fp.Pads():
        net = pad.GetNet()
        name = net.GetNetname() if net else ""
        if not name or name == "/GND" or name.startswith("unconnected"):
            continue
        cls = pad.GetNetClassName()
        sign = 1 if cls == "HV" else -1
        p = pad.GetPosition()
        t = pcbnew.PCB_TRACK(board)
        t.SetStart(p)
        t.SetEnd(pcbnew.VECTOR2I(p.x + sign * pcbnew.FromMM(length), p.y))
        t.SetLayer(pcbnew.F_Cu)
        t.SetWidth(pcbnew.FromMM(widths.get(cls, ROUTE["kicad_default"][2])))
        t.SetNet(net)
        t.SetLocked(True)
        board.Add(t)


def _closest_on_segment(p, a, b):
    ax, ay, bx, by = a.x, a.y, b.x, b.y
    dx, dy = bx - ax, by - ay
    L2 = dx * dx + dy * dy
    u = 0.0 if L2 == 0 else max(0.0, min(1.0, ((p.x - ax) * dx + (p.y - ay) * dy) / L2))
    cx, cy = ax + u * dx, ay + u * dy
    return int(round(cx)), int(round(cy)), ((p.x - cx) ** 2 + (p.y - cy) ** 2) ** 0.5


def heal_t_junctions(board, reach_mm=1.6):
    """After shrinking, a track that joined another track mid-way (a T-junction) can end beside
    the narrowed copper. Snap such dangling ends onto the centre line of the same-net track."""
    traces = [t for t in board.GetTracks() if t.Type() == pcbnew.PCB_TRACE_T]
    vias = [t for t in board.GetTracks() if t.Type() == pcbnew.PCB_VIA_T]
    pads = [p for f in board.GetFootprints() for p in f.Pads()]
    reach = pcbnew.FromMM(reach_mm)
    fixed = 0
    for t in traces:
        for is_start in (True, False):
            p = t.GetStart() if is_start else t.GetEnd()
            net, layer = t.GetNetCode(), t.GetLayer()
            if any(pd.GetNetCode() == net and pd.IsOnLayer(layer) and pd.HitTest(p) for pd in pads):
                continue
            if any(v.GetNetCode() == net and v.GetPosition() == p for v in vias):
                continue
            best, touching = None, False
            for o in traces:
                if o is t or o.GetNetCode() != net or o.GetLayer() != layer:
                    continue
                if p in (o.GetStart(), o.GetEnd()):
                    touching = True
                    break
                cx, cy, d = _closest_on_segment(p, o.GetStart(), o.GetEnd())
                if d <= o.GetWidth() / 2:
                    touching = True
                    break
                if d <= reach and (best is None or d < best[2]):
                    best = (cx, cy, d)
            if touching or best is None:
                continue
            q = pcbnew.VECTOR2I(best[0], best[1])
            if is_start:
                t.SetStart(q)
            else:
                t.SetEnd(q)
            fixed += 1
    return fixed


def unconnected(board):
    board.BuildConnectivity()
    return board.GetConnectivity().GetUnconnectedCount(False)


def remove_dangling_stubs(board):
    """Drop duplicate escape stubs (the SES brings them back), then delete each stub whose
    removal does not leave anything unconnected (the router may join a stub anywhere along it)."""
    seen = set()
    for t in [t for t in board.GetTracks() if t.Type() == pcbnew.PCB_TRACE_T and t.IsLocked()]:
        key = (t.GetNetCode(), t.GetLayer(), t.GetStart().x, t.GetStart().y, t.GetEnd().x, t.GetEnd().y)
        if key in seen:
            board.Delete(t)
        seen.add(key)
    base = unconnected(board)
    for t in [t for t in board.GetTracks() if t.Type() == pcbnew.PCB_TRACE_T and t.IsLocked()]:
        s, e = t.GetStart(), t.GetEnd()       # copy values: t is freed by Delete
        keep = (s.x, s.y, e.x, e.y, t.GetLayer(), t.GetWidth(), t.GetNetCode())
        board.Delete(t)
        if unconnected(board) > base:          # it was needed: put it back
            n = pcbnew.PCB_TRACK(board)
            n.SetStart(pcbnew.VECTOR2I(keep[0], keep[1])); n.SetEnd(pcbnew.VECTOR2I(keep[2], keep[3]))
            n.SetLayer(keep[4]); n.SetWidth(keep[5])
            board.Add(n)
            n.SetNet(board.FindNet(keep[6])); n.SetLocked(True)


def patch_via_padstacks(text):
    """Enlarge the via pads in the DSN library (vias are set back to their real size after import)."""
    out = []
    in_via = False
    for line in text.splitlines(keepends=True):
        if line.lstrip().startswith('(padstack "Via'):
            in_via = True
        elif in_via and line.lstrip().startswith("(shape (circle"):
            line = re.sub(r"(\(circle \S+ )[\d.]+", lambda m: m.group(1) + str(VIA_ROUTE_UM), line)
        elif in_via and line.strip() == ")":
            in_via = False
        out.append(line)
    return "".join(out)


def patch_class_rules(text):
    """Replace each class's width/clearance in the DSN with the routing values."""
    out, pos = [], 0
    for m in re.finditer(r"\(class (\S+) ", text):
        name = m.group(1)
        start = m.start()
        rule = text.index("(rule", start)
        end = text.index(")", text.index("(clearance", rule)) + 1
        width, clear, _ = ROUTE[name]
        block = text[rule:end]
        block = re.sub(r"\(width [\d.]+\)", f"(width {width})", block)
        block = re.sub(r"\(clearance [\d.]+\)", f"(clearance {clear})", block, count=1)
        out.append(text[pos:rule] + block)
        pos = end
    out.append(text[pos:])
    return "".join(out)


def shrink_tracks(board):
    final = {("Default" if k == "kicad_default" else k): pcbnew.FromMM(v[2]) for k, v in ROUTE.items()}
    for t in board.GetTracks():
        if t.Type() == pcbnew.PCB_VIA_T:
            t.SetWidth(pcbnew.FromMM(VIA_FINAL_MM))
            t.SetDrill(pcbnew.FromMM(VIA_DRILL_MM))
        elif t.Type() == pcbnew.PCB_TRACE_T:
            name = t.GetNetClassName()
            t.SetWidth(final.get(name, final["Default"]))


def main():
    WORK.mkdir(exist_ok=True)
    dsn, ses = WORK / "board.dsn", WORK / "board.ses"

    board = pcbnew.LoadBoard(str(BOARD))
    try:
        tracks = list(board.GetTracks())
    except TypeError:
        tracks = []
    for t in tracks:
        board.Delete(t)
    for z in list(board.Zones()):
        board.Delete(z)                 # GND is routed as tracks; the pour is added afterwards
    add_connector_stubs(board)
    if not pcbnew.ExportSpecctraDSN(board, str(dsn)):
        raise SystemExit("DSN export failed")

    text = dsn.read_text(encoding="utf-8")
    classes = re.findall(r"\(class (\S+) ", text)
    text = patch_class_rules(text)
    text = patch_via_padstacks(text)
    extra = "".join(
        f"    (class_class (classes HV {c})\n      (rule (clearance {HV_TO_LV_UM}))\n    )\n"
        for c in classes if c != "HV")
    # insert after the last (class ...) block inside (network ...)
    last = text.rfind("    (class ")
    depth, i = 0, last
    while True:
        if text[i] == "(":
            depth += 1
        elif text[i] == ")":
            depth -= 1
            if depth == 0:
                break
        i += 1
    text = text[:i + 1] + "\n" + extra + text[i + 1:]
    dsn.write_text(text, encoding="utf-8")

    ses.unlink(missing_ok=True)
    subprocess.run([str(JAVA), "-jar", str(JAR), "-de", str(dsn), "-do", str(ses),
                    "-mp", str(PASSES), "-l", "en"], check=True, cwd=WORK)   # logs stay in route/
    if not ses.exists():
        raise SystemExit("Freerouting produced no session file")

    if not pcbnew.ImportSpecctraSES(board, str(ses)):
        raise SystemExit("SES import failed")
    shrink_tracks(board)
    print("healed T-junctions:", heal_t_junctions(board))
    remove_dangling_stubs(board)
    add_gnd_plane(board)
    pcbnew.ZONE_FILLER(board).Fill(board.Zones())
    board.Save(str(BOARD))
    print("routed and saved", BOARD.name)


if __name__ == "__main__":
    main()