mittens_engine/engine/ecs/system/

ik_system.rs

use crate::engine::ecs::component::ik_chain::TwoBoneIkDebugVisuals;
use crate::engine::ecs::component::{
    AvatarControlComponent, BoneRestPoseComponent, ColorComponent, EmissiveComponent,
    IKChainComponent, IKSolver, OverlayComponent, QueryRootMode, RenderableComponent,
    TransformComponent, resolve_component_ref,
};
use crate::engine::ecs::{ComponentId, IntentValue, SignalEmitter, World};
use crate::utils::math::{
    mat_to_quat, mat4_inverse, quat_conjugate, quat_from_axis_angle, quat_mul, quat_nlerp,
    quat_rotate_vec3, quat_rotation_y, quat_to_axis_angle, shortest_arc_quat, vec3_add, vec3_cross,
    vec3_dot, vec3_len, vec3_lerp, vec3_normalize, vec3_scale, vec3_sub,
};
use std::collections::HashSet;

#[derive(Debug, Default)]
pub struct IKSystem {
    chains: HashSet<ComponentId>,
}

impl IKSystem {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn tick(&mut self, world: &mut World, emit: &mut dyn SignalEmitter, _dt_sec: f32) {
        let ids: Vec<_> = self.chains.iter().copied().collect();
        for id in ids {
            // Lazy resolution of `target_source` / `end_effector_source` into
            // the actual ComponentId fields. Matches AnimationSystem's
            // behavior for ActionComponent and supports forward refs
            // (sources authored before the referent exists). No-op when the
            // ids are already filled (either by registry-time resolve at
            // call-construction, or by a previous tick).
            resolve_ik_chain_refs(world, id);
            resolve_avc_ancestor(world, id);
            tick_chain(id, world, emit);
        }
    }

    pub fn register(&mut self, component: ComponentId) {
        self.chains.insert(component);
    }

    pub fn remove(&mut self, component: ComponentId) {
        self.chains.remove(&component);
    }
}

/// Best-effort resolve of IKChain's `target_source` / `end_effector_source`
/// into the matching `target_id` / `end_effector_id` slots when the latter
/// are null. Silent no-op on miss — the IK solver short-circuits on a null
/// target anyway, so a still-unresolved chain just skips that tick. A future
/// tick may succeed once the referent spawns.
fn resolve_ik_chain_refs(world: &mut World, id: ComponentId) {
    use crate::engine::ecs::component::ComponentRef;
    use slotmap::Key;
    let (target_src, end_src, target_id, end_id) = {
        let Some(ik) = world.get_component_by_id_as::<IKChainComponent>(id) else {
            return;
        };
        (
            ik.target_source.clone(),
            ik.end_effector_source.clone(),
            ik.target_id,
            ik.end_effector_id,
        )
    };

    let resolve = |src: &ComponentRef| -> Option<ComponentId> {
        resolve_component_ref(world, src, Some(id), QueryRootMode::SelfSubtree)
    };

    let new_target = if target_id.is_null() {
        target_src.as_ref().and_then(resolve)
    } else {
        None
    };
    let new_end = if end_id.is_null() {
        end_src.as_ref().and_then(resolve)
    } else {
        None
    };

    if new_target.is_none() && new_end.is_none() {
        return;
    }
    if let Some(ik) = world.get_component_by_id_as_mut::<IKChainComponent>(id) {
        if let Some(t) = new_target {
            ik.target_id = t;
        }
        if let Some(e) = new_end {
            ik.end_effector_id = e;
        }
    }
}

/// One-time ancestor walk: find the nearest `AvatarControlComponent` above
/// this IKChainComponent and cache it.  No-op once cached.
fn resolve_avc_ancestor(world: &mut World, id: ComponentId) {
    let needs_walk = world
        .get_component_by_id_as::<IKChainComponent>(id)
        .map(|c| c.avc_id.is_none())
        .unwrap_or(false);
    if !needs_walk {
        return;
    }
    let found = find_avc_ancestor(world, id);
    if let Some(ik) = world.get_component_by_id_as_mut::<IKChainComponent>(id) {
        ik.avc_id = found;
    }
}

/// Walk up the parent chain from `id` to find the nearest ancestor
/// `AvatarControlComponent`.  Returns `None` if no AVC is found within
/// 32 hops.
fn find_avc_ancestor(world: &World, id: ComponentId) -> Option<ComponentId> {
    let mut cur = id;
    for _ in 0..32 {
        let parent = world.parent_of(cur)?;
        if world
            .get_component_by_id_as::<AvatarControlComponent>(parent)
            .is_some()
        {
            return Some(parent);
        }
        cur = parent;
    }
    None
}

