#!/usr/bin/env python3
"""retire_block.py — remove a circuit block from a CM5-carrier KiCad sheet cleanly.

For the diff-from-reference port (see .claude/skills/kicad-port). Given a sheet and a
set of reference designators to drop plus the net-name globs they own, this:
  1. excises each named (symbol ...) block (paren-matched, by Reference property),
  2. floods the retired nets from their (label ...) coords through (wire ...) segments,
  3. removes those labels + every wire in the flooded nets,
  4. places (no_connect ...) markers on the exposed module pins that sat on those nets,
  5. writes back only if parentheses stay balanced.

Module pin coordinates are computed from the lib_symbol pin locals + the placement
transform (validated abs = origin + (lx, -ly) for rot0/no-mirror; other orientations
are rejected rather than guessed). Idempotent-ish: re-running finds nothing to do.

Usage:
  python3 tools/retire_block.py <sheet.kicad_sch> \
      --symbols J7,U18 --nets 'SD_*' \
      --module-lib CM5IO:ComputeModule5-CM5 [--apply]
Without --apply it's a dry run (prints the plan, writes nothing).
"""
import re, sys, argparse, fnmatch, math, uuid as uuidlib

S = r'\s*'

def xform(px, py, rot, mirror, lx, ly):
    """Schematic abs coord of a lib pin local (lx,ly) for a placement. Validated:
    mirror, then CCW rotate, then lib-Y-up -> schematic-Y-down flip."""
    X, Y = lx, ly
    if mirror == 'y': X = -X
    if mirror == 'x': Y = -Y
    r = math.radians(rot)
    rx = X * math.cos(r) - Y * math.sin(r)
    ry = X * math.sin(r) + Y * math.cos(r)
    return (round(px + rx, 2), round(py - ry, 2))

def paren_end(t, s):
    d = 0
    for i in range(s, len(t)):
        if t[i] == '(': d += 1
        elif t[i] == ')':
            d -= 1
            if d == 0: return i + 1
    raise ValueError("unbalanced from offset %d" % s)

def fnum(x): return round(float(x), 2)

def find_symbol_blocks(t):
    """Yield (start,end,reference) for every placed (symbol (lib_id ...)) block."""
    for m in re.finditer(r'\(symbol' + S + r'\(lib_id' + S + r'"', t):
        s = m.start(); e = paren_end(t, s); seg = t[s:e]
        ref = re.search(r'\(property' + S + r'"Reference"' + S + r'"([^"]+)"', seg)
        if ref: yield (s, e, ref.group(1))

def module_pins_abs(t, module_lib):
    """Compute {pin_number:(x,y)} for the placed module instance on this sheet."""
    inst = re.search(r'\(symbol' + S + r'\(lib_id' + S + r'"' + re.escape(module_lib) + r'"\)(.{0,500}?)\(unit' + S + r'(\d+)\)', t, re.S)
    if not inst: return {}, None
    body, unit = inst.group(1), inst.group(2)
    at = re.search(r'\(at' + S + r'([\d.-]+)' + S + r'([\d.-]+)' + S + r'([\d.-]+)\)', body)
    mir = re.search(r'\(mirror' + S + r'(\w+)\)', body)
    px, py, rot = float(at.group(1)), float(at.group(2)), float(at.group(3))
    if rot != 0 or mir:
        sys.exit(f"REFUSING: module instance has rot={rot} mirror={mir and mir.group(1)} — transform only validated for rot0/no-mirror.")
    libname = module_lib.split(':')[-1]
    pins = {}
    for pm in re.finditer(r'\(pin\b.*?\(at' + S + r'([\d.-]+)' + S + r'([\d.-]+)' + S + r'[\d.-]+\).*?\(number' + S + r'"([^"]+)"', t, re.S):
        lx, ly, num = float(pm.group(1)), float(pm.group(2)), pm.group(3)
        pins.setdefault(num, (round(px + lx, 2), round(py - ly, 2)))
    return pins, unit

