#!/usr/bin/env node
// Bake a combined "Run_Punch" animation: legs from Jog_Fwd_Loop + upper body from Punch_Jab.
//
// No Blender required — edits the glTF JSON chunk only. The new animation reuses the source
// GLB's accessors/buffer verbatim (same skeleton); arm-lift and wind-up append fresh keyframes
// to the binary chunk.
//
//   --mode arms-only  : spine+neck+head+legs from RUN (torso upright, gaze forward), arms from PUNCH
//   --arm-lift DEG    : pitch both upperarms up (Y axis, mirrored per arm) → fist at eye level
//   --windup DEG      : synthesize a wind-up on the PUNCHING (left) upperarm: neutral → cocked
//                       (pulled back DEG on Z axis) → extended → recovered. Adds anticipation.
//   --windup-dur S    : wind-up (pull-back) duration in seconds (default 0.18)
//
// Usage:  node tools/bake_run_punch.mjs --mode arms-only --arm-lift 35 --windup 45

import fs from "node:fs";
import path from "node:path";

const argv = process.argv.slice(2);
const getOpt = (name, fallback) => {
  const i = argv.indexOf(name);
  return i >= 0 ? argv[i + 1] : fallback;
};
const mode = getOpt("--mode", "arms-only");
const armLiftDeg = parseFloat(getOpt("--arm-lift", "0"));
const windupDeg = parseFloat(getOpt("--windup", "0"));
const windupDur = parseFloat(getOpt("--windup-dur", "0.18"));

const SRC = "assets/quaternius/anim-lib-1/Universal Animation Library[Standard]/Unreal-Godot/UAL1_Standard.glb";
const OUT = "assets/quaternius/anim-combined/Run_Punch.glb";
const RUN_NAME = "Jog_Fwd_Loop";
const PUNCH_NAME = "Punch_Jab";
const OUT_CLIP = "Run_Punch";
const THROW_DUR = 0.15;   // cocked → extended (the fast thrust)

function parseGlb(buf) {
  if (buf.toString("ascii", 0, 4) !== "glTF") throw new Error("not a GLB");
  const length = buf.readUInt32LE(8);
  let o = 12, json = null, bin = null;
  while (o < length) {
    const clen = buf.readUInt32LE(o);
    const ctype = buf.toString("ascii", o + 4, o + 8);
    const data = buf.subarray(o + 8, o + 8 + clen);
    if (ctype === "JSON") json = JSON.parse(data.toString("utf8"));
    else if (ctype.startsWith("BIN")) bin = Buffer.from(data);
    o += 8 + clen;
  }
  return { json, bin };
}

function writeGlb(jsonObj, binBuf, outPath) {
  const jsonBytes = Buffer.from(JSON.stringify(jsonObj), "utf8");
  const jsonPad = (4 - (jsonBytes.length % 4)) % 4;
  const jsonPadded = Buffer.concat([jsonBytes, Buffer.alloc(jsonPad, 0x20)]);
  const binPad = (4 - (binBuf.length % 4)) % 4;
  const binPadded = Buffer.concat([binBuf, Buffer.alloc(binPad, 0x00)]);
  const total = 12 + 8 + jsonPadded.length + 8 + binPadded.length;
  const out = Buffer.alloc(total);
  let p = 0;
  out.write("glTF", p, "ascii"); p += 4;
  out.writeUInt32LE(2, p); p += 4;
  out.writeUInt32LE(total, p); p += 4;
  out.writeUInt32LE(jsonPadded.length, p); p += 4; out.write("JSON", p, "ascii"); p += 4;
  jsonPadded.copy(out, p); p += jsonPadded.length;
  out.writeUInt32LE(binPadded.length, p); p += 4; out.write("BIN\0", p, "ascii"); p += 4;
  binPadded.copy(out, p); p += binPadded.length;
  fs.mkdirSync(path.dirname(outPath), { recursive: true });
  fs.writeFileSync(outPath, out);
  return total;
}

