import { Bone, Group, Quaternion, Skeleton, Vector3 } from 'three';

// Three.js integration; the model and dense inference runtime are engine-independent.

export function validateRig(rig) {
  const n = rig.names?.length;
  if (!Number.isInteger(n) || n < 1 || new Set(rig.names).size !== n ||
      rig.parents?.length !== n || rig.offsets?.length !== n || rig.end_sites?.length !== n) {
    throw new Error('Rig requires unique names and equally sized hierarchy, offset and end-site arrays.');
  }
  if (rig.units !== 'metres' || rig.up_axis !== 'Y' || rig.handedness !== 'right') {
    throw new Error('Convert the rig to right-handed Y-up metres before binding.');
  }
  for (let i = 0; i < n; i++) {
    const p = rig.parents[i];
    if (!Number.isInteger(p) || (i === 0 ? p !== -1 : p < 0 || p >= i)) {
      throw new Error(`Invalid parent for joint ${rig.names[i]}.`);
    }
    if (typeof rig.end_sites[i] !== 'boolean' || rig.offsets[i]?.length !== 3 ||
        !rig.offsets[i].every(Number.isFinite)) throw new Error(`Invalid joint ${rig.names[i]}.`);
  }
  return rig;
}

export function createSourceRig(rig) {
  validateRig(rig);
  const group = new Group();
  group.name = 'SuperHQMotion';
  const bones = rig.names.map((name, i) => {
    const bone = new Bone();
    bone.name = name;
    bone.position.fromArray(rig.offsets[i]);
    return bone;
  });
  bones.forEach((bone, i) => (rig.parents[i] < 0 ? group : bones[rig.parents[i]]).add(bone));
  group.updateMatrixWorld(true);
  const skeleton = new Skeleton(bones);
  return { group, skeleton, adapter: bindSkeleton(skeleton, rig) };
}

/** Apply local rotations to a compatible Three.js Skeleton. Retargeting is a separate operation. */
export function bindSkeleton(skeleton, rig, nameMap = {}) {
  validateRig(rig);
  if (!(skeleton instanceof Skeleton)) throw new Error('Expected a Three.js Skeleton.');
  const byName = new Map();
  for (const bone of skeleton.bones) {
    if (byName.has(bone.name)) throw new Error(`Duplicate target bone name: ${bone.name}`);
    byName.set(bone.name, bone);
  }
  const mapped = rig.names.map((name, i) => {
    const bone = byName.get(nameMap[name] ?? name);
    if (!bone && !rig.end_sites[i]) throw new Error(`Missing articulated joint: ${name}`);
    return bone;
  });
  if (new Set(mapped.filter(Boolean)).size !== mapped.filter(Boolean).length) {
    throw new Error('Multiple source joints map to the same target bone.');
  }
  mapped.forEach((bone, i) => {
    if (!bone) return;
    if (rig.parents[i] >= 0 && bone.parent !== mapped[rig.parents[i]]) {
      throw new Error(`Hierarchy mismatch: ${rig.names[i]}`);
    }
    // This initial adapter requires the canonical bind pose. It does not silently retarget proportions.
    if (bone.position.distanceTo(new Vector3(...rig.offsets[i])) > 1e-4) {
      throw new Error(`Rest offset mismatch: ${rig.names[i]}. Retarget before binding.`);
    }
  });
  const root = mapped[0];
  const work = new Quaternion();
  return {
    skeleton,
    rig,
    applyPose(pose) {
      if (pose.rootPosition?.length !== 3 || !pose.rootPosition.every(Number.isFinite) ||
          pose.localQuaternions?.length !== rig.names.length) throw new Error('Invalid pose dimensions.');
      // Validate the whole pose before mutating any bone.
      for (const q of pose.localQuaternions) {
        if (q?.length !== 4 || !q.every(Number.isFinite) || Math.hypot(...q) < 1e-8) {
          throw new Error('Nonfinite or degenerate pose rotation.');
        }
      }
      root.position.fromArray(pose.rootPosition);
      mapped.forEach((bone, i) => {
        if (bone) bone.quaternion.copy(work.fromArray(pose.localQuaternions[i]).normalize());
      });
      root.updateWorldMatrix(true, true);
    },
  };
}

/** Playback of saved source or prediction sequences. This does not run inference. */
export function samplePoseSequence(clip, seconds) {
  if (!['recorded_training_motion', 'saved_model_rollout'].includes(clip.kind) || clip.partition !== 'train' ||
      !Number.isFinite(seconds) || seconds < 0 || !(clip.dt > 0) || clip.root_positions.length < 1) {
    throw new Error('Expected a training pose sequence and a finite nonnegative time.');
  }
  const t = Math.min(seconds / clip.dt, clip.root_positions.length - 1);
  const a = Math.floor(t), b = Math.min(a + 1, clip.root_positions.length - 1), alpha = t - a;
  const rootPosition = new Vector3(...clip.root_positions[a]).lerp(new Vector3(...clip.root_positions[b]), alpha).toArray();
  const localQuaternions = clip.local_quaternions[a].map((q, i) => new Quaternion(...q)
    .slerp(new Quaternion(...clip.local_quaternions[b][i]), alpha).toArray());
  return { rootPosition, localQuaternions };
}

export function sampleRecordedMotion(clip, seconds) {
  if (clip.kind !== 'recorded_training_motion') throw new Error('Expected recorded source motion.');
  return samplePoseSequence(clip, seconds);
}

