#!/usr/bin/env python3
"""bake_board_glb.py — turn a KiCad board STEP-derived GLB into a viewer-ready, tagged model.

The Path B model-prep step (see SKILL.md §5c). What it does, and why each step exists:

1. TAG every solid `section|REFDES|i` — **refdes comes from the GLB node hierarchy**: KiCad's
   STEP export names each component's assembly node with its refdes (R39, U2 …); both converters
   (step2glb, cascadio) preserve it. Geometry-bearing LEAF nodes are OCCT tags (`=>[0:1:1:41]`),
   so walk UP the parent chain to the first node matching a refdes that exists in the .kicad_pcb.
   Verified 297/297 solids on a 143-footprint board — no positional guessing.
   (Do NOT match by nearest centroid as the primary method: on the same board it provably
   mistagged 4/297 — a big IC's pin-1 dimple sits closer to a neighbouring 0402 than to its own
   origin, and stale geometry lands wherever it was when the STEP was exported.)
2. CROSS-CHECK each solid's centroid against its refdes's .kicad_pcb position → a distance over
   ~2 mm means the STEP predates a component move (**stale geometry**). Keep the (correct) refdes
   tag, but flag it in the report — the page must say the model is stale for those refs.
3. SECTION each refdes. Default: **derive sheet membership from the design files** — the PCB's
   `(path "/<sheet-uuid>/<symbol-uuid>")` entries against the root .kicad_sch's sheet
   declarations (works without the sub-sheet files). An explicit blocks.json (author's finer
   sections) overrides; refs in neither get "other".
4. BAKE node transforms into vertex data (KiCad places purely by node transform; the Adom viewer
   bundle does NOT apply node transforms — unbaked, the whole BOM collapses at the origin), then
   RECENTER on the board-outline centre. Verified after export: every node matrix is identity.
   Coordinate convention: STEP is METRES with Y NEGATED vs the PCB (glb_y = -pcb_y/1000).
5. DE-FIGHT the layers (measured on a real board: mask sheet sits 10-15 µm off copper, silk 10 µm
   under component bottoms — about ONE depth-buffer step at viewing distance, i.e. guaranteed
   shimmer): lift top MASK +15 µm, top SILK +15 µm, components +30 µm (skip parts that wrap the
   board edge); mirror negative on the bottom. Order (silk above mask above copper) is preserved.
   Also force the PCB material OPAQUE — exporters emit it at alpha 0.98 BLEND, which re-blends
   the exactly-coplanar buried copper faces every frame (the worst clip of all, unfixable by
   offsets since one copper mesh carries both slabs). Viewer-side, also raise the camera's
   near plane (minZ = radius/100) — see the board-viewer template.
6. COVERAGE REPORT: real parts with NO solid (invisible in the viewer!) vs legitimately-bodiless
   footprints, plus DNP refs and stale-geometry refs. Print it and put it on the page (§7B.4).

Usage:
  bake_board_glb.py <raw.glb> <board.kicad_pcb> <out.glb>
                    [--sch root.kicad_sch] [--blocks blocks.json]
                    [--report coverage.json] [--index board-index.json]
                    [--silk-um 15] [--mask-um 15] [--comp-um 30]

NOTE: feed this a **cascadio/float32 GLB**. step2glb output is KHR_mesh_quantization-compressed,
which trimesh misreads by 32767x — dequantize first (gltf-transform dequantize) or convert with
cascadio for the bake and keep step2glb for direct-to-viewer uses.
Requires: trimesh + numpy.
"""
import re, json, collections, argparse
import numpy as np
import trimesh

BODILESS_FP = ('SolderPad', 'SolderPads', 'MountingHole', 'SolderJumper', 'TestPoint',
               'Fiducial', 'LOGO', 'Logo', 'Feeder_Pins')
REFPAT = re.compile(r'^([A-Z]{1,4}[0-9]+)(?:[ _].*)?$')

ap = argparse.ArgumentParser()
ap.add_argument('raw_glb'); ap.add_argument('kicad_pcb'); ap.add_argument('out_glb')
ap.add_argument('--sch'); ap.add_argument('--blocks')
ap.add_argument('--report'); ap.add_argument('--index')
ap.add_argument('--silk-um', type=float, default=15.0)
ap.add_argument('--mask-um', type=float, default=15.0)
ap.add_argument('--comp-um', type=float, default=30.0)
ap.add_argument('--stale-mm', type=float, default=2.0)
A = ap.parse_args()

