// Generic MJCF rig builder, adapted from the microduck sandbox's
// src/game/duck.js. Loads kinematics.json (produced by
// tools/mjcf_to_kinematics.py) and builds an Object3D tree with one Group per
// body; joints become per-body local rotations driven by name.
//
// Differences from the duck version: the jaw hinge and its SITTING_POSE are
// gone (G1 has neither), meshes are keyed by the original glTF node name
// because three strips '.' from node names, and a missing GLB is fatal rather
// than optional since every G1 mesh is packed into one.

import * as THREE from "three";
import { GLTFLoader } from "three/addons/loaders/GLTFLoader.js";
import { STLLoader } from "three/addons/loaders/STLLoader.js";
import { mergeVertices, toCreasedNormals } from "three/addons/utils/BufferGeometryUtils.js";

const CREASE = Math.PI / 5; // 36 deg

export async function loadKinematics(url) {
  const r = await fetch(url);
  if (!r.ok) throw new Error(`kinematics fetch ${r.status} for ${url}`);
  return r.json();
}

// three's GLTFLoader runs node names through PropertyBinding.sanitizeNodeName,
// which strips '.' and ':' - so a node exported as "pelvis.STL" arrives as
// "pelvisSTL". The untouched original survives on userData.name, so register
// every alias we might be asked for.
function registerAliases(map, mesh, entry) {
  for (const key of [mesh.userData?.name, mesh.name, mesh.geometry?.name]) {
    if (key && !map.has(key)) map.set(key, entry);
  }
}

function prepare(raw) {
  // STL repeats every vertex per facet and the packed GLB is already welded;
  // normals are rebuilt here either way so creases stay sharp.
  raw.deleteAttribute("normal");
  const welded = raw.index ? raw : mergeVertices(raw, 1e-4);
  // toCreasedNormals hashes on a 0.01-unit grid. Meshes are in metres, so
  // scale a clone to mm (10 um grid) and back.
  const scaled = welded.clone();
  scaled.scale(1000, 1000, 1000);
  const display = toCreasedNormals(scaled, CREASE);
  display.scale(1e-3, 1e-3, 1e-3);
  return { display, welded };
}

export async function loadGlbGeometries(url) {
  const gltf = await new GLTFLoader().loadAsync(url);
  const map = new Map();
  gltf.scene.traverse((o) => {
    if (!o.isMesh || !o.geometry) return;
    registerAliases(map, o, prepare(o.geometry));
  });
  return map;
}

