roast2d_internal 0.4.0

//! Skinned mesh for skeletal animation.

use std::path::Path;

use wgpu::util::DeviceExt;

use crate::engine::Engine;
use glam::Mat4;

use crate::skeleton::{
    Animation, AnimationChannel, AnimationPlayer, Bone, BoneTransform, Keyframe, MAX_BONES,
    Skeleton,
};

/// Vertex with bone weights for skeletal animation.
///
/// Extends the basic Mesh3DVertex with bone indices and weights
/// for GPU skinning.
#[repr(C)]
#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable, Debug)]
pub struct SkinnedMesh3DVertex {
    pub position: [f32; 3],
    pub normal: [f32; 3],
    pub uv: [f32; 2],
    /// Bone indices (up to 4 influences per vertex)
    pub bone_indices: [u32; 4],
    /// Bone weights (must sum to 1.0)
    pub bone_weights: [f32; 4],
}

impl SkinnedMesh3DVertex {
    /// Create a new skinned vertex.
    pub fn new(
        position: glam::Vec3,
        normal: glam::Vec3,
        uv: glam::Vec2,
        bone_indices: [u8; 4],
        bone_weights: [f32; 4],
    ) -> Self {
        Self {
            position: position.to_array(),
            normal: normal.to_array(),
            uv: uv.to_array(),
            bone_indices: [
                bone_indices[0] as u32,
                bone_indices[1] as u32,
                bone_indices[2] as u32,
                bone_indices[3] as u32,
            ],
            bone_weights,
        }
    }
}

/// A skinned mesh with skeleton and animations for skeletal animation.
///
/// This type combines a mesh with bone weights, a skeleton hierarchy,
/// and animation clips. Use with `AnimationPlayer` to play animations.
///
/// # Example
/// ```ignore
/// // Load a skinned mesh from glTF
/// let mesh = SkinnedMesh3D::from_gltf(g, "character.glb")?;
///
/// // Create an animation player
/// let mut player = mesh.create_player();
/// player.play(0); // Play first animation
///
/// // In update loop
/// player.update(g.tick, &mesh.skeleton, &mesh.animations);
///
/// // In draw loop
/// g.draw3d_skinned(Draw3DSkinned::new(mesh.clone(), Mat4::IDENTITY, &player));
/// ```
#[derive(Clone)]
pub struct SkinnedMesh3D {
    pub vertex: wgpu::Buffer,
    pub index: wgpu::Buffer,
    pub index_count: u32,
    pub skeleton: Skeleton,
    pub animations: Vec<Animation>,
}