# ---------------- footprints (refdes, value, fp name, position, layer, dnp, sheet path)
s = open(A.kicad_pcb, encoding='utf8', errors='replace').read()
fps = []
for b in s.split('\n\t(footprint ')[1:]:
    g = lambda k: (re.search(r'\(property "%s" "([^"]*)"' % k, b) or [None, None])[1]
    at = re.search(r'\(at (-?[\d.]+) (-?[\d.]+)(?: (-?[\d.]+))?\)', b)
    lay = re.search(r'\(layer "([^"]+)"\)', b)
    head = b[:b.find('(pad ')] if '(pad ' in b else b
    pth = re.search(r'\(path "([^"]*)"\)', head)
    fps.append(dict(ref=g('Reference'), val=g('Value'),
                    fp=re.match(r'"([^"]+)"', b).group(1).split(':')[-1],
                    x=float(at.group(1)), y=float(at.group(2)),
                    rot=float(at.group(3)) if at.group(3) else 0.0,
                    layer=lay.group(1) if lay else '?',
                    dnp=bool(re.search(r'\(attr[^)]*\bdnp\b', head)),
                    path=pth.group(1) if pth else ''))
fps = [f for f in fps if f['ref'] and not f['ref'].startswith('G*')]
REFS = {f['ref'] for f in fps}
F = {f['ref']: f for f in fps}

# ---------------- sections: sheet-derived default, blocks.json override
sheet_of = {}
if A.sch:
    sch = open(A.sch, encoding='utf8', errors='replace').read()
    sheets = {}   # uuid -> sheet name
    for m in re.finditer(r'\(sheet\s(.*?)\n\t\)', sch, re.S):
        u = re.search(r'\(uuid "([^"]+)"\)', m.group(1))
        n = re.search(r'\(property "Sheetname" "([^"]+)"', m.group(1))
        if u and n:
            sheets[u.group(1)] = n.group(1)
    for f in fps:
        parts = [p for p in f['path'].split('/') if p]
        f['sheet'] = sheets.get(parts[0], 'root') if parts else 'root'
        sheet_of[f['ref']] = f['sheet']
BLK = {}
if A.blocks:
    BL = json.load(open(A.blocks))
    BLK = {r: k for k, v in BL.items() for r in v}
def section(ref):
    return BLK.get(ref) or sheet_of.get(ref, 'other') if (BLK or sheet_of) else 'other'

# ---------------- board outline centre
xs, ys = [], []
for m in re.finditer(r'\(gr_(?:line|arc|rect|circle|poly)\s(.*?)\n\t\)', s, re.S):
    if '"Edge.Cuts"' not in m.group(1):
        continue
    for x, y in re.findall(r'\((?:start|end|mid|center|xy)\s+(-?[\d.]+)\s+(-?[\d.]+)\)', m.group(1)):
        xs.append(float(x)); ys.append(float(y))
CX, CY = (min(xs) + max(xs)) / 2 / 1000.0, -(min(ys) + max(ys)) / 2 / 1000.0
W = (max(xs) - min(xs)) / 1000.0

# ---------------- load + refdes from node ancestry
sc = trimesh.load(A.raw_glb)
parent = {c: p for p, c, _ in sc.graph.to_edgelist()}
def refdes_of(node):
    cur = node
    while True:
        m = REFPAT.match(str(cur))
        if m and m.group(1) in REFS:
            return m.group(1)
        if cur not in parent:
            return None
        cur = parent[cur]

SUBKIND = [('PCB', 'PCB'), ('copper', 'COPPER'), ('soldermask', 'MASK'), ('silkscreen', 'SILK')]
insts = []
z_top = 0.0
for node in sc.graph.nodes_geometry:
    T, gname = sc.graph[node]
    m = sc.geometry[gname].copy(); m.apply_transform(T)
    kind = None
    if (m.bounds[1][0] - m.bounds[0][0]) > W * 0.9:
        base = gname.rsplit('_', 1)[0] if gname.rsplit('_', 1)[-1].isdigit() else gname
        for key, tag in SUBKIND:
            if key in base:
                kind = tag
        if kind == 'PCB':
            z_top = float(m.bounds[1][2])
    insts.append((node, gname, m, kind))

