import { readFile, writeFile } from "node:fs/promises";
import path from "node:path";
import { Accessor, NodeIO } from "@gltf-transform/core";
import { Matrix3, Matrix4, Quaternion, Vector3 } from "three";

const [inputGlb, rigJsonFile, outputGlb, outputReport] = process.argv.slice(2);
if (!inputGlb || !rigJsonFile || !outputGlb || !outputReport) {
  throw new Error(
    "Usage: node inject-actor-rig.mjs <input.glb> <actor-rig.json> <output.glb> <report.json>",
  );
}

const rig = JSON.parse(await readFile(rigJsonFile, "utf8"));
if (rig.schemaVersion !== 2 || rig.status !== "animated") {
  throw new Error("Expected an animated schemaVersion 2 actor rig.");
}
const binary = await readFile(path.join(path.dirname(rigJsonFile), rig.binary));
const io = new NodeIO();
const document = await io.read(inputGlb);
const root = document.getRoot();
const buffer = root.listBuffers()[0] ?? document.createBuffer("actor-rig");

async function readObjMeshPositions(file) {
  const groups = new Map();
  let active = null;
  for (const line of (await readFile(file, "utf8")).split(/\r?\n/)) {
    if (line.startsWith("g ")) {
      active = line.slice(2).trim();
      if (!groups.has(active)) groups.set(active, []);
      continue;
    }
    if (!active || !line.startsWith("v ")) continue;
    const [x, y, z] = line.trim().split(/\s+/).slice(1).map(Number);
    if (![x, y, z].every(Number.isFinite)) throw new Error("Invalid OBJ position in " + active + ".");
    // Blender's OBJ importer stores Rune/OBJ Y-up coordinates in this local
    // accessor basis, then applies its +90-degree X node transform.
    groups.get(active).push([x, z, -y]);
  }
  return groups;
}

function positionBucket(position, tolerance) {
  return position.map((value) => Math.floor(value / tolerance)).join(",");
}

function mapTargetVertices(sourcePositions, targetPositions, meshName) {
  const tolerance = 1e-4;
  const buckets = new Map();
  for (let source = 0; source < sourcePositions.length; source += 1) {
    const key = positionBucket(sourcePositions[source], tolerance);
    if (!buckets.has(key)) buckets.set(key, []);
    buckets.get(key).push(source);
  }
  const mapping = new Uint32Array(targetPositions.length / 3);
  for (let target = 0; target < mapping.length; target += 1) {
    const position = Array.from(targetPositions.subarray(target * 3, target * 3 + 3));
    const base = position.map((value) => Math.floor(value / tolerance));
    let best = -1;
    let bestDistanceSquared = Number.POSITIVE_INFINITY;
    for (let dx = -1; dx <= 1; dx += 1) {
      for (let dy = -1; dy <= 1; dy += 1) {
        for (let dz = -1; dz <= 1; dz += 1) {
          const candidates = buckets.get([base[0] + dx, base[1] + dy, base[2] + dz].join(",")) ?? [];
          for (const source of candidates) {
            const candidate = sourcePositions[source];
            const distanceSquared = (position[0] - candidate[0]) ** 2
              + (position[1] - candidate[1]) ** 2
              + (position[2] - candidate[2]) ** 2;
            if (distanceSquared < bestDistanceSquared) {
              best = source;
              bestDistanceSquared = distanceSquared;
            }
          }
        }
      }
    }
    if (best < 0 || bestDistanceSquared > tolerance ** 2) {
      throw new Error(
        "Could not map Blender vertex " + target + " back to " + meshName
          + "; nearest accepted distance is " + tolerance + ".",
      );
    }
    mapping[target] = best;
  }
  return mapping;
}

const objMeshPositions = await readObjMeshPositions(path.join(path.dirname(rigJsonFile), "scene.obj"));

function typedBlock(descriptor, Type, bytesPerElement) {
  const start = descriptor.byteOffset;
  const end = start + descriptor.count * bytesPerElement;
  if (!Number.isInteger(start) || start < 0 || end > binary.byteLength) {
    throw new Error("Rig binary block is outside actor-rig.bin.");
  }
  return new Type(binary.buffer.slice(binary.byteOffset + start, binary.byteOffset + end));
}

