"""Post-pass on a Blender-exported binary FBX: remove the duplicate shape-key animation that Blender
writes as Geometry-property curves (AnimationCurveNode named after the shape key, connected OP to the
Geometry). The DeformPercent curves on the BlendShapeChannel deformers are kept. Result loads in
three.js FBXLoader (which otherwise throws on the duplicates) and stays valid for Blender/Unity.
usage: python3 fbx_clean.py in.fbx out.fbx
"""
import struct, sys

def read_fbx(data):
    ver = struct.unpack('<I', data[23:27])[0]; big = ver >= 7500
    def read_node(pos):
        if big:
            end, nprops, plen, nlen = struct.unpack_from('<QQQB', data, pos); pos += 25
        else:
            end, nprops, plen, nlen = struct.unpack_from('<IIIB', data, pos); pos += 13
        if end == 0: return None, pos
        name = data[pos:pos+nlen]; pos += nlen
        props = []
        for _ in range(nprops):
            t = data[pos:pos+1]; pos += 1
            if t in b'YCIFDL':
                sz = {b'Y':2,b'C':1,b'I':4,b'F':4,b'D':8,b'L':8}[t]
                props.append((t, data[pos:pos+sz])); pos += sz
            elif t in b'SR':
                ln = struct.unpack_from('<I', data, pos)[0]; pos += 4
                props.append((t, data[pos:pos+ln])); pos += ln
            else:
                alen, enc, clen = struct.unpack_from('<III', data, pos)
                props.append((t, data[pos:pos+12+clen])); pos += 12 + clen
        children = []
        while pos < end:
            c, pos = read_node(pos)
            if c is None: break
            children.append(c)
        return [name, props, children], end
    pos = 27; nodes = []
    while pos < len(data):
        n, pos = read_node(pos)
        if n is None: break
        nodes.append(n)
    footer = data[pos:]
    return ver, big, nodes, footer

def write_fbx(ver, big, nodes, footer):
    out = bytearray(b'Kaydara FBX Binary  \x00\x1a\x00' + struct.pack('<I', ver))
    def null_rec():
        return b'\x00' * (25 if big else 13)
    def write_node(n):
        name, props, children = n
        start = len(out)
        out.extend(b'\x00' * (25 if big else 13))  # header placeholder
        out.extend(name)
        pstart = len(out)
        for t, raw in props:
            out.extend(t)
            if t in b'SR': out.extend(struct.pack('<I', len(raw)))
            out.extend(raw)
        plen = len(out) - pstart
        if children:
            for c in children: write_node(c)
            out.extend(null_rec())
        end = len(out)
        hdr = struct.pack('<QQQB', end, len(props), plen, len(name)) if big else struct.pack('<IIIB', end, len(props), plen, len(name))
        out[start:start+len(hdr)] = hdr
    for n in nodes: write_node(n)
    out.extend(null_rec())
    out.extend(footer)
    return bytes(out)

def pid(n):  # first int64 prop
    t, raw = n[1][0]; return struct.unpack('<q', raw)[0]

def clean(data):
    ver, big, nodes, footer = read_fbx(data)
    objs = next(n for n in nodes if n[0] == b'Objects'); conns = next(n for n in nodes if n[0] == b'Connections')
    geom_ids = {pid(o) for o in objs[2] if o[0] == b'Geometry'}
    # curve nodes connected (OP) to a Geometry → duplicates
    dup_nodes = set()
    for c in conns[2]:
        p = c[1]
        if len(p) >= 4 and p[0][1] == b'OP':
            child, parent = struct.unpack('<q', p[1][1])[0], struct.unpack('<q', p[2][1])[0]
            if parent in geom_ids:
                dup_nodes.add(child)
    acn = {pid(o) for o in objs[2] if o[0] == b'AnimationCurveNode'}
    dup_nodes &= acn
    # curves whose parent is a dup node
    dup_curves = set()
    for c in conns[2]:
        p = c[1]
        if len(p) >= 3:
            child, parent = struct.unpack('<q', p[1][1])[0], struct.unpack('<q', p[2][1])[0]
            if parent in dup_nodes: dup_curves.add(child)
    kill = dup_nodes | dup_curves
    objs[2] = [o for o in objs[2] if not (o[0] in (b'AnimationCurveNode', b'AnimationCurve') and pid(o) in kill)]
    def keep_conn(c):
        p = c[1]
        if len(p) < 3: return True
        a, b = struct.unpack('<q', p[1][1])[0], struct.unpack('<q', p[2][1])[0]
        return a not in kill and b not in kill
    conns[2] = [c for c in conns[2] if keep_conn(c)]
    # fix Definitions counts (Count fields) — optional for most importers; update AnimationCurveNode/AnimationCurve counts
    defs = next((n for n in nodes if n[0] == b'Definitions'), None)
    if defs:
        counts = {}
        for o in objs[2]: counts[o[0]] = counts.get(o[0], 0) + 1
        for ot in defs[2]:
            if ot[0] == b'ObjectType' and ot[1] and ot[1][0][1] in counts:
                for ch in ot[2]:
                    if ch[0] == b'Count': ch[1] = [(b'I', struct.pack('<i', counts[ot[1][0][1]]))]
    return write_fbx(ver, big, nodes, footer), len(dup_nodes), len(dup_curves)

if __name__ == '__main__':
    src, dst = sys.argv[1], sys.argv[2]
    out, nn, nc = clean(open(src, 'rb').read())
    open(dst, 'wb').write(out)
    print(f"[fbx_clean] removed {nn} duplicate curve nodes, {nc} curves; {len(out)} bytes")