"""Place + route the SandTable controller and write the KiCad board."""
import math
import os
import sys
import time

sys.path.insert(0, os.path.dirname(__file__))
import design as D
from board import build, check_placement, write_pcb, write_project, keepout_zone, text_node
import router as R

OUT = os.path.expanduser('~/Documents/projects/sandTable-local/sandtable-controller')

# routing width per net
WIDTH = {}
for n in D.NETS:
    WIDTH[n] = 0.2
for n in ('VIN_RAW', 'VIN_FUSED', '+12V'):
    WIDTH[n] = 0.8
WIDTH['GND'] = 0.5
for n in ('X_1A', 'X_1B', 'X_2A', 'X_2B', 'Y_1A', 'Y_1B', 'Y_2A', 'Y_2B'):
    WIDTH[n] = 0.5
for n in ('VBUS', '+3V3'):
    WIDTH[n] = 0.4
for n in ('+5V', '+5V_REG'):          # up to 1 A from the 5 V regulator (LED strip)
    WIDTH[n] = 0.6
FALLBACK = [0.3, 0.2]
# nets that must keep their full width: fail loudly instead of necking down
NO_FALLBACK = {'VIN_RAW', 'VIN_FUSED', '+12V', 'X_1A', 'X_1B', 'X_2A', 'X_2B', 'Y_1A', 'Y_1B', 'Y_2A', 'Y_2B'}
VIA_PAD_GAP = 0.2      # min copper gap between a via and any pad (keeps holes out of solder pads)


def radius(w):
    return round(w / 2 + R.CLEAR + R.MARGIN, 3)


def via_radius(w):
    return max(radius(w), round(R.VIA_D / 2 + R.CLEAR + R.MARGIN, 3))


# Hand-routed USB-C escape (pads are on a 0.5 mm pitch). Board coords.
USB_X = 7.445
PREROUTE = {
    'USB_DP': dict(tracks=[(USB_X, 43.75, 9.0, 43.75, 0.2, 0), (USB_X, 44.75, 9.0, 44.75, 0.2, 0),
                           (9.0, 43.75, 9.0, 44.75, 0.2, 0)], ports=[(9.0, 44.25, (0,))]),
    'USB_DM': dict(tracks=[(USB_X, 43.25, 6.15, 43.25, 0.2, 0), (USB_X, 44.25, 6.15, 44.25, 0.2, 0),
                           (6.15, 43.25, 6.15, 44.25, 0.2, 1)],
                   vias=[(6.15, 43.25), (6.15, 44.25)], ports=[(6.15, 43.75, (1,))]),
    'CC1': dict(tracks=[(USB_X, 42.75, 10.675, 42.75, 0.2, 0)]),
    'CC2': dict(tracks=[(USB_X, 45.75, 10.675, 45.75, 0.2, 0)]),
    'GND': dict(tracks=[(USB_X, 40.75, 9.3, 40.75, 0.4, 0), (USB_X, 47.25, 9.3, 47.25, 0.4, 0),
                        # driver VMOT ceramics (C16, C17): straight down from the socket's GND pin 7
                        (22.86, 1.9, 22.86, 5.4, 0.5, 0), (43.86, 1.9, 43.86, 5.4, 0.5, 0),
                        # boxed-in GND pads get their own via: U5 pin 2 (between the USB lines), C13/C14
                        (11.863, 49.5, 10.55, 49.5, 0.3, 0),
                        (61.775, 18.5, 61.775, 19.6, 0.3, 0), (61.775, 19.6, 61.775, 20.7, 0.3, 0)],
                vias=[(10.55, 49.5), (61.775, 19.6)],
                ports=[(9.3, 40.75, (0,)), (9.3, 47.25, (0,))]),
    # ... and from the socket's VMOT pin 8
    '+12V': dict(tracks=[(20.32, 1.9, 20.32, 5.4, 0.8, 0), (41.32, 1.9, 41.32, 5.4, 0.8, 0)]),
    'VBUS': dict(tracks=[(USB_X, 41.55, 9.3, 41.55, 0.4, 0), (USB_X, 46.45, 9.3, 46.45, 0.4, 0)],
                 ports=[(9.3, 41.55, (0,)), (9.3, 46.45, (0,))]),
}

ORDER_FIRST = ['USB_DP', 'USB_DM', 'CC1', 'CC2', 'VBUS', 'VIN_RAW',
               # motor coils first, outermost lane first (J3.1 / J3.4 / J4.1 / J4.4), then the 12 V trunk
               'X_2B', 'X_1B', 'X_2A', 'X_1A', 'Y_1B', 'Y_2B', 'Y_1A', 'Y_2A', '+12V',
               'STAT',                      # short straight run under the LEDs; boxed in if left late
               '+5V_REG', '+5V', '+3V3']
ORDER_LAST = ['GND']