fn tick_chain(id: ComponentId, world: &mut World, emit: &mut dyn SignalEmitter) {
    let (solver, target_id, end_effector_id, weight, avc_id, xr_pose_driver) = {
        let Some(c) = world.get_component_by_id_as::<IKChainComponent>(id) else {
            return;
        };
        (
            c.solver,
            c.target_id,
            c.end_effector_id,
            c.weight,
            c.avc_id,
            c.xr_pose_driver,
        )
    };
    if weight <= 0.0 {
        return;
    }
    if let Some(driver) = xr_pose_driver {
        let valid = world
            .get_component_by_id_as::<crate::engine::ecs::component::InputXRComponent>(driver)
            .map(|component| component.pose_valid)
            .or_else(|| {
                world
                    .get_component_by_id_as::<crate::engine::ecs::component::ControllerXRComponent>(
                        driver,
                    )
                    .map(|component| component.pose_valid)
            })
            .unwrap_or(false);
        if !valid {
            return;
        }
    }

    // For AimConstraint / Fabrik, root joint TC = parent of IKChainComponent.
    // TwoBoneIK ignores this and uses the explicit joint IDs on the variant.
    let root_tc_opt = world.parent_of(id).filter(|&p| {
        world
            .get_component_by_id_as::<TransformComponent>(p)
            .is_some()
    });

    match solver {
        IKSolver::AimConstraint {
            offset_yaw,
            copy_position,
            target_position_offset,
        } => {
            let Some(root_tc) = root_tc_opt else { return };
            solve_aim(
                world,
                emit,
                root_tc,
                target_id,
                offset_yaw,
                copy_position,
                target_position_offset,
                weight,
            );
        }
        IKSolver::TwoBoneIK {
            root_joint_id,
            mid_joint_id,
            pole_direction,
            copy_end_rotation,
        } => {
            // Explicit 3-node chain — no topology discovery. Sibling helper /
            // collider / cloth nodes under the arm bones are ignored.
            use slotmap::Key;
            if root_joint_id.is_null() || mid_joint_id.is_null() || end_effector_id.is_null() {
                return;
            }
            if world
                .get_component_by_id_as::<TransformComponent>(root_joint_id)
                .is_none()
                || world
                    .get_component_by_id_as::<TransformComponent>(mid_joint_id)
                    .is_none()
                || world
                    .get_component_by_id_as::<TransformComponent>(end_effector_id)
                    .is_none()
            {
                return;
            }
            let chain = [root_joint_id, mid_joint_id, end_effector_id];
            solve_two_bone(
                world,
                emit,
                id,
                &chain,
                target_id,
                pole_direction,
                copy_end_rotation,
                weight,
                avc_id,
            );
        }
        IKSolver::Fabrik {
            max_iterations,
            tolerance,
            target_position_offset,
        } => {
            let Some(root_tc) = root_tc_opt else { return };
            let chain = collect_tc_chain(world, root_tc, end_effector_id);
            if chain.len() < 2 {
                return;
            }
            solve_fabrik(
                world,
                emit,
                &chain,
                target_id,
                target_position_offset,
                max_iterations,
                tolerance,
                weight,
            );
        }
    }
}

// ---------------------------------------------------------------------------
// Chain collection
// ---------------------------------------------------------------------------

/// Walk UP from `end_id` via TC parents until reaching `root`.  Returns the
/// collected IDs in root-to-end order.
///
/// Walking up is unique (each TC has one parent) so this picks out a single
/// topological path even when intermediate joints have multiple TC children
/// (e.g. the spine fork into clavicles).  Returns empty if `end_id` is not
/// a TC descendant of `root` within 32 hops.
fn collect_tc_chain(world: &World, root: ComponentId, end_id: ComponentId) -> Vec<ComponentId> {
    if root == end_id {
        return vec![root];
    }
    let mut up: Vec<ComponentId> = vec![end_id];
    let mut cur = end_id;
    for _ in 0..32 {
        let Some(parent) = world.parent_of(cur) else {
            return Vec::new();
        };
        if world
            .get_component_by_id_as::<TransformComponent>(parent)
            .is_none()
        {
            return Vec::new();
        }
        up.push(parent);
        if parent == root {
            up.reverse();
            return up;
        }
        cur = parent;
    }
    Vec::new()
}

// ---------------------------------------------------------------------------
// AimConstraint solver
// ---------------------------------------------------------------------------

/// Orient `root_tc` so its world rotation matches `target_id`'s world rotation
/// post-multiplied by `rot_y(offset_yaw)`.  Only modifies rotation; preserves
/// existing local translation and scale.
fn solve_aim(
    world: &World,
    emit: &mut dyn SignalEmitter,
    root_tc: ComponentId,
    target_id: ComponentId,
    offset_yaw: f32,
    copy_position: bool,
    target_position_offset: [f32; 3],
    weight: f32,
) {
    let Some(target_tc) = world.get_component_by_id_as::<TransformComponent>(target_id) else {
        return;
    };
    let target_world_rot = mat_to_quat(target_tc.transform.matrix_world);
    let desired_world_rot = quat_mul(target_world_rot, quat_rotation_y(offset_yaw));
    // Apply the offset in TARGET local frame, then add to target world position.
    // For an HMD target with offset = (0, -eye_height, 0), this drops the bone target
    // down along the HMD's local Y so the eye mesh (above the bone pivot) lines up
    // with the HMD position.
    let target_local_offset_world = quat_rotate_vec3(target_world_rot, target_position_offset);
    let target_world_pos = [
        target_tc.transform.matrix_world[3][0] + target_local_offset_world[0],
        target_tc.transform.matrix_world[3][1] + target_local_offset_world[1],
        target_tc.transform.matrix_world[3][2] + target_local_offset_world[2],
    ];

    let parent_world_mat = world
        .parent_of(root_tc)
        .and_then(|p| world.get_component_by_id_as::<TransformComponent>(p))
        .map(|t| t.transform.matrix_world)
        .unwrap_or([
            [1.0, 0.0, 0.0, 0.0],
            [0.0, 1.0, 0.0, 0.0],
            [0.0, 0.0, 1.0, 0.0],
            [0.0, 0.0, 0.0, 1.0],
        ]);
    let parent_world_rot = mat_to_quat(parent_world_mat);

    let full_local_rot = quat_mul(quat_conjugate(parent_world_rot), desired_world_rot);

    let local_rot = if weight < 1.0 {
        let cur = world
            .get_component_by_id_as::<TransformComponent>(root_tc)
            .map(|t| t.transform.rotation)
            .unwrap_or([0.0, 0.0, 0.0, 1.0]);
        quat_nlerp(cur, full_local_rot, weight)
    } else {
        full_local_rot
    };

    let (cur_t, s) = world
        .get_component_by_id_as::<TransformComponent>(root_tc)
        .map(|tc| (tc.transform.translation, tc.transform.scale))
        .unwrap_or(([0.0; 3], [1.0, 1.0, 1.0]));

    let local_t = if copy_position {
        // Invert parent world matrix to get local position from target world position.
        // Closed-form inverse of an affine TRS matrix: local_pos = R^T * (target_pos - parent_pos) / parent_scale.
        // Easier: use the inverse-transpose of the 3x3 rotation+scale block, then translate.
        let parent_pos = [
            parent_world_mat[3][0],
            parent_world_mat[3][1],
            parent_world_mat[3][2],
        ];
        let delta = [
            target_world_pos[0] - parent_pos[0],
            target_world_pos[1] - parent_pos[1],
            target_world_pos[2] - parent_pos[2],
        ];
        // Apply inverse parent rotation (assuming uniform scale; for non-uniform scale this would be approximate).
        let inv_parent_rot = quat_conjugate(parent_world_rot);
        let local_pos = quat_rotate_vec3(inv_parent_rot, delta);
        if weight < 1.0 {
            [
                cur_t[0] + (local_pos[0] - cur_t[0]) * weight,
                cur_t[1] + (local_pos[1] - cur_t[1]) * weight,
                cur_t[2] + (local_pos[2] - cur_t[2]) * weight,
            ]
        } else {
            local_pos
        }
    } else {
        cur_t
    };

    emit.push_intent_now(
        root_tc,
        IntentValue::UpdateTransform {
            component_ids: vec![root_tc],
            translation: local_t,
            rotation_quat_xyzw: local_rot,
            scale: s,
        },
    );
}

