lyra-engine/lyra-resource/src/loader/model.rs

use std::sync::Arc;

use base64::Engine;

use crate::{ResourceLoader, LoaderError, Mesh, Model, MeshVertexAttribute, VertexAttributeData, Resource, Material, PbrRoughness};

impl From<gltf::Error> for LoaderError {
    fn from(value: gltf::Error) -> Self {
        LoaderError::DecodingError(value.into())
    }
}

#[derive(Default)]
pub struct ModelLoader;

impl ModelLoader {
    fn parse_uri(uri: &str) -> Option<Vec<u8>> {
        let uri = uri.strip_prefix("data")?;
        let (mime, data) = uri.split_once(",")?;

        let (_mime, is_base64) = match mime.strip_suffix(";base64") {
            Some(mime) => (mime, true),
            None => (mime, false),
        };

        if is_base64 {
            base64::engine::general_purpose::STANDARD.decode(data).ok()
        } else {
            Some(data.as_bytes().to_vec())
        }
    }

    fn process_node(&self, buffers: &Vec<Vec<u8>>, materials: &Vec<Material>, node: gltf::Node<'_>) -> Vec<Mesh> {
        let mut meshes = vec![];
        if let Some(mesh) = node.mesh() {
            for prim in mesh.primitives() {
                let reader = prim.reader(|buf| Some(buffers[buf.index()].as_slice()));

                let mut new_mesh = Mesh::default();
                
                // read the positions
                if let Some(pos) = reader.read_positions() {
                    let pos: Vec<glam::Vec3> = pos.map(|t| t.into()).collect();
                    new_mesh.add_attribute(MeshVertexAttribute::Position, VertexAttributeData::Vec3(pos));
                }

                // read the normals
                if let Some(norms) = reader.read_normals() {
                    let norms: Vec<glam::Vec3> = norms.map(|t| t.into()).collect();
                    new_mesh.add_attribute(MeshVertexAttribute::Normals, VertexAttributeData::Vec3(norms));
                }

                // read the tangents
                if let Some(tangents) = reader.read_tangents() {
                    let tangents: Vec<glam::Vec4> = tangents.map(|t| t.into()).collect();
                    new_mesh.add_attribute(MeshVertexAttribute::Tangents, VertexAttributeData::Vec4(tangents));
                }

                // read tex coords
                if let Some(tex_coords) = reader.read_tex_coords(0) {
                    let tex_coords: Vec<glam::Vec2> = tex_coords.into_f32().map(|t| t.into()).collect(); // TODO: u16, u8
                    new_mesh.add_attribute(MeshVertexAttribute::TexCoords, VertexAttributeData::Vec2(tex_coords));
                }

                // read the indices
                if let Some(indices) = reader.read_indices() {
                    let indices: Vec<u32> = match indices {
                        gltf::mesh::util::ReadIndices::U8(i) => i.map(|i| i as u32).collect(),
                        gltf::mesh::util::ReadIndices::U16(i) => i.map(|i| i as u32).collect(),
                        gltf::mesh::util::ReadIndices::U32(i) => i.collect(),
                    };

                    new_mesh.indices = Some(indices);
                }

                prim.material().

                meshes.push(new_mesh);
            }
        }

        for child in node.children() {
            let mut child_meshes = self.process_node(buffers, materials, child);
            meshes.append(&mut child_meshes);
        }
        
        meshes
    }
}

impl ResourceLoader for ModelLoader {
    fn extensions(&self) -> &[&str] {
        &[
            "gltf"
        ]
    }

    fn load(&self, path: &str) -> Result<std::sync::Arc<dyn crate::ResourceStorage>, crate::LoaderError> {
        // check if the file is supported by this loader
        if !self.does_support_file(path) {
            return Err(LoaderError::UnsupportedExtension(path.to_string()));
        }

        let gltf = gltf::Gltf::open(path)?;

        let buffers: Vec<Vec<u8>> = gltf.buffers().map(|b| match b.source() {
                gltf::buffer::Source::Bin => gltf.blob.as_deref().map(|v| v.to_vec()),
                gltf::buffer::Source::Uri(uri) => ModelLoader::parse_uri(uri),
            }).flatten().collect();

        // TODO: Read in multiple scenes
        let scene = gltf.scenes().next().unwrap();

        // Load the materials
        let materials: Vec<Material> = gltf.materials()
            .map(|mat| Material {
                double_sided: mat.double_sided(),
                name: mat.name().map(|s| s.to_string()),
                shader_uuid: None,
                pbr: Some(PbrRoughness::from(mat.pbr_metallic_roughness())),
                texture: None
            }).collect();

        let meshes: Vec<Mesh> = scene.nodes()
            .map(|node| self.process_node(&buffers, &materials, node))
            .flatten().collect();

        Ok(Arc::new(Resource::with_data(path, Model::new(meshes))))
    }
}

#[cfg(test)]
mod tests {
    use crate::ResourceLoader;
    use super::*;

    fn test_file_path(path: &str) -> String {
        let manifest = std::env::var("CARGO_MANIFEST_DIR").unwrap();
        
        format!("{manifest}/test_files/gltf/{path}")
    }

    #[test]
    fn test_loading() {
        let path = test_file_path("test-embedded.gltf");

        let loader = ModelLoader::default();
        let model = loader.load(&path).unwrap();
        let model = Arc::downcast::<Model>(model.as_arc_any()).unwrap();
        assert_eq!(model.meshes.len(), 1); // There should only be 1 mesh
        let mesh = &model.meshes[0];
        assert!(mesh.position().unwrap().len() > 0);
        assert!(mesh.normals().unwrap().len() > 0);
        assert!(mesh.tex_coords().unwrap().len() > 0);
        assert!(mesh.indices.clone().unwrap().len() > 0);
    }
}