function semanticClipName(name) {
  const normalized = String(name).trim().toLowerCase();
  if (normalized === "death") return "Death - " + name;
  if (normalized === "dead") return "Terminal Pose - " + name;
  if (normalized.startsWith("stand_")) return "Stand - " + name;
  if (normalized.includes("idle")) return "Idle - " + name;
  if (normalized.startsWith("walk")) return "Walk - " + name;
  if (normalized.startsWith("run")) return "Run - " + name;
  if (normalized.includes("attack")) return "Attack - " + name;
  // ACT motion libraries use both `casting*` and shorter combat-cast names.
  // Keep fishing_cast as a non-combat action, but expose cast01, cast_sp01,
  // continuous_cast, and their variants to the runtime cast selector.
  if (
    normalized.startsWith("casting")
    || normalized === "cast"
    || /^cast(?:\d|_)/.test(normalized)
    || normalized.startsWith("continuous_cast")
  ) return "Cast - " + name;
  if (normalized === "hurt") return "Wound - " + name;
  return "Action - " + name;
}

function deinterleave(source, frameCount, boneCount, boneIndex, itemSize) {
  const result = new Float32Array(frameCount * itemSize);
  for (let frame = 0; frame < frameCount; frame += 1) {
    const sourceOffset = (frame * boneCount + boneIndex) * itemSize;
    result.set(source.subarray(sourceOffset, sourceOffset + itemSize), frame * itemSize);
  }
  return result;
}

function isConstant(values, itemSize, tolerance = 1e-6) {
  for (let offset = itemSize; offset < values.length; offset += itemSize) {
    for (let component = 0; component < itemSize; component += 1) {
      if (Math.abs(values[offset + component] - values[component]) > tolerance) return false;
    }
  }
  return true;
}

const skinnedMeshNodes = [];
for (const sourceMesh of rig.meshes) {
  const candidates = root.listNodes().filter((node) => {
    const mesh = node.getMesh();
    if (!mesh) return false;
    return node.getName() === sourceMesh.name || mesh.getName() === sourceMesh.name;
  });
  if (candidates.length !== 1) {
    throw new Error(
      "Expected one GLB mesh named " + sourceMesh.name + "; found " + candidates.length + ".",
    );
  }
  const node = candidates[0];
  const primitives = node.getMesh().listPrimitives();
  if (primitives.length !== 1) {
    throw new Error("Skinned actor meshes must contain exactly one primitive.");
  }
  const sourceJoints = typedBlock(sourceMesh.joints, Uint16Array, 2);
  const sourceWeights = typedBlock(sourceMesh.weights, Float32Array, 4);
  const positionAccessor = primitives[0].getAttribute("POSITION");
  const targetVertexCount = positionAccessor.getCount();
  const sourcePositions = objMeshPositions.get(sourceMesh.name);
  if (!sourcePositions || sourcePositions.length !== sourceMesh.vertexCount) {
    throw new Error("OBJ source vertex accounting failed for " + sourceMesh.name + ".");
  }
  const joints = new Uint16Array(targetVertexCount * 4);
  const weights = new Float32Array(targetVertexCount * 4);
  const vertexMapping = mapTargetVertices(sourcePositions, positionAccessor.getArray(), sourceMesh.name);
  for (let vertex = 0; vertex < targetVertexCount; vertex += 1) {
    const sourceVertex = vertexMapping[vertex];
    joints.set(sourceJoints.subarray(sourceVertex * 4, sourceVertex * 4 + 4), vertex * 4);
    weights.set(sourceWeights.subarray(sourceVertex * 4, sourceVertex * 4 + 4), vertex * 4);
  }
  for (let vertex = 0; vertex < targetVertexCount; vertex += 1) {
    let total = 0;
    for (let influence = 0; influence < 4; influence += 1) {
      const offset = vertex * 4 + influence;
      if (!Number.isFinite(weights[offset]) || weights[offset] < 0) {
        throw new Error("Invalid skin weight in " + sourceMesh.name + ".");
      }
      if (weights[offset] <= 1e-8) {
        weights[offset] = 0;
        joints[offset] = 0;
      }
      if (weights[offset] > 0 && joints[offset] >= rig.skeleton.boneCount) {
        throw new Error("Skin joint index exceeds the exported GR2 skeleton.");
      }
      total += weights[offset];
    }
    if (total <= 1e-8) {
      joints[vertex * 4] = 0;
      weights[vertex * 4] = 1;
    } else {
      for (let influence = 0; influence < 4; influence += 1) {
        weights[vertex * 4 + influence] /= total;
      }
    }
  }
  const jointAccessor = document.createAccessor(sourceMesh.name + "-joints", buffer)
    .setType(Accessor.Type.VEC4)
    .setArray(joints);
  const weightAccessor = document.createAccessor(sourceMesh.name + "-weights", buffer)
    .setType(Accessor.Type.VEC4)
    .setArray(weights);
  primitives[0].setAttribute("JOINTS_0", jointAccessor);
  primitives[0].setAttribute("WEIGHTS_0", weightAccessor);
  skinnedMeshNodes.push(node);
}

