import { CANONICAL_HIERARCHY, type CanonicalBone } from "@psx/shared";
import type { AnimationClip } from "../anim/procedural.js";
import { computeFlatNormals, vertexCount, type MeshData } from "../mesh/MeshBuilder.js";
import type { Skeleton } from "../rig/skeleton.js";

/**
 * Linear blend skinning on the CPU, so a clip can be inspected offscreen.
 *
 * Skin weights are only wrong in ways you can see once the skeleton moves —
 * a heel split 50/50 between shin and foot looks perfect in bind pose and
 * shears into a wedge the moment the ankle pitches. Rendering a posed mesh is
 * the only way to catch that without launching the game.
 *
 * Matches the runtime's convention: rest rotations are identity, rest
 * translations are the joint's local offset, and the inverse bind matrix is a
 * translation by the negated rest world position.
 */

/** Column-major 4x4, matching glTF. */
type Mat4 = Float32Array;

function identity(): Mat4 {
  const m = new Float32Array(16);
  m[0] = 1;
  m[5] = 1;
  m[10] = 1;
  m[15] = 1;
  return m;
}

function multiply(a: Mat4, b: Mat4): Mat4 {
  const out = new Float32Array(16);
  for (let col = 0; col < 4; col++) {
    for (let row = 0; row < 4; row++) {
      let sum = 0;
      for (let k = 0; k < 4; k++) sum += a[k * 4 + row]! * b[col * 4 + k]!;
      out[col * 4 + row] = sum;
    }
  }
  return out;
}

/** Translation composed with an Euler XYZ rotation — the clips' channel form. */
function trs(translation: readonly number[], euler: readonly [number, number, number]): Mat4 {
  const [x, y, z] = euler;
  const cx = Math.cos(x);
  const sx = Math.sin(x);
  const cy = Math.cos(y);
  const sy = Math.sin(y);
  const cz = Math.cos(z);
  const sz = Math.sin(z);

  // R = Rz * Ry * Rx, matching eulerToQuaternion in anim/procedural.ts.
  const m = identity();
  m[0] = cy * cz;
  m[1] = cy * sz;
  m[2] = -sy;
  m[4] = sx * sy * cz - cx * sz;
  m[5] = sx * sy * sz + cx * cz;
  m[6] = sx * cy;
  m[8] = cx * sy * cz + sx * sz;
  m[9] = cx * sy * sz - sx * cz;
  m[10] = cx * cy;
  m[12] = translation[0]!;
  m[13] = translation[1]!;
  m[14] = translation[2]!;
  return m;
}

/** Sample a channel's keyframes, lerping the Euler triples. */
function sampleEuler(clip: AnimationClip, bone: CanonicalBone, time: number): [number, number, number] {
  const channel = clip.channels.find((c) => c.bone === bone);
  if (!channel || channel.keyframes.length === 0) return [0, 0, 0];
  const keys = channel.keyframes;

  if (time <= keys[0]!.time) return keys[0]!.rotation;
  const last = keys[keys.length - 1]!;
  if (time >= last.time) return last.rotation;

  for (let i = 0; i < keys.length - 1; i++) {
    const a = keys[i]!;
    const b = keys[i + 1]!;
    if (time > b.time) continue;
    const span = b.time - a.time;
    const f = span > 1e-6 ? (time - a.time) / span : 0;
    return [
      a.rotation[0] + (b.rotation[0] - a.rotation[0]) * f,
      a.rotation[1] + (b.rotation[1] - a.rotation[1]) * f,
      a.rotation[2] + (b.rotation[2] - a.rotation[2]) * f,
    ];
  }
  return last.rotation;
}

/** Skinning matrices, one per joint, in the skeleton's own joint order. */
export function poseMatrices(skeleton: Skeleton, clip: AnimationClip, time: number): Mat4[] {
  const world: Mat4[] = [];
  const skin: Mat4[] = [];

  for (const [i, joint] of skeleton.joints.entries()) {
    const local = trs(joint.local, sampleEuler(clip, joint.name, time));
    const parentName = CANONICAL_HIERARCHY[joint.name];
    const parentIndex = parentName === null ? -1 : (skeleton.index.get(parentName) ?? -1);
    world[i] = parentIndex >= 0 ? multiply(world[parentIndex]!, local) : local;

    // Inverse bind is a pure translation by the negated rest world position.
    const inverseBind = identity();
    inverseBind[12] = -joint.world[0];
    inverseBind[13] = -joint.world[1];
    inverseBind[14] = -joint.world[2];
    skin[i] = multiply(world[i]!, inverseBind);
  }
  return skin;
}

/** A copy of `mesh` deformed into the clip's pose at `time`. */
export function poseMesh(
  mesh: MeshData,
  skeleton: Skeleton,
  clip: AnimationClip,
  time: number,
): MeshData {
  const skin = poseMatrices(skeleton, clip, time);
  const count = vertexCount(mesh);
  const posed: MeshData = { ...mesh, positions: new Array<number>(count * 3).fill(0) };

  for (let v = 0; v < count; v++) {
    const x = mesh.positions[v * 3]!;
    const y = mesh.positions[v * 3 + 1]!;
    const z = mesh.positions[v * 3 + 2]!;
    let ox = 0;
    let oy = 0;
    let oz = 0;
    let total = 0;

    for (let i = 0; i < 4; i++) {
      const weight = mesh.weights[v * 4 + i] ?? 0;
      if (weight <= 0) continue;
      const m = skin[mesh.joints[v * 4 + i]!];
      if (!m) continue;
      total += weight;
      ox += weight * (m[0]! * x + m[4]! * y + m[8]! * z + m[12]!);
      oy += weight * (m[1]! * x + m[5]! * y + m[9]! * z + m[13]!);
      oz += weight * (m[2]! * x + m[6]! * y + m[10]! * z + m[14]!);
    }

    if (total <= 0) {
      ox = x;
      oy = y;
      oz = z;
    }
    posed.positions[v * 3] = ox;
    posed.positions[v * 3 + 1] = oy;
    posed.positions[v * 3 + 2] = oz;
  }

  posed.normals = [];
  computeFlatNormals(posed);
  return posed;
}