def collect(t, pat):
    """Return list of (start,end,coord,text) for top-level objects named `pat` token."""
    out = []
    for m in re.finditer(r'\(' + pat + r'\b', t):
        s = m.start(); e = paren_end(t, s); seg = t[s:e]
        at = re.search(r'\(at' + S + r'([\d.-]+)' + S + r'([\d.-]+)', seg)
        out.append((s, e, (fnum(at.group(1)), fnum(at.group(2))) if at else None, seg))
    return out

def wire_pts(seg):
    return [(fnum(a), fnum(b)) for a, b in re.findall(r'\(xy' + S + r'([\d.-]+)' + S + r'([\d.-]+)\)', seg)]

def lib_pins(t, lib_id):
    """Local (lx,ly,angle) for every pin of a lib_symbol definition."""
    name = lib_id.split(':')[-1]
    lm = re.search(r'\(symbol' + S + r'"' + re.escape(lib_id) + r'"', t) or \
         re.search(r'\(symbol' + S + r'"' + re.escape(name) + r'"', t)
    if not lm: return []
    seg = t[lm.start():paren_end(t, lm.start())]
    return [(float(a), float(b), float(c)) for a, b, c in
            re.findall(r'\(pin\b.*?\(at' + S + r'([\d.-]+)' + S + r'([\d.-]+)' + S + r'([\d.-]+)\)', seg, re.S)]

def placed_symbols(t):
    """Yield dict per placed symbol: start,end,ref,lib,at(px,py,rot),mirror."""
    for m in re.finditer(r'\(symbol' + S + r'\(lib_id' + S + r'"([^"]+)"\)', t):
        s = m.start(); e = paren_end(t, s); seg = t[s:e]
        ref = re.search(r'\(property' + S + r'"Reference"' + S + r'"([^"]+)"', seg)
        at = re.search(r'\(at' + S + r'([\d.-]+)' + S + r'([\d.-]+)' + S + r'([\d.-]+)\)', seg)
        mir = re.search(r'\(mirror' + S + r'(\w+)\)', seg)
        yield dict(s=s, e=e, ref=ref.group(1) if ref else None, lib=m.group(1),
                   at=(float(at.group(1)), float(at.group(2)), float(at.group(3))) if at else None,
                   mirror=mir.group(1) if mir else None)

def sym_pin_coords(t, sym):
    if not sym['at']: return []
    px, py, rot = sym['at']
    return [xform(px, py, rot, sym['mirror'], lx, ly) for lx, ly, _ in lib_pins(t, sym['lib'])]