if (!rig.skeleton?.bones?.length || !skinnedMeshNodes.length) {
  throw new Error("Animated rig is missing bones or skinned meshes.");
}
const bones = [...rig.skeleton.bones].sort((left, right) => left.index - right.index);
if (bones.some((bone, index) => bone.index !== index)) {
  throw new Error("GR2 bones must use a dense, ordered index space.");
}
const jointNodes = bones.map((bone) => document.createNode(bone.name || "bone-" + bone.index)
  .setTranslation(bone.translation)
  .setRotation(bone.rotation)
  .setScale(bone.scale));
const rigRoot = document.createNode("RuneWaker actor rig")
  .setTranslation(skinnedMeshNodes[0].getTranslation())
  .setRotation([0, 0, 0, 1])
  .setScale(skinnedMeshNodes[0].getScale());
for (const bone of bones) {
  if (bone.parent === -1) rigRoot.addChild(jointNodes[bone.index]);
  else {
    if (!jointNodes[bone.parent]) throw new Error("Invalid GR2 bone parent index.");
    jointNodes[bone.parent].addChild(jointNodes[bone.index]);
  }
}
const scene = root.listScenes()[0];
if (!scene) throw new Error("Actor GLB contains no scene.");
scene.addChild(rigRoot);
const meshBasisMatrix = new Matrix4().fromArray(skinnedMeshNodes[0].getMatrix());
for (const meshNode of skinnedMeshNodes) {
  const meshMatrix = new Matrix4().fromArray(meshNode.getMatrix());
  const transformDelta = meshMatrix.clone()
    .premultiply(meshBasisMatrix.clone().invert());
  const identity = new Matrix4().elements;
  if (transformDelta.elements.some((value, index) => Math.abs(value - identity[index]) > 1e-5)) {
    throw new Error("Skinned actor mesh nodes do not share one Blender basis transform.");
  }
  const normalMatrix = new Matrix3().getNormalMatrix(meshMatrix);
  const vector = new Vector3();
  for (const primitive of meshNode.getMesh().listPrimitives()) {
    const position = primitive.getAttribute("POSITION");
    const sourcePositions = position.getArray();
    const bakedPositions = new Float32Array(sourcePositions.length);
    for (let offset = 0; offset < sourcePositions.length; offset += 3) {
      vector.fromArray(sourcePositions, offset).applyMatrix4(meshMatrix).toArray(bakedPositions, offset);
    }
    position.setArray(bakedPositions);
    const normal = primitive.getAttribute("NORMAL");
    if (normal) {
      const sourceNormals = normal.getArray();
      const bakedNormals = new Float32Array(sourceNormals.length);
      for (let offset = 0; offset < sourceNormals.length; offset += 3) {
        vector.fromArray(sourceNormals, offset).applyMatrix3(normalMatrix)
          .normalize().toArray(bakedNormals, offset);
      }
      normal.setArray(bakedNormals);
    }
  }
  meshNode.setTranslation([0, 0, 0]).setRotation([0, 0, 0, 1]).setScale([1, 1, 1]);
}