// ---------------------------------------------------------------------------
// TwoBoneIK solver
// ---------------------------------------------------------------------------

/// Closed-form 2-bone IK.
///
/// `chain` must have length ≥ 3: [root (upper arm), mid (lower arm), end (hand)].
/// `target_id`: TC providing the desired hand world position.
/// `pole_direction`: direction hint for the elbow.  Interpreted in body-local
///   space when `avc_id` is `Some` (transformed to world via model root rotation);
///   world-space when `avc_id` is `None`.
/// `copy_end_rotation`: if true, also aligns the end bone to the target's world rotation.
/// `avc_id`: cached ancestor `AvatarControlComponent` id, or `None` for world-space pole.
fn solve_two_bone(
    world: &mut World,
    emit: &mut dyn SignalEmitter,
    ik_chain_id: ComponentId,
    chain: &[ComponentId],
    target_id: ComponentId,
    pole_direction: [f32; 3],
    copy_end_rotation: bool,
    weight: f32,
    avc_id: Option<ComponentId>,
) {
    let (root_tc, mid_tc, end_tc) = (chain[0], chain[1], chain[2]);

    // FK world positions — bone lengths are measured here each tick.
    let root_pos = tc_world_pos(world, root_tc);
    let mid_pos = tc_world_pos(world, mid_tc);
    let end_pos = tc_world_pos(world, end_tc);
    let target_pos = tc_world_pos(world, target_id);

    let root_world_rot = tc_world_rot(world, root_tc);
    let mid_world_rot = tc_world_rot(world, mid_tc);

    let upper_len = vec3_len(vec3_sub(mid_pos, root_pos)).max(1e-6);
    let lower_len = vec3_len(vec3_sub(end_pos, mid_pos)).max(1e-6);

    let root_parent_world_rot = world
        .parent_of(root_tc)
        .and_then(|p| world.get_component_by_id_as::<TransformComponent>(p))
        .map(|t| mat_to_quat(t.transform.matrix_world))
        .unwrap_or([0.0, 0.0, 0.0, 1.0]);

    // Triangle solve — clamp reach to avoid degenerate case beyond full extension.
    let to_target = vec3_sub(target_pos, root_pos);
    let raw_reach = vec3_len(to_target);
    let reach = raw_reach.min(upper_len + lower_len - 1e-4).max(1e-6);
    let reach_dir = if raw_reach > 1e-6 {
        vec3_scale(to_target, 1.0 / raw_reach)
    } else {
        [0.0, 1.0, 0.0]
    };

    let cos_upper = ((upper_len * upper_len + reach * reach - lower_len * lower_len)
        / (2.0 * upper_len * reach))
        .clamp(-1.0, 1.0);
    let upper_angle = cos_upper.acos();

    // Build elbow plane from pole hint.
    // Transform body-local pole to world space when under an AVC.
    let pole = match avc_id {
        Some(avc) => world
            .get_component_by_id_as::<AvatarControlComponent>(avc)
            .and_then(|c| c.model_root_id)
            .map(|root_id| {
                let root_rot = tc_world_rot(world, root_id);
                quat_rotate_vec3(root_rot, pole_direction)
            })
            .unwrap_or(pole_direction),
        None => pole_direction,
    };
    let cross_tp = vec3_cross(to_target, pole);
    let plane_normal = if vec3_len(cross_tp) > 1e-6 {
        vec3_normalize(cross_tp)
    } else {
        let fallback = if reach_dir[0].abs() < 0.9 {
            [1.0, 0.0, 0.0]
        } else {
            [0.0, 1.0, 0.0]
        };
        vec3_normalize(vec3_cross(to_target, fallback))
    };
    let perp = vec3_normalize(vec3_cross(plane_normal, reach_dir));

    let elbow_dir = vec3_normalize(vec3_add(
        vec3_scale(reach_dir, upper_angle.cos()),
        vec3_scale(perp, upper_angle.sin()),
    ));
    let elbow_pos = vec3_add(root_pos, vec3_scale(elbow_dir, upper_len));

    maybe_update_two_bone_debug(
        world,
        emit,
        ik_chain_id,
        avc_id,
        root_pos,
        target_pos,
        pole,
        plane_normal,
        elbow_pos,
    );

    // --- Upper arm ---
    let old_upper_fwd = if vec3_len(vec3_sub(mid_pos, root_pos)) > 1e-6 {
        vec3_normalize(vec3_sub(mid_pos, root_pos))
    } else {
        [0.0, 0.0, 1.0]
    };
    let delta_upper = shortest_arc_quat(old_upper_fwd, elbow_dir);
    let new_upper_world_rot = quat_mul(delta_upper, root_world_rot);

    let full_upper_local = quat_mul(quat_conjugate(root_parent_world_rot), new_upper_world_rot);
    let upper_local = if weight < 1.0 {
        let cur = world
            .get_component_by_id_as::<TransformComponent>(root_tc)
            .map(|t| t.transform.rotation)
            .unwrap_or([0.0, 0.0, 0.0, 1.0]);
        quat_nlerp(cur, full_upper_local, weight)
    } else {
        full_upper_local
    };

    // --- Lower arm ---
    let old_lower_fwd = if vec3_len(vec3_sub(end_pos, mid_pos)) > 1e-6 {
        vec3_normalize(vec3_sub(end_pos, mid_pos))
    } else {
        [0.0, 0.0, 1.0]
    };
    // After the upper arm delta, the lower arm's FK forward rotates by delta_upper too.
    let lower_fwd_after_upper = quat_rotate_vec3(delta_upper, old_lower_fwd);
    let new_lower_fwd = if vec3_len(vec3_sub(target_pos, elbow_pos)) > 1e-6 {
        vec3_normalize(vec3_sub(target_pos, elbow_pos))
    } else {
        lower_fwd_after_upper
    };
    let delta_lower = shortest_arc_quat(lower_fwd_after_upper, new_lower_fwd);
    let mut new_lower_world_rot = quat_mul(delta_lower, quat_mul(delta_upper, mid_world_rot));
    let end_local_rest = read_bone_rest_rot(world, end_tc);
    let solved_forearm_axis = if vec3_len(vec3_sub(target_pos, elbow_pos)) > 1e-6 {
        vec3_normalize(vec3_sub(target_pos, elbow_pos))
    } else {
        new_lower_fwd
    };
    let solved_hand_world_rot = quat_mul(new_lower_world_rot, end_local_rest);
    let target_world_rot = tc_world_rot(world, target_id);
    let target_from_solved_hand = quat_mul(target_world_rot, quat_conjugate(solved_hand_world_rot));
    let forearm_twist_world =
        extract_twist_about_axis(target_from_solved_hand, solved_forearm_axis);
    let (_, forearm_twist_angle) = quat_to_axis_angle(forearm_twist_world);
    if forearm_twist_angle.abs() > 1e-4 {
        new_lower_world_rot = quat_mul(forearm_twist_world, new_lower_world_rot);
    }

    // Parent of lower arm is now upper arm with new_upper_world_rot.
    let full_lower_local = quat_mul(quat_conjugate(new_upper_world_rot), new_lower_world_rot);
    let lower_local = if weight < 1.0 {
        let cur = world
            .get_component_by_id_as::<TransformComponent>(mid_tc)
            .map(|t| t.transform.rotation)
            .unwrap_or([0.0, 0.0, 0.0, 1.0]);
        quat_nlerp(cur, full_lower_local, weight)
    } else {
        full_lower_local
    };

    // Emit UpdateTransform for upper arm (preserve existing translation/scale).
    let (rt, rs) = world
        .get_component_by_id_as::<TransformComponent>(root_tc)
        .map(|t| (t.transform.translation, t.transform.scale))
        .unwrap_or(([0.0; 3], [1.0, 1.0, 1.0]));
    emit.push_intent_now(
        root_tc,
        IntentValue::UpdateTransform {
            component_ids: vec![root_tc],
            translation: rt,
            rotation_quat_xyzw: upper_local,
            scale: rs,
        },
    );

    // Emit UpdateTransform for lower arm.
    let (mt, ms) = world
        .get_component_by_id_as::<TransformComponent>(mid_tc)
        .map(|t| (t.transform.translation, t.transform.scale))
        .unwrap_or(([0.0; 3], [1.0, 1.0, 1.0]));
    emit.push_intent_now(
        mid_tc,
        IntentValue::UpdateTransform {
            component_ids: vec![mid_tc],
            translation: mt,
            rotation_quat_xyzw: lower_local,
            scale: ms,
        },
    );

    // Optionally copy target rotation to end-effector bone.
    if copy_end_rotation {
        let full_end_local = quat_mul(quat_conjugate(new_lower_world_rot), target_world_rot);
        let end_local = if weight < 1.0 {
            let cur = world
                .get_component_by_id_as::<TransformComponent>(end_tc)
                .map(|t| t.transform.rotation)
                .unwrap_or([0.0, 0.0, 0.0, 1.0]);
            quat_nlerp(cur, full_end_local, weight)
        } else {
            full_end_local
        };
        let (et, es) = world
            .get_component_by_id_as::<TransformComponent>(end_tc)
            .map(|t| (t.transform.translation, t.transform.scale))
            .unwrap_or(([0.0; 3], [1.0, 1.0, 1.0]));
        emit.push_intent_now(
            end_tc,
            IntentValue::UpdateTransform {
                component_ids: vec![end_tc],
                translation: et,
                rotation_quat_xyzw: end_local,
                scale: es,
            },
        );
    }
}