// ── Quaternion + vector math (xyzw) ──
function qMul(a, b) {
  const [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];
}
function qAxis(axis, deg) {
  const sign = deg < 0 ? -1 : 1;
  const r = Math.abs(deg) * Math.PI / 180, h = Math.sin(r / 2) * sign;
  const v = [0, 0, 0]; v[axis] = h;
  return [v[0], v[1], v[2], Math.cos(r / 2)];
}
function qSlerp(a, b, t) {
  let [ax, ay, az, aw] = a;
  let [bx, by, bz, bw] = b;
  let dot = ax * bx + ay * by + az * bz + aw * bw;
  if (dot < 0) { bx = -bx; by = -by; bz = -bz; bw = -bw; dot = -dot; }
  if (dot > 0.9995) {
    return [ax + t * (bx - ax), ay + t * (by - ay), az + t * (bz - az), aw + t * (bw - aw)].map(v => v);
  }
  const theta = Math.acos(dot), sinT = Math.sin(theta);
  const s0 = Math.sin((1 - t) * theta) / sinT, s1 = Math.sin(t * theta) / sinT;
  return [s0 * ax + s1 * bx, s0 * ay + s1 * by, s0 * az + s1 * bz, s0 * aw + s1 * bw];
}
// Lift quaternion per arm side (Y axis, mirrored: right +, left −). Raises the fist.
const qLiftFor = (bone, deg) => bone.endsWith("_l") ? qAxis(1, -deg) : qAxis(1, deg);
// Retract (wind-up) quaternion per arm side (Z axis, mirrored). Pulls the fist backward.
const qRetractFor = (bone, deg) => bone.endsWith("_l") ? qAxis(2, -deg) : qAxis(2, deg);

const isLowerBody = b => !b ? false : b === "root" || b === "pelvis" || /thigh|calf|foot|ball/i.test(b);
const isArm = b => !b ? false : /clavicle|upperarm|lowerarm|^hand_|hand_l|hand_r|index_|middle_|pinky_|ring_|thumb_/i.test(b);
const fromRun = mode === "arms-only" ? (b => !isArm(b))
              : mode === "no-head" ? (b => isLowerBody(b) || b === "neck_01" || b === "Head")
              : isLowerBody;

const { json, bin } = parseGlb(fs.readFileSync(SRC));
const nodes = json.nodes;
const anims = json.animations;
const accessors = json.accessors;
const bufferViews = json.bufferViews;
const run = anims.find(a => (a.name || "") === RUN_NAME);
const punch = anims.find(a => (a.name || "") === PUNCH_NAME);
if (!run || !punch) throw new Error(`source clips not found`);

// parent map (glTF uses children arrays)
const parentOf = new Array(nodes.length).fill(-1);
for (let i = 0; i < nodes.length; i++) if (nodes[i].children) for (const c of nodes[i].children) parentOf[c] = i;
const ni = nm => nodes.findIndex(n => (n.name || "") === nm);

const newSamplers = [];
const newChannels = [];
const runBones = new Set(), punchBones = new Set();

function addChannels(srcAnim, include, set) {
  for (const ch of srcAnim.channels) {
    const bone = nodes[ch.target.node]?.name;
    if (!include(bone)) continue;
    set.add(bone);
    const ss = srcAnim.samplers[ch.sampler];
    newSamplers.push({ input: ss.input, output: ss.output, interpolation: ss.interpolation });
    newChannels.push({ sampler: newSamplers.length - 1, target: { node: ch.target.node, path: ch.target.path } });
  }
}
addChannels(run, fromRun, runBones);
addChannels(punch, b => !fromRun(b), punchBones);

// ── Binary-append helpers (append float data + create bufferView + accessor) ──
let outBin = bin;
function appendFloats(floats, type, count) {
  const comps = type === "SCALAR" ? 1 : type === "VEC3" ? 3 : 4;
  const byteOffset = outBin.length;
  const buf = Buffer.alloc(count * comps * 4);
  for (let i = 0; i < floats.length; i++) buf.writeFloatLE(floats[i], i * 4);
  outBin = Buffer.concat([outBin, buf]);
  const bvIdx = bufferViews.length;
  bufferViews.push({ buffer: 0, byteOffset, byteLength: count * comps * 4 });
  const accIdx = accessors.length;
  const min = new Array(comps).fill(Infinity), max = new Array(comps).fill(-Infinity);
  for (let i = 0; i < count; i++) for (let c = 0; c < comps; c++) {
    const v = floats[i * comps + c];
    if (v < min[c]) min[c] = v; if (v > max[c]) max[c] = v;
  }
  accessors.push({ bufferView: bvIdx, componentType: 5126, count, type, min, max });
  return accIdx;
}

