#!/usr/bin/env python3
"""PROCITY R41 §41.1 — stage 2 of 2: per-clip skeleton-only GLBs -> <=6 GROUPED clip GLBs.

    python3 pipeline/clips_pack.py --stage <dir from clips_export.py> \
        --list pipeline/clips_r41.json --out web/models/clips \
        --manifest web/assets/motion_manifest.json

Pure python3 stdlib — no Blender, no deps. Operates at the glTF level, which is safe here because
every per-clip GLB in the staging dir carries the IDENTICAL 66-node mixamorig skeleton (asserted,
not assumed: `--check` fails loudly on any node-name mismatch). Merging is therefore: keep group
clip #1's node graph + skin verbatim, then append every other clip's animation with its channel
targets remapped by node NAME.

TRACK FILTER (`--tracks`, default `rot`).
PROCITY's consumer is `web/js/citizens/rigs.js` `_rotOnly`, which keeps ONLY `.quaternion` tracks and
additionally drops `Hips.quaternion` — every translation and scale track a clip GLB ships is decoded
by three.js and then thrown away on the same frame. Shipping them is pure boot cost. So the default
packs rotation channels only. `--tracks all` reproduces the unfiltered T/R/S file if a future round
needs root motion back (the pipeline is re-runnable from `clips_r41.json` in ~1 minute).

LOOP SEAM is measured, not asserted: for every clip we compute the maximum per-bone angle between
the first and last rotation keyframe (`loopSeamDeg`). Small = the clip cycles cleanly and can be
played on repeat; large = it is a one-shot and should be played once and blended out. The manifest
carries both the curated `loopable` flag and the measured number so Lane D can second-guess it.
"""
import argparse
import json
import math
import os
import re
import struct

COMP = {5120: ("b", 1), 5121: ("B", 1), 5122: ("h", 2), 5123: ("H", 2),
        5125: ("I", 4), 5126: ("f", 4)}
NCOMP = {"SCALAR": 1, "VEC2": 2, "VEC3": 3, "VEC4": 4, "MAT4": 16}


# ---------------------------------------------------------------- GLB read/write
def glb_read(path):
    with open(path, "rb") as f:
        data = f.read()
    if data[:4] != b"glTF":
        raise ValueError("%s: not a GLB" % path)
    js, bin_ = None, b""
    off = 12
    while off < len(data):
        ln, ty = struct.unpack_from("<II", data, off)
        off += 8
        if ty == 0x4E4F534A:
            js = json.loads(data[off:off + ln])
        elif ty == 0x004E4942:
            bin_ = data[off:off + ln]
        off += ln
    return js, bin_


def glb_write(path, js, bin_):
    j = json.dumps(js, separators=(",", ":")).encode()
    j += b" " * ((4 - len(j) % 4) % 4)
    b = bytes(bin_) + b"\0" * ((4 - len(bin_) % 4) % 4)
    total = 12 + 8 + len(j) + (8 + len(b) if b else 0)
    out = bytearray()
    out += struct.pack("<4sII", b"glTF", 2, total)
    out += struct.pack("<II", len(j), 0x4E4F534A) + j
    if b:
        out += struct.pack("<II", len(b), 0x004E4942) + b
    with open(path, "wb") as f:
        f.write(out)
    return total


def acc_bytes(js, bin_, ai):
    """Return (tightly packed bytes, count, type, componentType, normalized) for accessor ai."""
    a = js["accessors"][ai]
    ct, n = a["componentType"], a["count"]
    fmt, sz = COMP[ct]
    nc = NCOMP[a["type"]]
    elem = sz * nc
    bv = js["bufferViews"][a["bufferView"]]
    base = bv.get("byteOffset", 0) + a.get("byteOffset", 0)
    stride = bv.get("byteStride") or elem
    if stride == elem:
        raw = bin_[base:base + elem * n]
    else:
        raw = b"".join(bin_[base + i * stride: base + i * stride + elem] for i in range(n))
    return raw, n, a["type"], ct, a.get("normalized", False), a.get("min"), a.get("max")