fn extract_twist_about_axis(delta_rot: [f32; 4], axis_world: [f32; 3]) -> [f32; 4] {
    let axis = vec3_normalize(axis_world);
    if vec3_len(axis) <= 1e-6 {
        return [0.0, 0.0, 0.0, 1.0];
    }
    let projected = vec3_scale(
        axis,
        vec3_dot([delta_rot[0], delta_rot[1], delta_rot[2]], axis),
    );
    let twist = [projected[0], projected[1], projected[2], delta_rot[3]];
    let len_sq =
        twist[0] * twist[0] + twist[1] * twist[1] + twist[2] * twist[2] + twist[3] * twist[3];
    if len_sq <= 1e-12 {
        [0.0, 0.0, 0.0, 1.0]
    } else {
        let inv_len = len_sq.sqrt().recip();
        [
            twist[0] * inv_len,
            twist[1] * inv_len,
            twist[2] * inv_len,
            twist[3] * inv_len,
        ]
    }
}

fn maybe_update_two_bone_debug(
    world: &mut World,
    emit: &mut dyn SignalEmitter,
    ik_chain_id: ComponentId,
    avc_id: Option<ComponentId>,
    root_pos: [f32; 3],
    target_pos: [f32; 3],
    pole_world: [f32; 3],
    plane_normal: [f32; 3],
    elbow_pos: [f32; 3],
) {
    let enabled = avc_id
        .and_then(|id| world.get_component_by_id_as::<AvatarControlComponent>(id))
        .map(|avc| avc.ik_debug)
        .unwrap_or(false);
    if !enabled {
        return;
    }

    let visuals = ensure_two_bone_debug_visuals(world, ik_chain_id);
    let pole_tip = vec3_add(root_pos, vec3_scale(vec3_normalize_or_z(pole_world), 0.30));
    let normal_tip = vec3_add(
        root_pos,
        vec3_scale(vec3_normalize_or_z(plane_normal), 0.22),
    );

    update_debug_segment(
        world,
        emit,
        visuals.target_line,
        root_pos,
        target_pos,
        0.010,
    );
    update_debug_segment(world, emit, visuals.pole_line, root_pos, pole_tip, 0.008);
    update_debug_segment(
        world,
        emit,
        visuals.plane_normal_line,
        root_pos,
        normal_tip,
        0.008,
    );
    update_debug_segment(world, emit, visuals.elbow_line, root_pos, elbow_pos, 0.010);
    update_debug_point(
        world,
        emit,
        visuals.elbow_point,
        elbow_pos,
        [0.035, 0.035, 0.035],
    );
}