/** Rest-aligned world-rotation retargeting. This does not solve foot contacts. */
export function createRetargeter(skeleton, rig, {nameMap, directionMap, restCorrectionMap, scale=1, rotation=[0,0,0,1], offset=[0,0,0]}={}) {
  validateRig(rig);
  const scaleVector=typeof scale==='number'?[scale,scale,scale]:scale;
  if(!(skeleton instanceof Skeleton)||!nameMap||scaleVector?.length!==3||!scaleVector.every(v=>Number.isFinite(v)&&v>0)||
     rotation?.length!==4||!rotation.every(Number.isFinite)||Math.hypot(...rotation)<1e-8||
     offset?.length!==3||!offset.every(Number.isFinite))throw Error('Invalid retargeting configuration.');
  const byName=new Map();
  for(const bone of skeleton.bones){if(byName.has(bone.name))throw Error('Duplicate target bone name.');byName.set(bone.name,bone);}
  const alignment=new Quaternion(...rotation).normalize(),inverseAlignment=alignment.clone().invert();
  const entries=Object.entries(nameMap).map(([source,target])=>{
    const index=rig.names.indexOf(source),bone=byName.get(target);
    if(index<0||rig.end_sites[index]||!bone)throw Error(`Invalid retarget mapping: ${source} -> ${target}`);
    bone.updateWorldMatrix(true,false);
    return {index,bone,rest:bone.getWorldQuaternion(new Quaternion())};
  });
  if(!entries.some(e=>e.index===0)||new Set(entries.map(e=>e.bone)).size!==entries.length)throw Error('Map the source root and use unique target bones.');
  function depth(bone){let n=0;for(let p=bone.parent;p;p=p.parent)n++;return n;}
  entries.sort((a,b)=>depth(a.bone)-depth(b.bone));
  const root=entries.find(e=>e.index===0).bone;
  if(entries.some(e=>e.bone!==root&&!isBelow(e.bone,root)))throw Error('Mapped target bones must descend from the mapped root.');
  function isBelow(bone,parent){for(let p=bone.parent;p;p=p.parent)if(p===parent)return true;return false;}
  if(directionMap||restCorrectionMap){
    const sourceRest=[];
    rig.offsets.forEach((v,i)=>sourceRest.push(new Vector3(...v).add(rig.parents[i]<0?new Vector3():sourceRest[rig.parents[i]])));
    const corrections=new Map();
    for(const entry of [...entries].sort((a,b)=>a.index-b.index)){
      const spec=directionMap?.[rig.names[entry.index]],explicit=restCorrectionMap?.[rig.names[entry.index]];let correction;
      if(explicit){
        if(explicit.length!==4||!explicit.every(Number.isFinite)||Math.hypot(...explicit)<1e-8)throw Error('Invalid rest correction quaternion.');
        correction=new Quaternion(...explicit).normalize();
      }else if(spec){
        const child=rig.names.indexOf(spec.sourceChild),targetChild=byName.get(spec.targetChild);
        if(child<0||!targetChild||targetChild.parent!==entry.bone)throw Error('Direction calibration requires a source descendant and direct target child.');
        let ancestor=rig.parents[child];while(ancestor>=0&&ancestor!==entry.index)ancestor=rig.parents[ancestor];
        if(ancestor!==entry.index)throw Error('Source calibration child must descend from its mapped bone.');
        const from=targetChild.getWorldPosition(new Vector3()).sub(entry.bone.getWorldPosition(new Vector3()));
        const to=sourceRest[child].clone().sub(sourceRest[entry.index]).applyQuaternion(alignment);
        if(from.length()<1e-8||to.length()<1e-8)throw Error('Degenerate calibration direction.');
        correction=new Quaternion().setFromUnitVectors(from.normalize(),to.normalize());
      }else{
        let ancestor=rig.parents[entry.index];while(ancestor>=0&&!corrections.has(ancestor))ancestor=rig.parents[ancestor];
        correction=(corrections.get(ancestor)??new Quaternion()).clone();
      }
      corrections.set(entry.index,correction);entry.rest.premultiply(correction);
    }
  }
  return {skeleton,rig,
    applyPose(pose) {
      if(pose.rootPosition?.length!==3||!pose.rootPosition.every(Number.isFinite)||pose.localQuaternions?.length!==rig.names.length||
         pose.localQuaternions.some(q=>q?.length!==4||!q.every(Number.isFinite)||Math.hypot(...q)<1e-8))throw Error('Invalid source pose.');
      const sourceWorld=[];
      pose.localQuaternions.forEach((q,i)=>{
        const local=new Quaternion(...q).normalize();
        sourceWorld.push(rig.parents[i]<0?local:local.premultiply(sourceWorld[rig.parents[i]]));
      });
      const rootWorld=new Vector3(...pose.rootPosition).multiply(new Vector3(...scaleVector)).applyQuaternion(alignment).add(new Vector3(...offset));
      if(root.parent){root.parent.updateWorldMatrix(true,false);root.position.copy(root.parent.worldToLocal(rootWorld));}
      else root.position.copy(rootWorld);
      for(const {index,bone,rest} of entries){
        const world=alignment.clone().multiply(sourceWorld[index]).multiply(inverseAlignment).multiply(rest);
        if(bone.parent){bone.parent.updateWorldMatrix(true,false);world.premultiply(bone.parent.getWorldQuaternion(new Quaternion()).invert());}
        bone.quaternion.copy(world).normalize();bone.updateWorldMatrix(false,false);
      }
      root.updateWorldMatrix(true,true);
    },
  };
}