const worldMatrices = new Array(bones.length);
function boneWorldMatrix(index) {
  if (worldMatrices[index]) return worldMatrices[index];
  const bone = bones[index];
  const local = new Matrix4().compose(
    new Vector3(...bone.translation),
    new Quaternion(...bone.rotation),
    new Vector3(...bone.scale),
  );
  worldMatrices[index] = bone.parent === -1
    ? local
    : boneWorldMatrix(bone.parent).clone().multiply(local);
  return worldMatrices[index];
}
const inverseBindValues = new Float32Array(bones.length * 16);
const rigRootMatrix = new Matrix4().compose(
  new Vector3(...rigRoot.getTranslation()),
  new Quaternion(...rigRoot.getRotation()),
  new Vector3(...rigRoot.getScale()),
);
for (const bone of bones) {
  rigRootMatrix.clone().multiply(boneWorldMatrix(bone.index)).invert()
    .toArray(inverseBindValues, bone.index * 16);
}
const inverseBindAccessor = document.createAccessor("RuneWaker inverse bind matrices", buffer)
  .setType(Accessor.Type.MAT4)
  .setArray(inverseBindValues);
const skin = document.createSkin("RuneWaker actor skin")
  .setInverseBindMatrices(inverseBindAccessor);
for (const jointNode of jointNodes) skin.addJoint(jointNode);
for (const meshNode of skinnedMeshNodes) meshNode.setSkin(skin);

const animationNames = [];
let animationChannelCount = 0;
for (const motion of rig.motions) {
  const clipName = semanticClipName(motion.name);
  const animation = document.createAnimation(clipName);
  const times = typedBlock(motion.times, Float32Array, 4);
  const timeAccessor = document.createAccessor(clipName + "-time", buffer)
    .setType(Accessor.Type.SCALAR)
    .setArray(times);
  const sourceByPath = {
    translation: typedBlock(motion.translations, Float32Array, 4),
    rotation: typedBlock(motion.rotations, Float32Array, 4),
    scale: typedBlock(motion.scales, Float32Array, 4),
  };
  const pathDetails = [
    ["translation", 3, Accessor.Type.VEC3],
    ["rotation", 4, Accessor.Type.VEC4],
    ["scale", 3, Accessor.Type.VEC3],
  ];
  for (const bone of bones) {
    for (const [targetPath, itemSize, accessorType] of pathDetails) {
      const values = deinterleave(
        sourceByPath[targetPath],
        motion.frameCount,
        bones.length,
        bone.index,
        itemSize,
      );
      if (Array.from(values).some((value) => !Number.isFinite(value))) {
        throw new Error("Non-finite " + targetPath + " sample in " + clipName + ".");
      }
      const rest = bone[targetPath];
      const restConstant = isConstant(values, itemSize)
        && rest.every((value, component) => Math.abs(values[component] - value) <= 1e-6);
      if (restConstant) continue;
      const outputAccessor = document
        .createAccessor(clipName + "-" + bone.index + "-" + targetPath, buffer)
        .setType(accessorType)
        .setArray(values);
      const sampler = document.createAnimationSampler()
        .setInput(timeAccessor)
        .setOutput(outputAccessor)
        .setInterpolation("LINEAR");
      const channel = document.createAnimationChannel()
        .setTargetNode(jointNodes[bone.index])
        .setTargetPath(targetPath)
        .setSampler(sampler);
      animation.addSampler(sampler).addChannel(channel);
      animationChannelCount += 1;
    }
  }
  if (!animation.listChannels().length) {
    animation.dispose();
    throw new Error("Original motion exported with no changing tracks: " + motion.name);
  }
  animationNames.push(clipName);
}

if (!animationNames.length) throw new Error("No native animation clips were injected.");
await io.write(outputGlb, document);
const report = {
  schemaVersion: 1,
  status: "animated-authentic-model",
  sourceRig: path.basename(rigJsonFile),
  bones: bones.length,
  skins: 1,
  skinnedMeshes: skinnedMeshNodes.length,
  animationChannels: animationChannelCount,
  animations: animationNames,
  warnings: [],
};
await writeFile(outputReport, JSON.stringify(report, null, 2) + "\n", "utf8");
console.log(JSON.stringify(report, null, 2));