fn ensure_two_bone_debug_visuals(
    world: &mut World,
    ik_chain_id: ComponentId,
) -> TwoBoneIkDebugVisuals {
    if let Some(existing) = world
        .get_component_by_id_as::<IKChainComponent>(ik_chain_id)
        .and_then(|ik| ik.two_bone_debug_visuals)
    {
        return existing;
    }

    let parent = world
        .get_component_by_id_as::<IKChainComponent>(ik_chain_id)
        .and_then(|ik| ik.avc_id)
        .unwrap_or(ik_chain_id);

    let visuals = TwoBoneIkDebugVisuals {
        target_line: spawn_debug_cube(
            world,
            parent,
            "ik_debug_target_line",
            (1.0, 0.85, 0.15, 0.95),
        ),
        pole_line: spawn_debug_cube(world, parent, "ik_debug_pole_line", (0.10, 0.95, 1.0, 0.95)),
        plane_normal_line: spawn_debug_cube(
            world,
            parent,
            "ik_debug_plane_normal_line",
            (1.0, 0.35, 0.80, 0.95),
        ),
        elbow_line: spawn_debug_cube(
            world,
            parent,
            "ik_debug_elbow_line",
            (0.20, 1.0, 0.35, 0.95),
        ),
        elbow_point: spawn_debug_cube(world, parent, "ik_debug_elbow_point", (1.0, 1.0, 1.0, 0.95)),
    };

    if let Some(ik) = world.get_component_by_id_as_mut::<IKChainComponent>(ik_chain_id) {
        ik.two_bone_debug_visuals = Some(visuals);
    }
    visuals
}

fn spawn_debug_cube(
    world: &mut World,
    parent: ComponentId,
    name: &str,
    color: (f32, f32, f32, f32),
) -> ComponentId {
    let overlay = world.add_component(OverlayComponent::new());
    let transform =
        world.add_component_boxed_named(name.to_string(), Box::new(TransformComponent::new()));
    let renderable = world.add_component(RenderableComponent::cube());
    let color = world.add_component(ColorComponent::rgba(color.0, color.1, color.2, color.3));
    let emissive = world.add_component(EmissiveComponent::on());
    let _ = world.add_child(overlay, transform);
    let _ = world.add_child(transform, renderable);
    let _ = world.add_child(renderable, color);
    let _ = world.add_child(renderable, emissive);
    let _ = world.add_child(parent, overlay);
    transform
}

fn update_debug_segment(
    world: &World,
    emit: &mut dyn SignalEmitter,
    component_id: ComponentId,
    start: [f32; 3],
    end: [f32; 3],
    thickness: f32,
) {
    let delta = vec3_sub(end, start);
    let len = vec3_len(delta);
    let dir = if len > 1e-6 {
        vec3_scale(delta, 1.0 / len)
    } else {
        [0.0, 0.0, 1.0]
    };
    let mid = vec3_scale(vec3_add(start, end), 0.5);
    let rot = shortest_arc_quat([0.0, 0.0, 1.0], dir);
    let local_mid = world_point_to_local(world, component_id, mid);
    let local_rot = world_quat_to_local(world, component_id, rot);
    emit.push_intent_now(
        component_id,
        IntentValue::UpdateTransform {
            component_ids: vec![component_id],
            translation: local_mid,
            rotation_quat_xyzw: local_rot,
            scale: [thickness, thickness, len.max(thickness)],
        },
    );
}

fn update_debug_point(
    world: &World,
    emit: &mut dyn SignalEmitter,
    component_id: ComponentId,
    position: [f32; 3],
    scale: [f32; 3],
) {
    let local_pos = world_point_to_local(world, component_id, position);
    emit.push_intent_now(
        component_id,
        IntentValue::UpdateTransform {
            component_ids: vec![component_id],
            translation: local_pos,
            rotation_quat_xyzw: [0.0, 0.0, 0.0, 1.0],
            scale,
        },
    );
}

fn world_point_to_local(
    world: &World,
    component_id: ComponentId,
    world_point: [f32; 3],
) -> [f32; 3] {
    let parent_world = nearest_ancestor_transform_world(world, component_id);
    let Some(inv_parent) = parent_world.and_then(mat4_inverse) else {
        return world_point;
    };
    let local = mat4_mul_vec4(
        inv_parent,
        [world_point[0], world_point[1], world_point[2], 1.0],
    );
    if local[3].abs() > 1e-6 {
        [
            local[0] / local[3],
            local[1] / local[3],
            local[2] / local[3],
        ]
    } else {
        [local[0], local[1], local[2]]
    }
}

fn world_quat_to_local(world: &World, component_id: ComponentId, world_rot: [f32; 4]) -> [f32; 4] {
    let parent_world_rot = nearest_ancestor_transform_world(world, component_id)
        .map(mat_to_quat)
        .unwrap_or([0.0, 0.0, 0.0, 1.0]);
    quat_mul(quat_conjugate(parent_world_rot), world_rot)
}

