"""Repair cross-midline finger skin bindings by INPAINTING from the mesh's own healthy
neighbours.

Mako's shipped rig binds ~444 left-hand verts to RIGHT middle-finger bones (many at weight
1.0, lever arm ~1.8 m), so any middle-finger rotation hurls them across the body. A blunt
_r -> _l mirror remap does NOT fix it: those verts sit 5.8-10.5 cm from the mirrored bone and
their nearest correct joints are thumb/pinky, so remapping would bind thumb skin to the middle
finger and tear. Instead we discard the corrupt influences and refill each vert from its
HEALTHY neighbours on the same mesh (topological BFS first, spatial fallback), which is exactly
what the surrounding 13k correctly-bound left-hand verts already encode.

usage:
  skin_crosshand_repair.py in.glb out.glb [--diagnose] [--k 8] [--report]

  --diagnose  analyse and print only, write nothing
  --k N       neighbours to blend per repaired vert (default 8)
"""
import json, struct, sys, math
from pathlib import Path
from collections import deque, defaultdict

FING = ("thumb", "index", "middle", "ring", "pinky")
MID = 0.02          # metres either side of x=0 that counts as "across the midline"


def read_glb(p):
    d = Path(p).read_bytes()
    length = struct.unpack_from("<I", d, 8)[0]
    off = 12
    chunks = []
    g = None
    while off < length:
        clen, ct = struct.unpack_from("<II", d, off)
        off += 8
        if ct == 0x4E4F534A:
            g = json.loads(d[off:off + clen].decode("utf-8"))
            chunks.append([ct, None])
        else:
            chunks.append([ct, bytearray(d[off:off + clen])])
        off += clen
    return g, chunks


def write_glb(path, g, chunks):
    js = json.dumps(g, separators=(",", ":")).encode("utf-8")
    js += b" " * ((4 - len(js) % 4) % 4)
    body = b""
    for ct, payload in chunks:
        if ct == 0x4E4F534A:
            payload = js
        body += struct.pack("<II", len(payload), ct) + bytes(payload)
    Path(path).write_bytes(struct.pack("<III", 0x46546C67, 2, 12 + len(body)) + body)


def acc_info(g, i):
    a = g["accessors"][i]
    bv = g["bufferViews"][a["bufferView"]]
    nc = {"SCALAR": 1, "VEC2": 2, "VEC3": 3, "VEC4": 4, "MAT4": 16}[a["type"]]
    fmt = {5121: "B", 5123: "H", 5125: "I", 5126: "f"}[a["componentType"]]
    size = struct.calcsize(fmt) * nc
    stride = bv.get("byteStride") or size
    off = bv.get("byteOffset", 0) + a.get("byteOffset", 0)
    return a, bv, nc, fmt, stride, off


def read_acc(g, buf, i):
    a, bv, nc, fmt, stride, off = acc_info(g, i)
    return [struct.unpack_from("<%d%s" % (nc, fmt), buf, off + k * stride)
            for k in range(a["count"])]


def quat_mat(q):
    x, y, z, w = q
    return [[1 - 2 * (y * y + z * z), 2 * (x * y - z * w), 2 * (x * z + y * w)],
            [2 * (x * y + z * w), 1 - 2 * (x * x + z * z), 2 * (y * z - x * w)],
            [2 * (x * z - y * w), 2 * (y * z + x * w), 1 - 2 * (x * x + y * y)]]


def node_local(nd):
    t = nd.get("translation", [0, 0, 0])
    r = nd.get("rotation", [0, 0, 0, 1])
    s = nd.get("scale", [1, 1, 1])
    R = quat_mat(r)
    return [[R[i][j] * s[j] for j in range(3)] + [t[i]] for i in range(3)] + [[0, 0, 0, 1]]


def matmul(A, B):
    return [[sum(A[i][k] * B[k][j] for k in range(4)) for j in range(4)] for i in range(4)]


def global_mats(g):
    loc = [node_local(nd) for nd in g["nodes"]]
    parent = {}
    for i, nd in enumerate(g["nodes"]):
        for c in nd.get("children", []):
            parent[c] = i
    memo = {}

    def gm(i):
        if i in memo:
            return memo[i]
        m = loc[i]
        p = parent.get(i)
        if p is not None:
            m = matmul(gm(p), m)
        memo[i] = m
        return m

    return [gm(i) for i in range(len(g["nodes"]))]


# ---------------------------------------------------------------- main
argv = sys.argv[1:]
src = argv[0]
dst = argv[1] if len(argv) > 1 and not argv[1].startswith("--") else None
DIAG = "--diagnose" in argv
K = int(argv[argv.index("--k") + 1]) if "--k" in argv else 8

g, chunks = read_glb(src)
buf = next(p for ct, p in chunks if ct == 0x004E4942)
names = [nd.get("name", "") for nd in g["nodes"]]
GM = global_mats(g)

