"""Bake a hand-pose JSON into a body GLB by rewriting finger node rest rotations.
IBMs untouched -> mesh deforms to the pose. usage: bake_preview.py body.glb pose.json out.glb"""
import json, struct, sys
from pathlib import Path

body, posef, out = sys.argv[1:4]
data = Path(body).read_bytes()
length = struct.unpack_from("<I", data, 8)[0]
off = 12
chunks = []
gltf = None
while off < length:
    clen, ctype = struct.unpack_from("<II", data, off)
    off += 8
    if ctype == 0x4E4F534A:
        gltf = json.loads(data[off:off+clen].decode("utf-8"))
        chunks.append([ctype, None])
    else:
        chunks.append([ctype, data[off:off+clen]])
    off += clen

pose = json.loads(Path(posef).read_text())["bones"]
byname = {nd.get("name"): nd for nd in gltf["nodes"]}

canon = None
if len(sys.argv) > 4:  # canonical-rest GLB: apply pose as rest-relative delta
    cdata = Path(sys.argv[4]).read_bytes()
    clen2 = struct.unpack_from("<I", cdata, 12)[0]
    cg = json.loads(cdata[20:20+clen2].decode("utf-8"))
    canon = {nd.get("name"): nd.get("rotation", [0, 0, 0, 1]) for nd in cg["nodes"]}

def qmul(a, b):
    ax, ay, az, aw = a; bx, by, bz, bw = b
    return [aw*bx + ax*bw + ay*bz - az*by,
            aw*by - ax*bz + ay*bw + az*bx,
            aw*bz + ax*by - ay*bx + az*bw,
            aw*bw - ax*bx - ay*by - az*bz]

n = 0
for bone, quat in pose.items():
    if bone in byname:
        if canon is not None:
            cr = canon.get(bone, [0, 0, 0, 1])
            delta = qmul([-cr[0], -cr[1], -cr[2], cr[3]], quat)  # canon_rest^-1 * pose
            body_rest = byname[bone].get("rotation", [0, 0, 0, 1])
            quat = qmul(body_rest, delta)
        byname[bone]["rotation"] = quat
        n += 1
# strip animations so nothing overrides the pose
gltf.pop("animations", None)

js = json.dumps(gltf, separators=(",", ":")).encode("utf-8")
js += b" " * ((4 - len(js) % 4) % 4)
body_out = b""
for ctype, payload in chunks:
    if ctype == 0x4E4F534A:
        payload = js
    body_out += struct.pack("<II", len(payload), ctype) + payload
hdr = struct.pack("<III", 0x46546C67, 2, 12 + len(body_out))
Path(out).write_bytes(hdr + body_out)
print(f"baked {n}/{len(pose)} bones -> {out}")