def sheet_port_spans(t, names):
    """On a ROOT sheet, find (pin "NAME" ...) hierarchical ports inside (sheet ...) blocks
    whose name is in `names`. Returns (pin spans to remove, their connection coords)."""
    spans, coords = [], set()
    for m in re.finditer(r'\(sheet\b', t):
        s = m.start(); e = paren_end(t, s)
        for pm in re.finditer(r'\(pin' + S + r'"([^"]+)"', t[s:e]):
            if pm.group(1) in names:
                ps = s + pm.start(); pe = paren_end(t, ps)
                at = re.search(r'\(at' + S + r'([\d.-]+)' + S + r'([\d.-]+)', t[ps:pe])
                spans.append((ps, pe))
                if at: coords.add((fnum(at.group(1)), fnum(at.group(2))))
    return spans, coords

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('sheet')
    ap.add_argument('--symbols', default='')
    ap.add_argument('--nets', default='')
    ap.add_argument('--nc-pins', default='', help='module pin NUMBERS to no-connect (from analyzer regressions)')
    ap.add_argument('--sheet-ports', default='', help='hierarchical net names: remove matching sheet-symbol pins + connecting root wires (run on the ROOT sheet)')
    ap.add_argument('--module-lib', default='CM5IO:ComputeModule5-CM5')
    ap.add_argument('--apply', action='store_true')
    a = ap.parse_args()
    drop_refs = set(filter(None, a.symbols.split(',')))
    net_globs = list(filter(None, a.nets.split(',')))
    nc_nums = list(filter(None, a.nc_pins.split(',')))
    port_names = set(filter(None, a.sheet_ports.split(',')))
    t = open(a.sheet).read()

    modpins, unit = module_pins_abs(t, a.module_lib)
    pin_at = {xy: n for n, xy in modpins.items()}
    nc_pins = [(n, modpins[n]) for n in nc_nums if n in modpins]
    missing_nc = [n for n in nc_nums if n not in modpins]
    if missing_nc: sys.exit(f"NC pins not found on module: {missing_nc}")

    # 1) symbols to remove + their abs pin coords (for stub cleanup)
    placed = list(placed_symbols(t))
    removed_syms = [sy for sy in placed if sy['ref'] in drop_refs]
    found_refs = {sy['ref'] for sy in removed_syms}
    sym_spans = [(sy['s'], sy['e']) for sy in removed_syms]
    removed_pin_coords = set()
    for sy in removed_syms:
        removed_pin_coords |= set(sym_pin_coords(t, sy))

    # 2) retired labels by glob — match ANY label type (local / hierarchical / global)
    labels = collect(t, 'label') + collect(t, 'hierarchical_label') + collect(t, 'global_label')
    retired_labels = [(s, e, c) for s, e, c, seg in labels
                      for nm in [re.search(r'"([^"]+)"', seg).group(1)]
                      if any(fnmatch.fnmatch(nm, g) for g in net_globs)]
    retired_coords = {c for _, _, c in retired_labels}

    # 3) anchor-aware wire removal via connected components.
    #    A wire is removed only if its component (a) is orphaned by THIS edit (touches a
    #    removed pin, retired label, or target NC pin) AND (b) reaches no KEPT anchor —
    #    a kept non-power component pin or a kept label. Shared nets (e.g. a camera FFC's
    #    I2C/power pins) stay anchored by the kept side, so they survive; pre-existing
    #    floating wires aren't seeded, so they're untouched.
    wires = collect(t, 'wire')
    wpts = [(s, e, wire_pts(seg)) for s, e, _, seg in wires]
    parent = {}
    def find(x):
        parent.setdefault(x, x)
        root = x
        while parent[root] != root: root = parent[root]
        while parent[x] != root: parent[x], x = root, parent[x]
        return root
    def union(a, b): parent[find(a)] = find(b)
    for s, e, pts in wpts:
        for p in pts[1:]:
            union(pts[0], p)

    nc_coords = {xy for _, xy in nc_pins}
    kept_label_coords = {c for s, e, c, seg in labels
                         if not any(fnmatch.fnmatch(re.search(r'"([^"]+)"', seg).group(1), g)
                                    for g in net_globs)}
    kept_pin_coords = set()
    for sy in placed:
        if sy['ref'] in drop_refs or sy['lib'].startswith('power:'): continue
        kept_pin_coords |= set(sym_pin_coords(t, sy))
    kept_pin_coords -= nc_coords          # pins we are NC'ing are being disconnected
    anchors = kept_pin_coords | kept_label_coords
    # sheet-symbol hierarchical ports to drop (root sheet) + their connection coords
    port_spans, port_coords = sheet_port_spans(t, port_names) if port_names else ([], set())
    seeds = set(removed_pin_coords) | set(retired_coords) | nc_coords | port_coords

    anchored_roots, seeded_roots = set(), set()
    for c in anchors:
        if c in parent: anchored_roots.add(find(c))
    for c in seeds:
        if c in parent: seeded_roots.add(find(c))
    removable_roots = seeded_roots - anchored_roots
    wire_remove = [(s, e, pts) for s, e, pts in wpts if find(pts[0]) in removable_roots]
    wire_remove_spans = {(s, e) for s, e, _ in wire_remove}
    freed = set(removed_pin_coords)
    for _, _, pts in wire_remove:
        freed |= set(pts)
    kept_wire_pts = set()
    for s, e, pts in wpts:
        if (s, e) not in wire_remove_spans:
            kept_wire_pts |= set(pts)
    flooded_wires = wire_remove  # naming compat for the summary print
    stub_wires = []

    # 5) orphaned power flags (power:*): pin freed by my edit, no kept wire left on it
    power_spans = []
    for sy in placed:
        if sy['ref'] in drop_refs or not sy['lib'].startswith('power:'): continue
        pcs = sym_pin_coords(t, sy)
        if pcs and all(pc in freed and pc not in kept_wire_pts for pc in pcs):
            power_spans.append((sy['s'], sy['e']))

    # 6) dangling no_connects sitting on a removed symbol's pin
    dangling_nc = [(s, e) for s, e, c, _ in collect(t, 'no_connect') if c in removed_pin_coords]

    # 7) labels orphaned by my edit (any label type): coord freed, no kept wire left,
    #    not already retired. Covers local (label), and cross-sheet (hierarchical_label /
    #    global_label) that only fed a removed connector (e.g. a camera FFC's SCL0/CAM_GPIO).
    retired_label_spans = {(s, e) for s, e, _ in retired_labels}
    all_label_objs = collect(t, 'label') + collect(t, 'hierarchical_label') + collect(t, 'global_label')
    orphan_labels = [(s, e) for s, e, c, _ in all_label_objs
                     if (s, e) not in retired_label_spans and c in freed and c not in kept_wire_pts]

    nc_pins = sorted(nc_pins, key=lambda kv: int(re.sub(r'\D', '', kv[0]) or 0))
    print(f"sheet: {a.sheet}  (module unit {unit}, {len(modpins)} pins)")
    print(f"symbols to remove: {sorted(found_refs)}  (requested {sorted(drop_refs)}; missing {sorted(drop_refs-found_refs)})")
    print(f"retired-net labels: {len(retired_labels)} | wires removed (anchor-aware): {len(wire_remove)}")
    print(f"orphaned power flags: {len(power_spans)} | dangling NC removed: {len(dangling_nc)} | orphaned labels: {len(orphan_labels)}")
    if port_names: print(f"sheet-symbol ports removed: {len(port_spans)} ({sorted(port_names)})")
    print(f"module pins to NC ({len(nc_pins)}): " + ", ".join(f"{n}@{xy}" for n, xy in nc_pins))

    # build new text: remove all spans, then append fresh no_connects on exposed module pins
    spans = sorted(set(sym_spans) | {(s, e) for s, e, _ in retired_labels}
                   | wire_remove_spans | set(power_spans) | set(dangling_nc) | set(orphan_labels)
                   | set(port_spans), key=lambda x: -x[0])
    nt = t
    for s, e in spans:
        p = s
        while p > 0 and nt[p-1] in '\t ': p -= 1
        if p > 0 and nt[p-1] == '\n': p -= 1
        nt = nt[:p] + nt[e:]
    # insert no_connect markers before the final closing paren of the sheet
    nc_block = "".join(
        f'\t(no_connect\n\t\t(at {xy[0]} {xy[1]})\n\t\t(uuid "{uuidlib.uuid5(uuidlib.NAMESPACE_DNS, a.sheet+n)}")\n\t)\n'
        for n, xy in nc_pins)
    last = nt.rstrip().rfind(')')
    nt = nt[:last] + nc_block + nt[last:]
    assert nt.count('(') == nt.count(')'), "UNBALANCED — aborting"

    if a.apply:
        open(a.sheet, 'w').write(nt)
        print(f"APPLIED. parens balanced ({nt.count('(')}). +{len(nc_pins)} no_connect, -{len(spans)} objects.")
    else:
        print(f"DRY RUN ok. parens would balance ({nt.count('(')}). Re-run with --apply to write.")

if __name__ == '__main__':
    main()