def acc_floats(js, bin_, ai):
    raw, n, ty, ct, _, _, _ = acc_bytes(js, bin_, ai)
    fmt, sz = COMP[ct]
    nc = NCOMP[ty]
    v = struct.unpack("<%d%s" % (n * nc, fmt), raw)
    return [v[i * nc:(i + 1) * nc] for i in range(n)]


def quantize_rot(raw, count, ct):
    """float32 VEC4 quaternions -> int16 normalized VEC4 (glTF-legal animation output).

    Halves the dominant payload. glTF 2.0 dequantises SHORT as max(c/32767, -1.0) and three.js
    r175 does exactly that in GLTFLoader `_getArrayFromAccessor` (verified in
    web/vendor/addons/loaders/GLTFLoader.js). Worst-case angular error is ~0.006 deg — three
    orders of magnitude under anything a 1.7 m citizen can show on screen.
    """
    if ct != 5126:
        return raw, ct, False
    v = struct.unpack("<%df" % (count * 4), raw)
    q = [max(-32767, min(32767, int(round(max(-1.0, min(1.0, x)) * 32767.0)))) for x in v]
    return struct.pack("<%dh" % len(q), *q), 5122, True


class Sink:
    """Accumulates tightly packed accessor payloads into one BIN chunk, deduped by content."""

    def __init__(self):
        self.bin = bytearray()
        self.views = []
        self.accs = []
        self._seen = {}
        self.dedup_saved = 0

    def add(self, raw, count, ty, ct, normalized, mn, mx):
        key = (bytes(raw), count, ty, ct, normalized)
        if key in self._seen:                      # identical payload already in this file
            self.dedup_saved += len(raw)
            return self._seen[key]
        while len(self.bin) % 4:
            self.bin += b"\0"
        off = len(self.bin)
        self.bin += raw
        self.views.append({"buffer": 0, "byteOffset": off, "byteLength": len(raw)})
        a = {"bufferView": len(self.views) - 1, "componentType": ct,
             "count": count, "type": ty}
        if normalized:
            a["normalized"] = True
        if mn is not None:
            a["min"], a["max"] = mn, mx
        self.accs.append(a)
        self._seen[key] = len(self.accs) - 1
        return len(self.accs) - 1


# The bank is not uniform about the mixamorig namespace: 45 of the 46 R41 clips export as
# `mixamorig1:Bone`, one (Idle_Smoking) as `mixamorig:Bone` — the exact split the house `_canon`
# trick in rigs.js exists to absorb (and the same split already shipped: idle.glb is mixamorig4:,
# sit.glb is mixamorig:). Joint NAMES and node ORDER are identical across all 46 (asserted below).
# Normalise to the bare `mixamorig:` namespace so the merge is name-exact and every group file
# reads the same; three.js sanitises the colon out anyway, so runtime binding is unchanged.
_NS = re.compile(r"mixamorig\d*:")


def canon_name(n):
    return _NS.sub("mixamorig:", n) if n else n


def quat_angle_deg(q0, q1):
    d = abs(sum(a * b for a, b in zip(q0, q1)))
    return math.degrees(2.0 * math.acos(max(-1.0, min(1.0, d))))


# ---------------------------------------------------------------- pack one group
FINGER = re.compile(r"Hand(Thumb|Index|Middle|Ring|Pinky)\d")