// ── --arm-lift: pitch both upperarms up (mirrored Y) ──
if (armLiftDeg !== 0) {
  for (const ch of newChannels) {
    const bone = nodes[ch.target.node]?.name;
    if ((bone !== "upperarm_l" && bone !== "upperarm_r") || ch.target.path !== "rotation") continue;
    const sampler = newSamplers[ch.sampler];
    const outAcc = accessors[sampler.output];
    const bv = bufferViews[outAcc.bufferView];
    const off = (bv.byteOffset || 0) + (outAcc.byteOffset || 0);
    const count = outAcc.count;
    const lift = qLiftFor(bone, armLiftDeg);
    const floats = [];
    for (let i = 0; i < count; i++) {
      const q = [bin.readFloatLE(off + (i * 4) * 4), bin.readFloatLE(off + (i * 4 + 1) * 4),
                 bin.readFloatLE(off + (i * 4 + 2) * 4), bin.readFloatLE(off + (i * 4 + 3) * 4)];
      const r = qMul(lift, q);
      floats.push(...r);
    }
    sampler.output = appendFloats(floats, "VEC4", count);
  }
}

// ── --windup: synthesize neutral → cocked → extended → recovered on the punching (left) upperarm ──
let windupApplied = false;
if (windupDeg !== 0) {
  const bone = "upperarm_l";
  const boneIdx = ni(bone);
  const ch = newChannels.find(c => c.target.node === boneIdx && c.target.path === "rotation");
  if (ch) {
    const sampler = newSamplers[ch.sampler];
    const outAcc = accessors[sampler.output];
    const inAcc = accessors[sampler.input];
    const bvO = bufferViews[outAcc.bufferView], bvI = bufferViews[inAcc.bufferView];
    const offO = (bvO.byteOffset || 0) + (outAcc.byteOffset || 0);
    const offI = (bvI.byteOffset || 0) + (inAcc.byteOffset || 0);
    const count = outAcc.count;
    const rots = [], times = [];
    for (let i = 0; i < count; i++) {
      rots.push([outBin.readFloatLE(offO + (i * 4) * 4), outBin.readFloatLE(offO + (i * 4 + 1) * 4),
                 outBin.readFloatLE(offO + (i * 4 + 2) * 4), outBin.readFloatLE(offO + (i * 4 + 3) * 4)]);
      times.push(outBin.readFloatLE(offI + i * 4));
    }
    // Note: rots/times read from outBin (post-lift) so the lift is preserved in neutral/extended/recover.
    const L = times[times.length - 1];          // clip length (~0.93s, matches run)
    const neutral = rots[0];
    const recovered = rots[rots.length - 1];
    // find the extension key = max forward (Z) hand via FK on the left arm chain
    const chain = [];
    { let i = boneIdx; while (i >= 0) { chain.unshift(i); i = parentOf[i]; } }
    function handZ(upperarmRot) {
      let M = [[1,0,0,0],[0,1,0,0],[0,0,1,0],[0,0,0,1]];
      for (const idx of chain) {
        const nm = nodes[idx].name;
        let rot = nodes[idx].rotation ? [...nodes[idx].rotation] : [0,0,0,1];
        if (nm === bone && upperarmRot) rot = upperarmRot;
        else {
          const pch = punch.channels.find(c => c.target.node === idx && c.target.path === "rotation");
          if (pch) rot = rots0(pch); // falls back below
        }
        const r = rot, t = nodes[idx].translation || [0,0,0];
        const q2m = q => { const[x,y,z,w]=q,x2=x*x,y2=y*y,z2=z*z,xy=x*y,xz=x*z,yz=y*z,wx=w*x,wy=w*y,wz=w*z;
          return[[1-2*(y2+z2),2*(xy-wz),2*(xz+wy)],[2*(xy+wz),1-2*(x2+z2),2*(yz-wx)],[2*(xz-wy),2*(yz+wx),1-2*(x2+y2)]]; };
        const m = q2m(r);
        M = mm(M, [[m[0][0],m[0][1],m[0][2],t[0]],[m[1][0],m[1][1],m[1][2],t[1]],[m[2][0],m[2][1],m[2][2],t[2]],[0,0,0,1]]);
      }
      return M[2][3]; // hand world Z (forward)
    }
    function mm(a,b){const r=[[0,0,0,0],[0,0,0,0],[0,0,0,0],[0,0,0,0]];for(let i=0;i<4;i++)for(let j=0;j<4;j++){let s=0;for(let k=0;k<4;k++)s+=a[i][k]*b[k][j];r[i][j]=s;}return r;}
    function rots0(pch){return null;} // placeholder — replaced by per-key lookup below
    // Proper FK: find extension key by scanning punch upperarm keys
    let extK = 0, extZ = -Infinity;
    for (let k = 0; k < count; k++) {
      // build chain with this key's upperarm rotation
      let M = [[1,0,0,0],[0,1,0,0],[0,0,1,0],[0,0,0,1]];
      for (const idx of chain) {
        const nm = nodes[idx].name;
        let rot = nodes[idx].rotation ? [...nodes[idx].rotation] : [0,0,0,1];
        if (nm === bone) rot = rots[k];
        else {
          const pch = punch.channels.find(c => c.target.node === idx && c.target.path === "rotation");
          if (pch) {
            const ps = punch.samplers[pch.sampler];
            const pa = accessors[ps.output]; const pbv = bufferViews[pa.bufferView];
            const poff = (pbv.byteOffset||0)+(pa.byteOffset||0);
            rot = [bin.readFloatLE(poff+(k*4)*4),bin.readFloatLE(poff+(k*4+1)*4),bin.readFloatLE(poff+(k*4+2)*4),bin.readFloatLE(poff+(k*4+3)*4)];
          }
        }
        const q2m = q => { const[x,y,z,w]=q,x2=x*x,y2=y*y,z2=z*z,xy=x*y,xz=x*z,yz=y*z,wx=w*x,wy=w*y,wz=w*z;
          return[[1-2*(y2+z2),2*(xy-wz),2*(xz+wy)],[2*(xy+wz),1-2*(x2+z2),2*(yz-wx)],[2*(xz-wy),2*(yz+wx),1-2*(x2+y2)]]; };
        const m = q2m(rot), t = nodes[idx].translation||[0,0,0];
        M = mm(M, [[m[0][0],m[0][1],m[0][2],t[0]],[m[1][0],m[1][1],m[1][2],t[1]],[m[2][0],m[2][1],m[2][2],t[2]],[0,0,0,1]]);
      }
      const z = M[2][3];
      if (z > extZ) { extZ = z; extK = k; }
    }
    const extended = rots[extK];
    // cocked = neutral pulled back (retract in bone-local space, pre-multiply)
    const retract = qRetractFor(bone, windupDeg);
    const cocked = qMul(retract, neutral);
    // 4-key synthesized curve: neutral@0 → cocked@windUpDur → extended@(+throw) → recovered@L
    const extT = windupDur + THROW_DUR;
    const sTimes = [0, windupDur, extT, L];
    const sRots = [neutral, cocked, extended, recovered];
    const floats = [];
    for (const q of sRots) floats.push(...q);
    sampler.input = appendFloats(sTimes, "SCALAR", 4);
    sampler.output = appendFloats(floats, "VEC4", 4);
    sampler.interpolation = "LINEAR";
    windupApplied = true;
    console.log(`  wind-up: punching arm '${bone}' → neutral@0 → cocked@${windupDur}s → extended@${extT}s (key ${extK}) → recovered@${L.toFixed(3)}s  (retract Z ${windupDeg}°)`);
  }
}

json.buffers[0].byteLength = outBin.length;
const runPunch = { name: OUT_CLIP, channels: newChannels, samplers: newSamplers };
const outJson = { ...json, animations: [runPunch] };
const bytes = writeGlb(outJson, outBin, OUT);

const combinedLen = Math.max(...newSamplers.map(s => accessors[s.input]?.max?.[0] ?? 0));
console.log(`baked ${OUT}  [mode=${mode}, arm-lift=${armLiftDeg}°, windup=${windupDeg}°]`);
console.log(`  from RUN(${runBones.size}) + PUNCH(${punchBones.size}) = ${runBones.size + punchBones.size} bones, len=${combinedLen.toFixed(3)}s, ${(bytes/1024/1024).toFixed(2)}MB`);