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

use std::{sync::Arc, path::PathBuf};

use thiserror::Error;

use crate::{ResourceLoader, LoaderError, Mesh, Model, MeshVertexAttribute, VertexAttributeData, ResHandle, Material, MeshIndices, ResourceManager, util};

use tracing::debug;

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

#[derive(Error, Debug)]
enum ModelLoaderError {
    #[error("The model ({0}) is missing the BIN section in the gltf file")]
    MissingBin(String),
    #[error("There was an error with decoding a uri defined in the model: '{0}'")]
    UriDecodingError(util::UriReadError),
}

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

#[allow(dead_code)]
pub(crate) struct GltfLoadContext<'a> {
    pub resource_manager: &'a mut ResourceManager,
    pub gltf: &'a gltf::Gltf,
    /// Path to the gltf
    pub gltf_path: &'a str,
    /// The path to the directory that the gltf is contained in.
    pub gltf_parent_path: &'a str,
    /// List of buffers in the gltf
    pub buffers: &'a Vec<Vec<u8>>,
}

#[derive(Default)]
pub struct ModelLoader;

impl ModelLoader {
    /* fn parse_uri(containing_path: &str, 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 {
            Some(base64::engine::general_purpose::STANDARD.decode(data).unwrap())
        } else {
            let full_path = format!("{containing_path}/{data}");
            let buf = std::fs::read(&full_path).unwrap();
            Some(buf)
        }
    } */

    fn process_node(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() {
                    if prim.mode() != gltf::mesh::Mode::Triangles {
                        todo!("Load position primitives that aren't triangles"); // TODO
                    }

                    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();
                    new_mesh.add_attribute(MeshVertexAttribute::TexCoords, VertexAttributeData::Vec2(tex_coords));
                }

                // read the indices
                if let Some(indices) = reader.read_indices() {
                    let indices: MeshIndices = match indices {
                        // wpgu doesn't support u8 indices, so those must be converted to u16
                        gltf::mesh::util::ReadIndices::U8(i) => MeshIndices::U16(i.map(|i| i as u16).collect()),
                        gltf::mesh::util::ReadIndices::U16(i) => MeshIndices::U16(i.collect()),
                        gltf::mesh::util::ReadIndices::U32(i) => MeshIndices::U32(i.collect()),
                    };

                    new_mesh.indices = Some(indices);
                }

                let mat = materials.get(prim.material().index().unwrap()).unwrap();
                new_mesh.set_material(mat.clone());

                meshes.push(new_mesh);
            }
        }

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

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

    fn mime_types(&self) -> &[&str] {
        &[]
    }

    fn load(&self, resource_manager: &mut ResourceManager, 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 mut parent_path = PathBuf::from(path);
        parent_path.pop();
        let parent_path = parent_path.display().to_string();

        /* let (document, buffers, images) = gltf::import(&path)?;
        let buffers: Vec<Vec<u8>> = buffers.into_iter().map(|b| b.0).collect();

        let scene = document.scenes().next().unwrap();

        let materials: Vec<Material> = document.materials()
            .map(|mat| Material::from_gltf(resource_manager, &parent_path, mat)).collect();

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

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

        let mut use_bin = false;
        let buffers: Vec<Vec<u8>> = gltf.buffers().flat_map(|b| match b.source() {
                gltf::buffer::Source::Bin => {
                    use_bin = true;
                    gltf.blob.as_deref().map(|v| v.to_vec())
                        .ok_or(ModelLoaderError::MissingBin(path.to_string()))
                },
                gltf::buffer::Source::Uri(uri) => util::gltf_read_buffer_uri(&parent_path, uri)
                    .map_err(ModelLoaderError::UriDecodingError),
            }).collect();

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

        let mut context = GltfLoadContext {
            resource_manager,
            gltf: &gltf,
            gltf_path: path,
            gltf_parent_path: &parent_path,
            buffers: &buffers,
        };

        let materials: Vec<Material> = gltf.materials()
            .map(|mat| Material::from_gltf(&mut context, mat)).collect();

        let meshes: Vec<Mesh> = scene.nodes()
            .flat_map(|node| ModelLoader::process_node(&buffers, &materials, node))
            .collect();
        debug!("Loaded {} meshes, and {} materials from '{}'", meshes.len(), materials.len(), path);

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

    #[allow(unused_variables)]
    fn load_bytes(&self, resource_manager: &mut ResourceManager, bytes: Vec<u8>, offset: usize, length: usize) -> Result<Arc<dyn crate::ResourceStorage>, LoaderError> {
        todo!()
    }
}

#[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("texture-embedded.gltf");

        let mut manager = ResourceManager::new();
        let loader = ModelLoader::default();
        let model = loader.load(&mut manager, &path).unwrap();
        let model = Arc::downcast::<ResHandle<Model>>(model.as_arc_any()).unwrap();
        let model = model.data_ref();
        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);
        assert!(mesh.material().base_color_texture.is_some());
        let _mesh_mat = mesh.material(); // inner panic if material was not loaded
    }
}