out = trimesh.Scene(); counts = collections.Counter(); idx = {}
stale = {}; unmatched = []
zmid = z_top / 2
for node, gname, m, kind in insts:
    if kind:                                   # substrate
        blk, ref = 'substrate', kind
        zc = float((m.bounds[0][2] + m.bounds[1][2]) / 2)
        if kind in ('MASK', 'SILK'):           # de-fight the zero-thickness sheets
            d = (A.silk_um if kind == 'SILK' else A.mask_um) * 1e-6
            m.apply_translation([0, 0, d if zc > zmid else -d])
        if kind == 'PCB':                      # kill the BLEND-over-buried-copper fight
            mat = getattr(m.visual, 'material', None)
            if mat is not None:
                if getattr(mat, 'baseColorFactor', None) is not None:
                    bcf = list(mat.baseColorFactor); bcf[3] = 255
                    mat.baseColorFactor = bcf
                mat.alphaMode = 'OPAQUE'
    else:
        ref = refdes_of(node)
        if ref is None:
            unmatched.append(gname); continue
        blk = section(ref)
        c = (m.bounds[0] + m.bounds[1]) / 2
        f = F[ref]
        d_mm = float(np.hypot(c[0] - f['x'] / 1000.0, c[1] - -f['y'] / 1000.0)) * 1000
        # big packages have sub-solids far from the origin; only flag if the WHOLE part drifts,
        # so track the minimum distance per ref and decide after the loop
        stale.setdefault(ref, []).append(d_mm)
        # lift components clear of the raised silk — except parts that wrap the board edge
        if m.bounds[0][2] < zmid:
            pass                               # e.g. a USB shell / through-hole pins: leave alone
        else:
            m.apply_translation([0, 0, A.comp_um * 1e-6])
    m.apply_translation([-CX, -CY, 0.0])
    name = f"{blk}|{ref}|{counts[(blk, ref)]}"; counts[(blk, ref)] += 1
    out.add_geometry(m, node_name=name, geom_name=name)
    b = m.bounds; k = f"{blk}|{ref}"
    e = idx.setdefault(k, {'blk': blk, 'ref': ref, 'lo': b[0].tolist(), 'hi': b[1].tolist(), 'n': 0})
    e['lo'] = np.minimum(e['lo'], b[0]).tolist(); e['hi'] = np.maximum(e['hi'], b[1]).tolist()
    e['n'] += 1

data = out.export(file_type='glb'); open(A.out_glb, 'wb').write(data)
r = trimesh.load(A.out_glb)
bad = [n for n in r.graph.nodes_geometry if not np.allclose(r.graph[n][0], np.eye(4), atol=1e-9)]
assert not bad, f'bake failed - nodes still carry transforms: {bad[:5]}'

# ---------------- coverage + staleness
stale_refs = {ref: round(min(ds), 2) for ref, ds in stale.items() if min(ds) > A.stale_mm}
tagged = {v['ref'] for v in idx.values() if v['blk'] != 'substrate'}
missing_real, bodiless = [], []
for f in fps:
    if f['ref'] in tagged:
        continue
    (bodiless if any(k in f['fp'] for k in BODILESS_FP) else missing_real).append(
        dict(ref=f['ref'], val=f['val'], fp=f['fp']))
report = dict(solids=sum(1 for i in insts), refs_tagged=len(tagged),
              missing_real=missing_real, bodiless_ok=bodiless,
              stale_geometry=stale_refs, unmatched_solids=unmatched,
              dnp=[f['ref'] for f in fps if f['dnp']],
              sections=dict(collections.Counter(section(f['ref']) for f in fps)))
print(f"baked {len(insts)} solids -> {A.out_glb} ({len(data)/1e6:.1f} MB); identity verified")
print(f"tagged {len(tagged)} refs | REAL parts with NO body ({len(missing_real)}): "
      f"{[m['ref'] for m in missing_real]}")
print(f"STALE geometry (moved after STEP export): {stale_refs}")
print(f"bodiless-ok {len(bodiless)} | unmatched {unmatched} | DNP {report['dnp']}")
if A.report:
    json.dump(report, open(A.report, 'w'), indent=1); print('report ->', A.report)
if A.index:
    FI = {f['ref']: dict(ref=f['ref'], val=f['val'], fp=f['fp'], layer=f['layer'], rot=f['rot'],
                         x=f['x'] / 1000.0 - CX, y=-f['y'] / 1000.0 - CY,
                         blk=section(f['ref']), sheet=f.get('sheet'), dnp=f['dnp'],
                         stale_mm=stale_refs.get(f['ref']),
                         has3d=(f"{section(f['ref'])}|{f['ref']}") in idx)
          for f in fps}
    json.dump({'groups': idx, 'fps': FI, 'centre_mm': [CX * 1000, -CY * 1000],
               'coverage': report},
              open(A.index, 'w'), indent=0)
    print('index ->', A.index)