fn nearest_ancestor_transform_world(
    world: &World,
    component_id: ComponentId,
) -> Option<[[f32; 4]; 4]> {
    let mut cur = component_id;
    while let Some(parent) = world.parent_of(cur) {
        if let Some(t) = world.get_component_by_id_as::<TransformComponent>(parent) {
            return Some(t.transform.matrix_world);
        }
        cur = parent;
    }
    None
}

fn mat4_mul_vec4(m: [[f32; 4]; 4], v: [f32; 4]) -> [f32; 4] {
    [
        m[0][0] * v[0] + m[1][0] * v[1] + m[2][0] * v[2] + m[3][0] * v[3],
        m[0][1] * v[0] + m[1][1] * v[1] + m[2][1] * v[2] + m[3][1] * v[3],
        m[0][2] * v[0] + m[1][2] * v[1] + m[2][2] * v[2] + m[3][2] * v[3],
        m[0][3] * v[0] + m[1][3] * v[1] + m[2][3] * v[2] + m[3][3] * v[3],
    ]
}

fn vec3_normalize_or_z(v: [f32; 3]) -> [f32; 3] {
    let len = vec3_len(v);
    if len > 1e-6 {
        vec3_scale(v, 1.0 / len)
    } else {
        [0.0, 0.0, 1.0]
    }
}

// ---------------------------------------------------------------------------
// FABRIK solver
// ---------------------------------------------------------------------------

fn solve_fabrik(
    world: &World,
    emit: &mut dyn SignalEmitter,
    chain: &[ComponentId],
    target_id: ComponentId,
    target_position_offset: [f32; 3],
    max_iterations: u32,
    tolerance: f32,
    weight: f32,
) {
    let n = chain.len();

    let mut positions: Vec<[f32; 3]> = chain.iter().map(|&tc| tc_world_pos(world, tc)).collect();
    let bone_lengths: Vec<f32> = (0..n - 1)
        .map(|i| vec3_len(vec3_sub(positions[i + 1], positions[i])).max(1e-6))
        .collect();
    let root_pos = positions[0];
    // Target = target world pos + R(target_world_rot) * target_position_offset.
    // Symmetric with solve_aim: lets a Y-only offset like (0, -eye_h, 0) drop the
    // chase point down along the target's local Y, regardless of target yaw.
    let target_pos = {
        let base = tc_world_pos(world, target_id);
        let target_rot = tc_world_rot(world, target_id);
        let off = quat_rotate_vec3(target_rot, target_position_offset);
        [base[0] + off[0], base[1] + off[1], base[2] + off[2]]
    };

    // FABRIK iterations.
    for _ in 0..max_iterations {
        if vec3_len(vec3_sub(*positions.last().unwrap(), target_pos)) < tolerance {
            break;
        }
        // Forward pass — pull end to target.
        *positions.last_mut().unwrap() = target_pos;
        for i in (0..n - 1).rev() {
            let d = vec3_len(vec3_sub(positions[i], positions[i + 1]));
            let t = if d > 1e-9 { bone_lengths[i] / d } else { 0.0 };
            positions[i] = vec3_lerp(positions[i + 1], positions[i], t);
        }
        // Backward pass — pin root.
        positions[0] = root_pos;
        for i in 0..n - 1 {
            let d = vec3_len(vec3_sub(positions[i + 1], positions[i]));
            let t = if d > 1e-9 { bone_lengths[i] / d } else { 0.0 };
            positions[i + 1] = vec3_lerp(positions[i], positions[i + 1], t);
        }
    }

    // Convert solved bone directions to local rotations and emit.
    let mut parent_world_rot = world
        .parent_of(chain[0])
        .and_then(|p| world.get_component_by_id_as::<TransformComponent>(p))
        .map(|t| mat_to_quat(t.transform.matrix_world))
        .unwrap_or([0.0, 0.0, 0.0, 1.0f32]);

    // Detect neck bone by name — it should be one bone before the head in the chain.
    // Typical name patterns: "J_Bip_C_Neck", "Armature.Neck", "Neck", etc.
    let neck_index = chain.iter().position(|&id| {
        world
            .get_component_by_id_as::<TransformComponent>(id)
            .and_then(|_| Some(()))
            .is_some()
            && world.children_of(id).iter().any(|&child| {
                // Head is typically a direct child of neck (or wrapped in splice).
                let is_head = world
                    .get_component_by_id_as::<TransformComponent>(child)
                    .and_then(|_| {
                        // Just check if it's the next bone in the chain
                        chain.iter().find(|&&c| c == child).map(|_| ())
                    })
                    .is_some();
                is_head
            })
    });

    for i in 0..n - 1 {
        let tc = chain[i];
        let cur_world_rot = tc_world_rot(world, tc);

        let cur_fwd = {
            let from = tc_world_pos(world, tc);
            let to = tc_world_pos(world, chain[i + 1]);
            let d = vec3_sub(to, from);
            if vec3_len(d) > 1e-6 {
                vec3_normalize(d)
            } else {
                [0.0, 0.0, 1.0]
            }
        };
        let desired_fwd = {
            let d = vec3_sub(positions[i + 1], positions[i]);
            if vec3_len(d) > 1e-6 {
                vec3_normalize(d)
            } else {
                cur_fwd
            }
        };

        // If this is the neck bone, constrain its rotation to be minimal (keep it
        // rigid relative to the upper torso). The neck should inherit the upper body's
        // yaw/rotation without bending, so we clamp the rotation angle.
        let delta = if neck_index == Some(i) {
            // Neck constraint: allow only small deviations from the current forward direction.
            // Use a very tight cone (max ~15 degrees) to keep the neck stiff.
            let unconstrained_delta = shortest_arc_quat(cur_fwd, desired_fwd);
            let (axis, angle) = quat_to_axis_angle(unconstrained_delta);
            let max_neck_angle = 0.26f32; // ~15 degrees
            if angle.abs() > max_neck_angle {
                // Clamp the rotation angle
                let clamped_angle = angle.max(-max_neck_angle).min(max_neck_angle);
                quat_from_axis_angle(axis, clamped_angle)
            } else {
                unconstrained_delta
            }
        } else {
            shortest_arc_quat(cur_fwd, desired_fwd)
        };

        let new_world_rot = quat_mul(delta, cur_world_rot);
        let full_local = quat_mul(quat_conjugate(parent_world_rot), new_world_rot);

        let local_rot = if weight < 1.0 {
            let cur = world
                .get_component_by_id_as::<TransformComponent>(tc)
                .map(|t| t.transform.rotation)
                .unwrap_or([0.0, 0.0, 0.0, 1.0]);
            quat_nlerp(cur, full_local, weight)
        } else {
            full_local
        };

        let (t, s) = world
            .get_component_by_id_as::<TransformComponent>(tc)
            .map(|tc| (tc.transform.translation, tc.transform.scale))
            .unwrap_or(([0.0; 3], [1.0, 1.0, 1.0]));
        emit.push_intent_now(
            tc,
            IntentValue::UpdateTransform {
                component_ids: vec![tc],
                translation: t,
                rotation_quat_xyzw: local_rot,
                scale: s,
            },
        );

        parent_world_rot = new_world_rot;
    }
}