def net_order(parts):
    pads_by_net = {}
    for p in parts.values():
        for pad in p.pads():
            if pad.net:
                pads_by_net.setdefault(pad.net, []).append(pad)
    rest = [n for n in D.NETS if n not in ORDER_FIRST and n not in ORDER_LAST]

    def span(n):
        xs = [p.x for p in pads_by_net[n]]
        ys = [p.y for p in pads_by_net[n]]
        return (max(xs) - min(xs)) + (max(ys) - min(ys))
    rest.sort(key=span)
    return ORDER_FIRST + rest + ORDER_LAST, pads_by_net


def seg_cells(x1, y1, x2, y2, L):
    n = max(1, int(math.hypot(x2 - x1, y2 - y1) / (R.G / 2)))
    out = set()
    for i in range(n + 1):
        t = i / n
        ix, iy = int(round((x1 + t * (x2 - x1)) / R.G)), int(round((y1 + t * (y2 - y1)) / R.G))
        out.add((L, ix, iy))
    return out


def route_all(parts, order, pads_by_net, verbose=True):
    widths = set(WIDTH.values()) | set(FALLBACK) | {0.4}
    radii = {radius(w) for w in widths} | {via_radius(w) for w in widths}
    rt = R.Router(D.BOARD_W, D.BOARD_H, radii)
    rt.block_edges()
    rt.block_rect(18.75, 56.9, 67.25, D.BOARD_H)          # ESP32 antenna keep-out
    for p in parts.values():
        for pad in p.pads():
            rt.add_pad(pad)
            rt.keep_vias_off_pad(pad, VIA_PAD_GAP)
    pre_components = {}     # net -> list of cell sets (each one electrically connected)
    for net, pr in PREROUTE.items():
        own = rt.net_cells.setdefault(net, set())
        items = []          # (cells, endpoints)
        for (x, y) in pr.get('vias', []):
            rt.add_via(x, y, net)
            ix, iy = rt.cell(x, y)
            items.append(({(0, ix, iy), (1, ix, iy)}, {(round(x, 3), round(y, 3))}))
        for (x1, y1, x2, y2, w, L) in pr['tracks']:
            rt.add_track(x1, y1, x2, y2, w, L, net)
            items.append((seg_cells(x1, y1, x2, y2, L), {(round(x1, 3), round(y1, 3)), (round(x2, 3), round(y2, 3))}))
        # union items sharing an endpoint
        comps = []
        for cells, ends in items:
            merged = [c for c in comps if c[1] & ends]
            for c in merged:
                comps.remove(c)
                cells = cells | c[0]
                ends = ends | c[1]
            comps.append((cells, ends))
        ports = pr.get('ports')
        pc = []
        for c in comps:
            own |= c[0]
            if ports:
                pcells = set()
                for (px, py, Ls) in ports:
                    ix, iy = rt.cell(px, py)
                    if any((L, ix, iy) in c[0] for L in (0, 1)) or any(abs(px - e[0]) < 1e-6 and abs(py - e[1]) < 1e-6 for e in c[1]):
                        for L in Ls:
                            pcells.add((L, ix, iy))
                pc.append((c[0], pcells if pcells else None))
            else:
                pc.append((c[0], None))
        pre_components[net] = pc

    failures = []
    stats = {}
    t0 = time.time()
    for net in order:
        pads = pads_by_net[net]
        w0 = WIDTH[net]
        # terminals: pre-routed components (with the pads they touch) and the remaining pads
        terms = []
        used = set()
        for comp, port_cells in pre_components.get(net, []):
            cells = set(comp)
            pts = []
            for p in pads:
                pc = rt.pad_cells(p)
                if pc & comp:
                    cells |= pc
                    used.add(id(p))
                    pts.append((p.x, p.y))
            if port_cells:
                cells = set(port_cells)
                c = next(iter(port_cells))
                pts = [(c[1] * R.G, c[2] * R.G)]
            if not pts:
                c = next(iter(comp))
                pts = [(c[1] * R.G, c[2] * R.G)]
            terms.append((cells, pts[0][0], pts[0][1], f'pre:{net}'))
        for p in pads:
            if id(p) not in used:
                terms.append((rt.pad_cells(p), p.x, p.y, f'{p.ref}.{p.num}'))
        first = min(terms, key=lambda t: (t[1] + t[2]))
        tree = set(first[0])
        remaining = [t for t in terms if t is not first]
        nseg = nvia = 0
        while remaining:
            sample = list(tree)[::max(1, len(tree) // 400)]

            def dist_to_tree(t):
                return min(abs(t[1] - c[1] * R.G) + abs(t[2] - c[2] * R.G) for c in sample)
            t = min(remaining, key=dist_to_tree)
            targets = t[0]
            remaining.remove(t)
            if targets & tree:          # already connected (pads touching / overlapping)
                tree |= targets
                continue
            path = None
            tries = [w0] if net in NO_FALLBACK else [w0] + [fw for fw in FALLBACK if fw < w0]
            for w in tries:
                path = rt.search(net, radius(w), via_radius(w), tree, targets, t[1], t[2])
                if path:
                    break
            if not path:
                failures.append((net, t[3]))
                continue
            a, b = rt.commit_path(net, path, w)
            nseg += a; nvia += b
            tree |= set(path) | targets
            rt.net_cells[net] |= targets
        stats[net] = (nseg, nvia)
        if verbose:
            print(f'  {net:10s} pads={len(pads):2d} segs={nseg:3d} vias={nvia:2d}  t={time.time() - t0:6.1f}s', flush=True)
    return rt, failures


def stitch_vias(rt, pitch=4.0, dense=((33.0, 37.0, 53.5, 56.5, 2.0),)):
    """GND stitching vias: one per pitch x pitch block wherever a via fits and the ground pour
    reaches it (connected) on both layers. Denser blocks under the ESP32 module."""
    import pour
    r_map = radius(0.2)
    k = rt.nid('GND')
    nx = rt.nx
    keep = {}
    for L in (0, 1):
        mask = pour.fill_mask(rt, 'GND', L, r_map, open_cells=2)
        anchors = [iy * nx + ix for (cl, ix, iy) in rt.net_cells.get('GND', ()) if cl == L]
        anchors += [rt.cell(x, y)[1] * nx + rt.cell(x, y)[0] for (x, y, n) in rt.vias if n == 'GND']
        keep[L] = pour.keep_connected(mask, nx, rt.ny, anchors)
    mv = rt.maps[via_radius(0.2)]
    gnd_vias = [(x, y) for (x, y, n) in rt.vias if n == 'GND']

    def ok(ix, iy, min_gap):
        if ix < 1 or iy < 1 or ix >= nx - 1 or iy >= rt.ny - 1:
            return False
        i = iy * nx + ix
        if rt.no_via[i] or not (keep[0][i] and keep[1][i]):
            return False
        if any(mv[L][i] not in (-1, k) for L in (0, 1)):
            return False
        x, y = rt.xy(ix, iy)
        return all(math.hypot(x - vx, y - vy) >= min_gap for vx, vy in gnd_vias)

    def blocks(x0, y0, x1, y1, p):
        s = int(round(p / R.G))
        for by in range(int(y0 / R.G), int(y1 / R.G), s):
            for bx in range(int(x0 / R.G), int(x1 / R.G), s):
                yield bx, by, s

    added = 0
    passes = [(0.0, 0.0, D.BOARD_W, D.BOARD_H, pitch)] + list(dense)
    for (x0, y0, x1, y1, p) in passes:
        for bx, by, s in blocks(x0, y0, x1, y1, p):
            cx, cy = bx + s // 2, by + s // 2
            cands = sorted(((ix, iy) for iy in range(by, by + s) for ix in range(bx, bx + s)),
                           key=lambda c: (c[0] - cx) ** 2 + (c[1] - cy) ** 2)
            for ix, iy in cands:
                if ok(ix, iy, p * 0.6):
                    x, y = rt.xy(ix, iy)
                    rt.add_via(round(x, 3), round(y, 3), 'GND')
                    gnd_vias.append((x, y))
                    added += 1
                    break
    print(f'  stitching vias: {added}')
    return added


def gnd_pours(rt, parts):
    """GND copper pour on both layers, filled by pour.py (clearance 0.3 mm)."""
    import pour
    from board import uid, fmt, net_list
    from sexpr import QStr
    OX, OY = D.ORIGIN
    k = rt.nid('GND')
    r_map = radius(0.2)          # 0.38 mm: 0.3 pour clearance + half cell diagonal
    net_ids = net_list()
    zones = []
    for L, lname in ((1, 'B.Cu'), (0, 'F.Cu')):
        mask = pour.fill_mask(rt, 'GND', L, r_map, open_cells=2)
        # anchors: GND copper cells on this layer (pads, tracks, vias)
        anchors = []
        for (cl, ix, iy) in rt.net_cells.get('GND', ()):
            if cl == L:
                anchors.append(iy * rt.nx + ix)
        for (x, y, n) in rt.vias:
            if n == 'GND':
                ix, iy = rt.cell(x, y)
                anchors.append(iy * rt.nx + ix)
        keep = pour.keep_connected(mask, rt.nx, rt.ny, anchors)
        polys = pour.polygons(keep, rt.nx, rt.ny)
        filled = []
        for poly in polys:
            pts = pour.to_mm(poly, R.G)
            filled.append(['filled_polygon', ['layer', QStr(lname)],
                           ['pts'] + [['xy', fmt(x + OX), fmt(y + OY)] for x, y in pts]])
        cells = sum(keep)
        print(f'  GND pour {lname}: {len(polys)} islands, {cells * R.G * R.G:.0f} mm2')
        inset = 0.25
        outline = [(inset, inset), (D.BOARD_W - inset, inset), (D.BOARD_W - inset, D.BOARD_H - inset), (inset, D.BOARD_H - inset)]
        zones.append(['zone', ['net', str(net_ids['GND'])], ['net_name', QStr('/GND')], ['layer', QStr(lname)],
                      ['uuid', QStr(uid('gndzone', lname))], ['name', QStr('GND_' + lname[0])],
                      ['hatch', 'edge', '0.5'], ['connect_pads', 'yes', ['clearance', '0.3']],
                      ['min_thickness', '0.25'], ['filled_areas_thickness', 'no'],
                      ['fill', 'yes', ['thermal_gap', '0.5'], ['thermal_bridge_width', '0.5']],
                      ['polygon', ['pts'] + [['xy', fmt(x + OX), fmt(y + OY)] for x, y in outline]]] + filled)
    return zones


def silk_labels():
    T = []
    T.append(text_node('SANDTABLE', 53.8, 27.6, 'F.SilkS', size=1.1, bold=True))
    T.append(text_node(f'CTRL v{D.VERSION}', 53.8, 29.6, 'F.SilkS', size=1.0))
    for txt, x, y, ang in (('SPIN', 4.5, 15.2, 0), ('12V', 2.4, 17.6, 0), ('USB', 2.2, 37.3, 0),
                           ('ARM', 71.0, 7.15, 90), ('HOME', 71.0, 21.0, 90), ('HALL', 71.4, 33.6, 90),
                           ('LED', 71.0, 43.2, 90), ('EXP', 74.6, 55.8, 90)):
        T.append(text_node(txt, x, y, 'F.SilkS', size=0.9, angle=ang))
    # driver orientation: match the pin names printed on the DRV8825 module
    for x1 in (38.1, 59.1):                  # A1 / A2 pin-1 x; pins step 2.54 mm to the left
        T.append(text_node('VMOT', x1 - 7 * 2.54, 3.5, 'F.SilkS', size=0.8))
        T.append(text_node('EN', x1 - 7 * 2.54, 13.0, 'F.SilkS', size=0.8))
        T.append(text_node('DIR', x1, 13.0, 'F.SilkS', size=0.8))
    # back side: name plate
    T.append(text_node(f'SANDTABLE CONTROLLER v{D.VERSION}', 40.0, 30.0, 'B.SilkS', size=1.6, bold=True))
    T.append(text_node('ESP32 + 2x DRV8825  |  12V in  |  FluidNC', 40.0, 33.0, 'B.SilkS', size=1.0))
    T.append(text_node('lazypugs 2026', 40.0, 35.5, 'B.SilkS', size=1.0))
    # JP1 (under the theta driver): bridge only for A4988 modules
    jx, jy = D.PARTS['JP1'][3], D.PARTS['JP1'][4]
    T.append(text_node('A4988 ONLY', jx, jy + 1.9, 'F.SilkS', size=0.7))
    T.append(text_node('WayWayWay', 40.0, 38.0, 'B.SilkS', size=1.0))   # PCBWay prints its order number here
    return T


def main():
    parts = build()
    probs = check_placement(parts)
    for p in probs:
        print('  !', p)
    order, pads_by_net = net_order(parts)
    rt, failures = route_all(parts, order, pads_by_net)
    print(f'tracks={len(rt.tracks)} vias={len(rt.vias)} failures={len(failures)}')
    crowded = [(x, y, pd.ref, pd.num) for (x, y, n) in rt.vias for p in parts.values() for pd in p.pads()
               if pd.layers and pd.kind != 'np_thru_hole' and pd.dist(x, y) < R.VIA_D / 2 + VIA_PAD_GAP - 1e-6]
    print(f'vias closer than {VIA_PAD_GAP} mm to a pad: {len(crowded)}', crowded[:6])
    for f in failures:
        print('   FAILED', f)
    stitch_vias(rt)
    tracks = [(x1, y1, x2, y2, w, L, n) for (x1, y1, x2, y2, w, L, n) in rt.tracks]
    zones = [keepout_zone('ANTENNA', 18.75, 56.9, 67.25, D.BOARD_H)]
    zones += gnd_pours(rt, parts)
    os.makedirs(OUT, exist_ok=True)
    write_pcb(os.path.join(OUT, 'sandtable-controller.kicad_pcb'), parts, tracks=tracks, vias=rt.vias,
              zones=zones, texts=silk_labels())
    write_project(os.path.join(OUT, 'sandtable-controller.kicad_pro'), 'sandtable-controller')
    return failures


if __name__ == '__main__':
    main()