total_fixed = 0
for mi, mesh in enumerate(g.get("meshes", [])):
    for pi, prim in enumerate(mesh.get("primitives", [])):
        att = prim["attributes"]
        if "JOINTS_0" not in att:
            continue
        skin_idx = next((nd.get("skin") for nd in g["nodes"]
                         if nd.get("mesh") == mi and "skin" in nd), None)
        if skin_idx is None:
            continue
        joints = g["skins"][skin_idx]["joints"]
        jname = [names[j] for j in joints]
        jpos = [(GM[j][0][3], GM[j][1][3], GM[j][2][3]) for j in joints]

        P = read_acc(g, buf, att["POSITION"])
        J = [list(r) for r in read_acc(g, buf, att["JOINTS_0"])]
        Wr = read_acc(g, buf, att["WEIGHTS_0"])
        _, _, _, wfmt, _, _ = acc_info(g, att["WEIGHTS_0"])
        wsc = 1.0 if wfmt == "f" else (1 / 255 if wfmt == "B" else 1 / 65535)
        W = [[w * wsc for w in r] for r in Wr]

        is_fing = [any(t in n.lower() for t in FING) for n in jname]
        side = ["l" if n.lower().endswith("_l") else ("r" if n.lower().endswith("_r") else "")
                for n in jname]

        # ---- classify corrupt refs
        corrupt = defaultdict(list)     # vert -> [slot,...]
        for vi, (p, jrow, wrow) in enumerate(zip(P, J, W)):
            for s, (j, w) in enumerate(zip(jrow, wrow)):
                if w <= 0.001 or not is_fing[j]:
                    continue
                if (side[j] == "r" and p[0] > MID) or (side[j] == "l" and p[0] < -MID):
                    corrupt[vi].append(s)

        if not corrupt:
            print(f"  mesh[{mi}] prim{pi}: no cross-midline finger refs — nothing to do")
            continue

        bad_verts = set(corrupt)
        nref = sum(len(v) for v in corrupt.values())
        print(f"  mesh[{mi}] '{mesh.get('name','')}' prim{pi}: {len(P)} verts")
        print(f"    corrupt refs {nref} across {len(bad_verts)} verts")

        # ---- topology adjacency
        adj = defaultdict(set)
        if "indices" in prim:
            idx = [r[0] for r in read_acc(g, buf, prim["indices"])]
            for t in range(0, len(idx) - 2, 3):
                a_, b_, c_ = idx[t], idx[t + 1], idx[t + 2]
                adj[a_].update((b_, c_))
                adj[b_].update((a_, c_))
                adj[c_].update((a_, b_))

        # healthy = not corrupt AND has some weight
        def healthy(v):
            return v not in bad_verts and sum(W[v]) > 0.5

        # spatial fallback pool: healthy verts near the affected region
        cx = sum(P[v][0] for v in bad_verts) / len(bad_verts)
        cy = sum(P[v][1] for v in bad_verts) / len(bad_verts)
        cz = sum(P[v][2] for v in bad_verts) / len(bad_verts)
        pool = [v for v in range(len(P))
                if healthy(v) and abs(P[v][0] - cx) < 0.30
                and abs(P[v][1] - cy) < 0.30 and abs(P[v][2] - cz) < 0.30]
        print(f"    healthy donor pool near region: {len(pool)} verts")

        topo_used = spatial_used = 0
        newJ, newW = {}, {}

        for vi in sorted(bad_verts):
            # BFS out to healthy neighbours through the mesh
            found = []
            seen = {vi}
            q = deque([(vi, 0)])
            while q and len(found) < K:
                v, d = q.popleft()
                if d > 4:
                    continue
                for nb in adj.get(v, ()):
                    if nb in seen:
                        continue
                    seen.add(nb)
                    if healthy(nb):
                        found.append(nb)
                        if len(found) >= K:
                            break
                    q.append((nb, d + 1))
            if found:
                topo_used += 1
            else:
                # spatial fallback
                ds = sorted(((math.dist(P[vi], P[v]), v) for v in pool))[:K]
                found = [v for _, v in ds]
                spatial_used += 1

            # inverse-distance blend of neighbour weight sets
            accw = defaultdict(float)
            for nb in found:
                d = math.dist(P[vi], P[nb])
                wgt = 1.0 / max(d, 1e-4)
                for j, w in zip(J[nb], W[nb]):
                    if w > 0.001:
                        accw[j] += w * wgt
            # keep top 4, renormalise
            top = sorted(accw.items(), key=lambda kv: -kv[1])[:4]
            tot = sum(w for _, w in top)
            if tot <= 0:
                continue
            nj = [0, 0, 0, 0]
            nw = [0.0, 0.0, 0.0, 0.0]
            for s, (j, w) in enumerate(top):
                nj[s] = j
                nw[s] = w / tot
            newJ[vi] = nj
            newW[vi] = nw

        print(f"    repaired {len(newJ)} verts (topological {topo_used}, spatial fallback {spatial_used})")
        total_fixed += len(newJ)

        if DIAG:
            # show what the repair decided for a few verts
            for vi in sorted(newJ)[:6]:
                before = " ".join(f"{jname[j]}={w:.3f}" for j, w in zip(J[vi], W[vi]) if w > 0.001)
                after = " ".join(f"{jname[j]}={w:.3f}" for j, w in zip(newJ[vi], newW[vi]) if w > 0.001)
                print(f"      v{vi}\n        before: {before}\n        after : {after}")
            continue

        # ---- write back
        aJ, bvJ, ncJ, fmtJ, strideJ, offJ = acc_info(g, att["JOINTS_0"])
        aW, bvW, ncW, fmtW, strideW, offW = acc_info(g, att["WEIGHTS_0"])
        for vi in newJ:
            struct.pack_into("<4%s" % fmtJ, buf, offJ + vi * strideJ, *newJ[vi])
            if fmtW == "f":
                vals = newW[vi]
            elif fmtW == "B":
                vals = [max(0, min(255, int(round(w * 255)))) for w in newW[vi]]
                vals[0] += 255 - sum(vals)
            else:
                vals = [max(0, min(65535, int(round(w * 65535)))) for w in newW[vi]]
                vals[0] += 65535 - sum(vals)
            struct.pack_into("<4%s" % fmtW, buf, offW + vi * strideW, *vals)

if DIAG:
    print("\ndiagnose only — nothing written")
elif dst:
    write_glb(dst, g, chunks)
    print(f"\nwrote {dst}  ({total_fixed} verts repaired)")
else:
    print("\nno output path given — nothing written")