export async function buildRig(k, opts = {}) {
  // placer carries the world pose in three.js space (no axis conversion)...
  const placer = new THREE.Group();
  placer.name = "robot_placer";
  // ...and root applies the MJCF +Z up -> three.js +Y up convention fix, so
  // everything below it stays in MJCF coordinates.
  const root = new THREE.Group();
  root.name = "robot_root";
  root.rotation.x = -Math.PI / 2;
  placer.add(root);

  const bodies = new Map();
  const joints = new Map();
  const geomByName = opts.glbUrl ? await loadGlbGeometries(opts.glbUrl) : new Map();

  const stlCache = new Map();
  const loadMesh = (name) => {
    const entry = geomByName.get(name);
    if (entry) return Promise.resolve(entry);
    if (!stlCache.has(name)) {
      stlCache.set(
        name,
        new STLLoader().loadAsync(`${k.mesh_dir}/${name}`).then(prepare),
      );
    }
    return stlCache.get(name);
  };

  for (const b of k.bodies) {
    const g = new THREE.Group();
    g.name = b.name;
    g.position.set(b.pos[0], b.pos[1], b.pos[2]);
    g.quaternion.set(b.quat[1], b.quat[2], b.quat[3], b.quat[0]);
    bodies.set(b.name, g);
  }
  for (const b of k.bodies) {
    const g = bodies.get(b.name);
    if (b.parent && bodies.has(b.parent)) bodies.get(b.parent).add(g);
    else root.add(g);
  }
  for (const b of k.bodies) {
    if (!b.joint || (b.joint.type && b.joint.type !== "hinge")) continue;
    const g = bodies.get(b.name);
    joints.set(b.joint.name, {
      body: g,
      axis: new THREE.Vector3(...b.joint.axis).normalize(),
      baseQuat: g.quaternion.clone(),
      range: b.joint.range ?? null,
    });
  }

  // Identical PBR props share one GPU material instance.
  const matCache = new Map();
  const matFor = (rgba) => {
    const key = rgba.join(",");
    const cached = matCache.get(key);
    if (cached) return cached;
    const m = new THREE.MeshStandardMaterial({
      color: new THREE.Color(rgba[0], rgba[1], rgba[2]),
      roughness: opts.roughness ?? 0.55,
      metalness: opts.metalness ?? 0.25,
      transparent: (rgba[3] ?? 1) < 1,
      opacity: rgba[3] ?? 1,
    });
    matCache.set(key, m);
    return m;
  };

  const pending = [];
  const missing = new Set();
  for (const b of k.bodies) {
    const g = bodies.get(b.name);
    if (!g) continue;
    for (const geom of b.geoms) {
      if (geom.type && geom.type !== "mesh") continue;
      if (!geom.mesh) continue;
      pending.push(
        loadMesh(geom.mesh)
          .then(({ display }) => {
            const m = new THREE.Mesh(display, matFor(geom.color ?? [0.8, 0.8, 0.8, 1]));
            m.userData.meshName = geom.mesh;
            m.castShadow = true;
            m.receiveShadow = true;
            if (geom.pos) m.position.set(...geom.pos);
            if (geom.quat) m.quaternion.set(geom.quat[1], geom.quat[2], geom.quat[3], geom.quat[0]);
            g.add(m);
          })
          .catch(() => missing.add(geom.mesh)),
      );
    }
  }
  await Promise.all(pending);
  if (missing.size) console.warn(`[rig] ${missing.size} mesh(es) failed:`, [...missing]);

  return { placer, root, bodies, joints, missing, pose: Object.fromEntries([...joints.keys()].map(name => [name, 0])) };
}

const _q = new THREE.Quaternion();

// Set one named joint, clamped to its MJCF range when known.
export function setJoint(rig, name, angle) {
  const j = rig.joints.get(name);
  if (!j || !Number.isFinite(angle)) return false;
  let a = angle;
  if (j.range) a = Math.min(j.range[1], Math.max(j.range[0], a));
  j.body.quaternion.copy(j.baseQuat).multiply(_q.setFromAxisAngle(j.axis, a));
  rig.pose[name] = a;
  return true;
}

export function applyPose(rig, pose) {
  for (const [name, ang] of Object.entries(pose)) setJoint(rig, name, ang);
}

// MJCF is Z-up, three.js is Y-up. The rig's root already applies Rx(-90) to
// everything below it, so a base pose given in MJCF world coords must be
// mapped into three.js space before it lands on the placer: positions become
// (x, z, -y) and the orientation is conjugated by the same rotation.
const AXIS_FIX = new THREE.Quaternion().setFromAxisAngle(new THREE.Vector3(1, 0, 0), -Math.PI / 2);
const AXIS_FIX_INV = AXIS_FIX.clone().invert();
const _qb = new THREE.Quaternion();

export function setBasePose(rig, pos, quat) {
  rig.placer.position.set(pos[0], pos[2], -pos[1]);
  if (!quat) return;
  _qb.set(quat[1], quat[2], quat[3], quat[0]);
  rig.placer.quaternion.copy(AXIS_FIX).multiply(_qb).multiply(AXIS_FIX_INV);
}

/** three.js world point -> MJCF world point, the inverse of the mapping above. */
export function toMujoco(v) {
  return [v.x, -v.z, v.y];
}

/** MJCF world point -> three.js world point. */
export function toThree(p, out) {
  return out.set(p[0], p[2], -p[1]);
}

/** Walk up from a hit mesh to the rig body that owns it. */
export function bodyOf(rig, object) {
  for (let o = object; o; o = o.parent) {
    if (rig.bodies.get(o.name) === o) return o.name;
  }
  return null;
}