def pack_group(gname, clips, stage, out_dir, tracks, quant=True, drop_fingers=False):
    keep = {"rotation"} if tracks == "rot" else {"rotation", "translation", "scale"}
    base_js, base_bin = glb_read(os.path.join(stage, "%s.glb" % clips[0]["id"]))
    for n in base_js["nodes"]:
        if n.get("name"):
            n["name"] = canon_name(n["name"])
    names = [n.get("name") for n in base_js["nodes"]]
    if len(set(names)) != len(names):
        raise SystemExit("%s: duplicate node names in base skeleton" % gname)
    idx_of = {n: i for i, n in enumerate(names)}

    sink = Sink()
    out = {
        "asset": {"version": "2.0",
                  "generator": "PROCITY pipeline/clips_pack.py (R41 41.1) from %s"
                               % base_js["asset"].get("generator", "?")},
        "scene": base_js.get("scene", 0),
        "scenes": base_js["scenes"],
        "nodes": base_js["nodes"],
        "animations": [],
        "buffers": [{"byteLength": 0}],
    }
    if "skins" in base_js:
        skins = json.loads(json.dumps(base_js["skins"]))
        for sk in skins:
            if "inverseBindMatrices" in sk:
                sk["inverseBindMatrices"] = sink.add(*acc_bytes(base_js, base_bin,
                                                                sk["inverseBindMatrices"]))
        out["skins"] = skins

    report = []
    for c in clips:
        p = os.path.join(stage, "%s.glb" % c["id"])
        js, bin_ = glb_read(p)
        nm = [canon_name(n.get("name")) for n in js["nodes"]]
        if nm != names:
            raise SystemExit("%s/%s: skeleton mismatch vs group base (%d vs %d nodes)"
                             % (gname, c["id"], len(nm), len(names)))
        if len(js.get("animations", [])) != 1:
            raise SystemExit("%s/%s: expected 1 animation, got %d"
                             % (gname, c["id"], len(js.get("animations", []))))
        anim = js["animations"][0]
        smap, chans, samps = {}, [], []
        dur, seam, nrot = 0.0, 0.0, 0
        for ch in anim["channels"]:
            path = ch["target"]["path"]
            if path not in keep or "node" not in ch["target"]:
                continue
            tname = nm[ch["target"]["node"]]
            if drop_fingers and FINGER.search(tname):
                continue
            si = ch["sampler"]
            if si not in smap:
                s = anim["samplers"][si]
                t = acc_floats(js, bin_, s["input"])
                dur = max(dur, t[-1][0] if t else 0.0)
                oraw, on, oty, oct, onorm, omn, omx = acc_bytes(js, bin_, s["output"])
                if path == "rotation":
                    q = acc_floats(js, bin_, s["output"])
                    if len(q) > 1:
                        seam = max(seam, quat_angle_deg(q[0], q[-1]))
                    nrot += 1
                    if quant and s.get("interpolation", "LINEAR") != "CUBICSPLINE":
                        oraw, oct, onorm = quantize_rot(oraw, on, oct)
                        omn = omx = None
                smap[si] = len(samps)
                samps.append({"input": sink.add(*acc_bytes(js, bin_, s["input"])),
                              "output": sink.add(oraw, on, oty, oct, onorm, omn, omx),
                              "interpolation": s.get("interpolation", "LINEAR")})
            chans.append({"sampler": smap[si],
                          "target": {"node": idx_of[tname], "path": path}})
        out["animations"].append({"name": c["id"], "channels": chans, "samplers": samps})
        report.append(dict(id=c["id"], duration=round(dur, 4), channels=len(chans),
                           samplers=len(samps), rot_bones=nrot, loopSeamDeg=round(seam, 2)))

    out["bufferViews"] = sink.views
    out["accessors"] = sink.accs
    out["buffers"][0]["byteLength"] = len(sink.bin) + ((4 - len(sink.bin) % 4) % 4)
    os.makedirs(out_dir, exist_ok=True)
    path = os.path.join(out_dir, "%s.glb" % gname)
    size = glb_write(path, out, sink.bin)
    return path, size, report, len(names), sink.dedup_saved


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--stage", required=True)
    ap.add_argument("--list", default="pipeline/clips_r41.json")
    ap.add_argument("--out", default="web/models/clips")
    ap.add_argument("--manifest", default="web/assets/motion_manifest.json")
    ap.add_argument("--tracks", choices=["rot", "all"], default="rot")
    ap.add_argument("--no-quantize", action="store_true",
                    help="keep float32 rotation outputs (default: int16 normalized, ~2x smaller)")
    ap.add_argument("--drop-fingers", action="store_true",
                    help="drop the 40 finger-bone channels (~half the payload; hands then hold "
                         "their bind pose). Off by default — fidelity kept, lever documented.")
    ap.add_argument("--loop-seam-max", type=float, default=25.0,
                    help="a clip curated loopable but seaming worse than this is demoted")
    a = ap.parse_args()

    sel = json.load(open(a.list))
    man = {
        "_schema": ("PROCITY motion manifest v1 (R41 41.1). groups{<file>} -> one GLB under "
                    "web/models/clips/, holding N named animations. clips{<clipId>} -> "
                    "{group,category,duration,loopable,source,mixamo,pair?}. Load a group ONCE "
                    "(gltf.animations[] are named by clipId); index by clips[id].group."),
        "_licence": ("Mixamo (Adobe) via ultra ~/Documents/MOTIONLIB/green/mixamo_full/. GREEN: "
                     "free to use inside games/projects, NO attribution required. Never "
                     "redistribute as a clip pack. R41 ships this licence regime only (ruling 2)."),
        "_tracks": ("rotation-only" if a.tracks == "rot" else "translation+rotation+scale"),
        "_encoding": ("rotation outputs are %s; %s"
                      % ("float32" if a.no_quantize else "int16 NORMALIZED (glTF-legal, three.js "
                         "r175 de-normalises in GLTFLoader._getArrayFromAccessor)",
                         "finger channels DROPPED" if a.drop_fingers else
                         "all 65 joints animated incl. fingers")),
        "_note": ("Skeleton-only, 0 tris, 66 nodes (65 mixamorig joints + Armature) — the same "
                  "byte-shape as web/models/peds/idle.glb. rigs.js `_canon` folds the mixamorigN: "
                  "prefix, `_rotOnly` keeps quaternions and drops Hips.quaternion."),
        "fps": 30,
        "groups": {},
        "clips": {},
    }
    total, dedup = 0, 0
    for gname, g in sel["groups"].items():
        path, size, rep, nnodes, dd = pack_group(gname, g["clips"], a.stage, a.out, a.tracks,
                                                 quant=not a.no_quantize,
                                                 drop_fingers=a.drop_fingers)
        total += size
        dedup += dd
        byid = {r["id"]: r for r in rep}
        man["groups"]["%s.glb" % gname] = {
            "category": g["category"], "desc": g["desc"], "clips": [c["id"] for c in g["clips"]],
            "bytes": size, "nodes": nnodes,
            "durationSum": round(sum(r["duration"] for r in rep), 3),
        }
        for c in g["clips"]:
            r = byid[c["id"]]
            loop = bool(c.get("loop")) and r["loopSeamDeg"] <= a.loop_seam_max
            e = {"group": "%s.glb" % gname, "category": g["category"],
                 "duration": r["duration"], "loopable": loop,
                 "loopSeamDeg": r["loopSeamDeg"],
                 "source": c["file"], "mixamo": c["mixamo"],
                 "channels": r["channels"], "rotBones": r["rot_bones"]}
            if c.get("loop") and not loop:
                e["loopDemoted"] = "curated loopable but loop seam %.1f deg > %.1f" % (
                    r["loopSeamDeg"], a.loop_seam_max)
            if c.get("pair"):
                e["pair"] = c["pair"]
            man["clips"][c["id"]] = e
        print("[pack] %-11s %2d clips  %7d B  seam max %.1f deg"
              % (gname, len(rep), size, max(r["loopSeamDeg"] for r in rep)))
    man["totalBytes"] = total
    man["clipCount"] = len(man["clips"])
    man["groupCount"] = len(man["groups"])
    man["dedupSavedBytes"] = dedup
    os.makedirs(os.path.dirname(a.manifest), exist_ok=True)
    with open(a.manifest, "w") as f:
        json.dump(man, f, indent=1)
        f.write("\n")
    print("\n%d clips in %d group GLBs, %d B total (%.2f MB) -> %s"
          % (man["clipCount"], man["groupCount"], total, total / 1048576.0, a.out))


if __name__ == "__main__":
    main()