// ---------------------------------------------------------------------------
// World-matrix helpers
// ---------------------------------------------------------------------------

fn tc_world_pos(world: &World, id: ComponentId) -> [f32; 3] {
    world
        .get_component_by_id_as::<TransformComponent>(id)
        .map(|t| {
            let m = t.transform.matrix_world;
            [m[3][0], m[3][1], m[3][2]]
        })
        .unwrap_or([0.0; 3])
}

fn tc_world_rot(world: &World, id: ComponentId) -> [f32; 4] {
    world
        .get_component_by_id_as::<TransformComponent>(id)
        .map(|t| mat_to_quat(t.transform.matrix_world))
        .unwrap_or([0.0, 0.0, 0.0, 1.0])
}

fn read_bone_rest_rot(world: &World, bone_id: ComponentId) -> [f32; 4] {
    world
        .children_of(bone_id)
        .iter()
        .find_map(|&c| world.get_component_by_id_as::<BoneRestPoseComponent>(c))
        .map(|rest| rest.rotation)
        .or_else(|| {
            world
                .get_component_by_id_as::<TransformComponent>(bone_id)
                .map(|t| t.transform.rotation)
        })
        .unwrap_or([0.0, 0.0, 0.0, 1.0])
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::engine::ecs::CommandQueue;
    use crate::engine::ecs::IntentValue;
    use crate::engine::ecs::component::{
        BoneRestPoseComponent, ComponentRef, IKChainComponent, IKSolver, TransformComponent,
    };
    use slotmap::Key;

    #[derive(Default)]
    struct TestEmitter {
        intents: Vec<(ComponentId, IntentValue)>,
    }

    impl SignalEmitter for TestEmitter {
        fn push_event(&mut self, _scope: ComponentId, _event: crate::engine::ecs::EventSignal) {}

        fn push_intent(&mut self, scope: ComponentId, intent: crate::engine::ecs::IntentSignal) {
            self.intents.push((scope, intent.value));
        }
    }

    #[test]
    fn xr_driven_chain_skips_until_pose_is_valid() {
        let mut world = World::default();
        let driver = world.add_component(crate::engine::ecs::component::InputXRComponent::on());
        let target = world.add_component(TransformComponent::new());
        let root = world.add_component(TransformComponent::new());
        world.add_child(driver, target).unwrap();

        let mut chain = IKChainComponent::new(
            IKSolver::AimConstraint {
                offset_yaw: 0.0,
                copy_position: false,
                target_position_offset: [0.0; 3],
            },
            target,
            root,
        );
        chain.xr_pose_driver = Some(driver);
        let chain_id = world.add_component(chain);
        world.add_child(root, chain_id).unwrap();

        let mut emit = TestEmitter::default();
        tick_chain(chain_id, &mut world, &mut emit);
        assert!(emit.intents.is_empty());

        world
            .get_component_by_id_as_mut::<crate::engine::ecs::component::InputXRComponent>(driver)
            .unwrap()
            .pose_valid = true;
        tick_chain(chain_id, &mut world, &mut emit);
        assert!(!emit.intents.is_empty());
    }

    // Temporarily gated: see docs/bugs/ik-solver-api-drift-breaks-tests.md.
    #[cfg(any())]
    #[test]
    fn resolves_forward_reference_on_first_tick() {
        let mut w = World::default();

        // Author IKChain *before* the target/end_effector components exist.
        // Both ids start as null sentinels; only the sources carry the
        // selector strings.
        let ik_id = w.add_component(
            IKChainComponent::new(
                IKSolver::TwoBoneIK {
                    root_joint_id: ComponentId::null(),
                    mid_joint_id: ComponentId::null(),
                    pole_direction: [0.0, 1.0, 0.0],
                    copy_end_rotation: false,
                },
                ComponentId::null(),
                ComponentId::null(),
            )
            .with_target_source(ComponentRef::Query("#hand".to_string()))
            .with_end_effector_source(ComponentRef::Query("#elbow".to_string())),
        );

        // Now spawn the targets the IKChain refers to.
        let hand = w.add_component_boxed_named("hand", Box::new(TransformComponent::new()));
        let elbow = w.add_component_boxed_named("elbow", Box::new(TransformComponent::new()));

        // Sanity: nothing resolved yet.
        {
            let ik = w.get_component_by_id_as::<IKChainComponent>(ik_id).unwrap();
            assert!(ik.target_id.is_null());
            assert!(ik.end_effector_id.is_null());
        }

        // First tick triggers the deferred resolve.
        let mut emit = CommandQueue::new();
        let mut sys = IKSystem::new();
        sys.tick(&mut w, &mut emit, 0.016);

        let ik = w.get_component_by_id_as::<IKChainComponent>(ik_id).unwrap();
        assert_eq!(ik.target_id, hand);
        assert_eq!(ik.end_effector_id, elbow);
    }

    #[test]
    fn does_not_overwrite_already_resolved_ids() {
        let mut w = World::default();
        let pre_target = w.add_component(TransformComponent::new());
        let pre_ee = w.add_component(TransformComponent::new());
        let unrelated = w.add_component_boxed_named("hand", Box::new(TransformComponent::new()));

        let ik_id = w.add_component(
            IKChainComponent::new(
                IKSolver::AimConstraint {
                    offset_yaw: 0.0,
                    copy_position: false,
                    target_position_offset: [0.0, 0.0, 0.0],
                },
                pre_target,
                pre_ee,
            )
            // Source points at a different component named "hand" — but
            // since target_id / end_effector_id are already non-null, the
            // resolve pass should leave them alone.
            .with_target_source(ComponentRef::Query("#hand".to_string())),
        );

        let mut emit = CommandQueue::new();
        let mut sys = IKSystem::new();
        sys.tick(&mut w, &mut emit, 0.016);

        let ik = w.get_component_by_id_as::<IKChainComponent>(ik_id).unwrap();
        assert_eq!(ik.target_id, pre_target);
        assert_ne!(ik.target_id, unrelated);
        assert_eq!(ik.end_effector_id, pre_ee);
    }

    #[test]
    fn resolves_relative_parent_prefixed_sources() {
        let mut w = World::default();
        let root = w.add_component(TransformComponent::new());
        let hand = w.add_component_boxed_named("hand", Box::new(TransformComponent::new()));
        let elbow = w.add_component_boxed_named("elbow", Box::new(TransformComponent::new()));
        w.add_child(root, hand).unwrap();
        w.add_child(root, elbow).unwrap();

        let ik_id = w.add_component(
            IKChainComponent::new(
                IKSolver::TwoBoneIK {
                    root_joint_id: ComponentId::null(),
                    mid_joint_id: ComponentId::null(),
                    pole_direction: [0.0, 1.0, 0.0],
                    copy_end_rotation: false,
                },
                ComponentId::null(),
                ComponentId::null(),
            )
            .with_target_source(ComponentRef::Query("../#hand".to_string()))
            .with_end_effector_source(ComponentRef::Query("../#elbow".to_string())),
        );
        w.add_child(root, ik_id).unwrap();

        let mut emit = CommandQueue::new();
        let mut sys = IKSystem::new();
        sys.tick(&mut w, &mut emit, 0.016);

        let ik = w.get_component_by_id_as::<IKChainComponent>(ik_id).unwrap();
        assert_eq!(ik.target_id, hand);
        assert_eq!(ik.end_effector_id, elbow);
    }

    #[test]
    fn resolves_bare_sources_from_local_scope() {
        let mut w = World::default();

        let unrelated_root = w.add_component(TransformComponent::new());
        let unrelated_hand =
            w.add_component_boxed_named("hand", Box::new(TransformComponent::new()));
        w.add_child(unrelated_root, unrelated_hand).unwrap();

        let root = w.add_component(TransformComponent::new());

        let ik_id = w.add_component(
            IKChainComponent::new(
                IKSolver::TwoBoneIK {
                    root_joint_id: ComponentId::null(),
                    mid_joint_id: ComponentId::null(),
                    pole_direction: [0.0, 1.0, 0.0],
                    copy_end_rotation: false,
                },
                ComponentId::null(),
                ComponentId::null(),
            )
            .with_target_source(ComponentRef::Query("#hand".to_string()))
            .with_end_effector_source(ComponentRef::Query("#elbow".to_string())),
        );
        w.add_child(root, ik_id).unwrap();
        let hand = w.add_component_boxed_named("hand", Box::new(TransformComponent::new()));
        let elbow = w.add_component_boxed_named("elbow", Box::new(TransformComponent::new()));
        w.add_child(ik_id, hand).unwrap();
        w.add_child(ik_id, elbow).unwrap();

        let mut emit = CommandQueue::new();
        let mut sys = IKSystem::new();
        sys.tick(&mut w, &mut emit, 0.016);

        let ik = w.get_component_by_id_as::<IKChainComponent>(ik_id).unwrap();
        assert_eq!(ik.target_id, hand);
        assert_ne!(ik.target_id, unrelated_hand);
        assert_eq!(ik.end_effector_id, elbow);
    }

    #[test]
    fn two_bone_forearm_twist_uses_hand_rest_rotation_not_live_hand_local_rotation() {
        let mut w = World::default();
        let root = w.add_component(TransformComponent::new().with_position(0.0, 0.0, 0.0));
        let mid = w.add_component(TransformComponent::new().with_position(1.0, 0.0, 0.0));
        let hand = w.add_component(
            TransformComponent::new()
                .with_position(2.0, 0.0, 0.0)
                .with_rotation_quat(quat_from_axis_angle(
                    [1.0, 0.0, 0.0],
                    std::f32::consts::FRAC_PI_2,
                )),
        );
        let target = w.add_component(TransformComponent::new().with_position(2.0, 0.0, 0.0));
        let hand_rest = w.add_component(BoneRestPoseComponent::new(
            [0.0, 0.0, 0.0],
            [0.0, 0.0, 0.0, 1.0],
            [1.0, 1.0, 1.0],
        ));
        w.add_child(hand, hand_rest).unwrap();

        let mut emit = TestEmitter::default();
        solve_two_bone(
            &mut w,
            &mut emit,
            ComponentId::null(),
            &[root, mid, hand],
            target,
            [0.0, 1.0, 0.0],
            true,
            1.0,
            None,
        );

        let lower_update = emit
            .intents
            .iter()
            .find_map(|(scope, value)| match (scope, value) {
                (
                    scope_id,
                    IntentValue::UpdateTransform {
                        rotation_quat_xyzw, ..
                    },
                ) if *scope_id == mid => Some(*rotation_quat_xyzw),
                _ => None,
            })
            .expect("lower arm update intent");

        let (_, lower_angle) = quat_to_axis_angle(lower_update);
        assert!(
            lower_angle.abs() < 1e-4,
            "expected no forearm twist from live hand local rotation, got {lower_angle}"
        );
    }
}