impl SkinnedMesh3D {
    /// Create a new skinned mesh from vertices, indices, skeleton, and animations.
    pub fn new(
        g: &Engine,
        vertices: &[SkinnedMesh3DVertex],
        indices: &[u16],
        skeleton: Skeleton,
        animations: Vec<Animation>,
    ) -> Self {
        let state = g.backend_state();
        let vertex = state
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("skinned_mesh3d.vertices"),
                contents: bytemuck::cast_slice(vertices),
                usage: wgpu::BufferUsages::VERTEX,
            });
        let index = state
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("skinned_mesh3d.indices"),
                contents: bytemuck::cast_slice(indices),
                usage: wgpu::BufferUsages::INDEX,
            });
        Self {
            vertex,
            index,
            index_count: indices.len() as u32,
            skeleton,
            animations,
        }
    }

    /// Create an AnimationPlayer for this mesh.
    pub fn create_player(&self) -> AnimationPlayer {
        AnimationPlayer::new()
    }

    /// Get animation index by name.
    pub fn animation_index(&self, name: &str) -> Option<usize> {
        self.animations.iter().position(|a| a.name == name)
    }

    /// Load a skinned mesh with skeleton and animations from glTF.
    ///
    /// The glTF file must contain at least one skinned mesh (mesh with skin).
    pub fn from_gltf<P: AsRef<Path>>(g: &Engine, path: P) -> anyhow::Result<Self> {
        let (document, buffers, _images) = gltf::import(path.as_ref())?;
        Self::from_gltf_document(g, &document, &buffers)
    }

    /// Load a skinned mesh from glTF/GLB data in memory.
    pub fn from_gltf_bytes(g: &Engine, data: &[u8]) -> anyhow::Result<Self> {
        let (document, buffers, _images) = gltf::import_slice(data)?;
        Self::from_gltf_document(g, &document, &buffers)
    }

    fn from_gltf_document(
        g: &Engine,
        document: &gltf::Document,
        buffers: &[gltf::buffer::Data],
    ) -> anyhow::Result<Self> {
        // Find the first node with a skin
        let (node, skin) = document
            .nodes()
            .find_map(|n| n.skin().map(|s| (n, s)))
            .ok_or_else(|| anyhow::anyhow!("No skinned mesh found in glTF"))?;

        let mesh = node
            .mesh()
            .ok_or_else(|| anyhow::anyhow!("Skinned node has no mesh"))?;

        // Load skeleton
        let skeleton = Self::load_skeleton(&skin, buffers)?;

        // Load animations
        let animations = Self::load_animations(document, &skin, buffers)?;

        // Load mesh with skin weights
        let primitive = mesh
            .primitives()
            .next()
            .ok_or_else(|| anyhow::anyhow!("Mesh has no primitives"))?;

        if primitive.mode() != gltf::mesh::Mode::Triangles {
            anyhow::bail!("Only triangle primitives are supported");
        }

        let reader = primitive.reader(|buffer| Some(&buffers[buffer.index()]));

        let positions: Vec<[f32; 3]> = reader
            .read_positions()
            .ok_or_else(|| anyhow::anyhow!("No positions"))?
            .collect();

        let normals: Vec<[f32; 3]> = reader
            .read_normals()
            .map(|i| i.collect())
            .unwrap_or_else(|| vec![[0.0, 1.0, 0.0]; positions.len()]);

        let uvs: Vec<[f32; 2]> = reader
            .read_tex_coords(0)
            .map(|i| i.into_f32().collect())
            .unwrap_or_else(|| vec![[0.0, 0.0]; positions.len()]);

        // Read joint indices
        let joints: Vec<[u16; 4]> = reader
            .read_joints(0)
            .map(|i| i.into_u16().collect())
            .unwrap_or_else(|| vec![[0, 0, 0, 0]; positions.len()]);

        // Read weights
        let weights: Vec<[f32; 4]> = reader
            .read_weights(0)
            .map(|i| i.into_f32().collect())
            .unwrap_or_else(|| vec![[1.0, 0.0, 0.0, 0.0]; positions.len()]);

        let indices: Vec<u16> = reader
            .read_indices()
            .map(|i| i.into_u32().map(|x| x as u16).collect())
            .unwrap_or_else(|| (0..positions.len() as u16).collect());

        // Build vertices
        let vertices: Vec<SkinnedMesh3DVertex> = positions
            .iter()
            .zip(normals.iter())
            .zip(uvs.iter())
            .zip(joints.iter())
            .zip(weights.iter())
            .map(|((((pos, normal), uv), joint), weight)| {
                SkinnedMesh3DVertex::new(
                    glam::Vec3::from_array(*pos),
                    glam::Vec3::from_array(*normal),
                    glam::Vec2::from_array(*uv),
                    [
                        joint[0] as u8,
                        joint[1] as u8,
                        joint[2] as u8,
                        joint[3] as u8,
                    ],
                    *weight,
                )
            })
            .collect();

        Ok(Self::new(g, &vertices, &indices, skeleton, animations))
    }

    fn load_skeleton(
        skin: &gltf::Skin,
        buffers: &[gltf::buffer::Data],
    ) -> anyhow::Result<Skeleton> {
        let joints: Vec<_> = skin.joints().collect();

        if joints.len() > MAX_BONES {
            anyhow::bail!(
                "Skeleton has {} bones, but maximum is {}",
                joints.len(),
                MAX_BONES
            );
        }

        // Load inverse bind matrices
        let inverse_bind_matrices: Vec<Mat4> = if let Some(accessor) = skin.inverse_bind_matrices()
        {
            let view = accessor
                .view()
                .ok_or_else(|| anyhow::anyhow!("No buffer view"))?;
            let buffer_data = &buffers[view.buffer().index()];
            let offset = view.offset() + accessor.offset();
            let stride = accessor.size();

            (0..accessor.count())
                .map(|i| {
                    let start = offset + i * stride;
                    let data = &buffer_data[start..start + 64];
                    let floats: [f32; 16] = bytemuck::cast_slice(data).try_into().unwrap();
                    Mat4::from_cols_array(&floats)
                })
                .collect()
        } else {
            vec![Mat4::IDENTITY; joints.len()]
        };

        // Build joint index map (glTF node index -> joint array index)
        let joint_indices: hashbrown::HashMap<usize, usize> = joints
            .iter()
            .enumerate()
            .map(|(i, j)| (j.index(), i))
            .collect();

        // Build parent map by traversing node children
        // The gltf crate doesn't expose parent directly, so we build it from children
        let mut parent_map: hashbrown::HashMap<usize, usize> = hashbrown::HashMap::new();
        for joint in &joints {
            for child in joint.children() {
                parent_map.insert(child.index(), joint.index());
            }
        }

        let mut bones = Vec::with_capacity(joints.len());
        let mut roots = Vec::new();

        for (i, joint) in joints.iter().enumerate() {
            // Find parent in joint list using the parent map
            let parent = parent_map
                .get(&joint.index())
                .and_then(|pi| joint_indices.get(pi).copied());

            if parent.is_none() {
                roots.push(i);
            }

            let (translation, rotation, scale) = joint.transform().decomposed();

            bones.push(Bone {
                name: joint.name().unwrap_or("").to_string(),
                parent,
                inverse_bind_matrix: inverse_bind_matrices
                    .get(i)
                    .copied()
                    .unwrap_or(Mat4::IDENTITY),
                local_bind_pose: BoneTransform::new(
                    glam::Vec3::from_array(translation),
                    glam::Quat::from_array(rotation),
                    glam::Vec3::from_array(scale),
                ),
            });
        }

        Ok(Skeleton { bones, roots })
    }

    fn load_animations(
        document: &gltf::Document,
        skin: &gltf::Skin,
        buffers: &[gltf::buffer::Data],
    ) -> anyhow::Result<Vec<Animation>> {
        let joints: Vec<_> = skin.joints().collect();
        let joint_indices: hashbrown::HashMap<usize, usize> = joints
            .iter()
            .enumerate()
            .map(|(i, j)| (j.index(), i))
            .collect();

        let mut animations = Vec::new();

        for anim in document.animations() {
            let mut channels = Vec::new();
            let mut duration = 0.0f32;

            for channel in anim.channels() {
                let target = channel.target();
                let node_idx = target.node().index();

                // Skip if not a joint in our skin
                let Some(&bone_idx) = joint_indices.get(&node_idx) else {
                    continue;
                };

                let reader = channel.reader(|buffer| Some(&buffers[buffer.index()]));

                let times: Vec<f32> = reader
                    .read_inputs()
                    .ok_or_else(|| anyhow::anyhow!("No input times"))?
                    .collect();

                duration = duration.max(times.last().copied().unwrap_or(0.0));

                match target.property() {
                    gltf::animation::Property::Translation => {
                        let outputs = reader
                            .read_outputs()
                            .ok_or_else(|| anyhow::anyhow!("No outputs"))?;
                        if let gltf::animation::util::ReadOutputs::Translations(iter) = outputs {
                            let values: Vec<glam::Vec3> =
                                iter.map(glam::Vec3::from_array).collect();
                            let keyframes: Vec<Keyframe<glam::Vec3>> = times
                                .iter()
                                .zip(values.iter())
                                .map(|(&t, &v)| Keyframe::new(t, v))
                                .collect();
                            channels.push(AnimationChannel::Translation {
                                bone_index: bone_idx,
                                keyframes,
                            });
                        }
                    }
                    gltf::animation::Property::Rotation => {
                        let outputs = reader
                            .read_outputs()
                            .ok_or_else(|| anyhow::anyhow!("No outputs"))?;
                        if let gltf::animation::util::ReadOutputs::Rotations(iter) = outputs {
                            let values: Vec<glam::Quat> = iter
                                .into_f32()
                                .map(|[x, y, z, w]| glam::Quat::from_xyzw(x, y, z, w))
                                .collect();
                            let keyframes: Vec<Keyframe<glam::Quat>> = times
                                .iter()
                                .zip(values.iter())
                                .map(|(&t, &v)| Keyframe::new(t, v))
                                .collect();
                            channels.push(AnimationChannel::Rotation {
                                bone_index: bone_idx,
                                keyframes,
                            });
                        }
                    }
                    gltf::animation::Property::Scale => {
                        let outputs = reader
                            .read_outputs()
                            .ok_or_else(|| anyhow::anyhow!("No outputs"))?;
                        if let gltf::animation::util::ReadOutputs::Scales(iter) = outputs {
                            let values: Vec<glam::Vec3> =
                                iter.map(glam::Vec3::from_array).collect();
                            let keyframes: Vec<Keyframe<glam::Vec3>> = times
                                .iter()
                                .zip(values.iter())
                                .map(|(&t, &v)| Keyframe::new(t, v))
                                .collect();
                            channels.push(AnimationChannel::Scale {
                                bone_index: bone_idx,
                                keyframes,
                            });
                        }
                    }
                    _ => {} // Ignore morph targets
                }
            }

            if !channels.is_empty() {
                animations.push(Animation {
                    name: anim.name().unwrap_or("").to_string(),
                    duration,
                    channels,
                });
            }
        }

        Ok(animations)
    }
}