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

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

use glam::{Quat, Vec3};
use instant::Instant;
use lyra_math::Transform;
use thiserror::Error;

use crate::{gltf::GltfScene, util, LoaderError, ResHandle, ResourceLoader, ResourceManager};
use super::{GltfNode, Material, Mesh, MeshIndices, MeshVertexAttribute, VertexAttributeData};

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(ctx: &mut GltfLoadContext, materials: &Vec<ResHandle<Material>>, gnode: gltf::Node<'_>) -> GltfNode {
        let mut node = GltfNode::default();

        node.transform = {
            let gt = gnode.transform();
            let (pos, rot, scale) = gt.decomposed();

            Transform::new(Vec3::from(pos), Quat::from_array(rot), Vec3::from(scale))
        };
        node.name = gnode.name().map(str::to_string);

        if let Some(mesh) = gnode.mesh() {
            let mut new_mesh = Mesh::default();

            for prim in mesh.primitives() {
                let reader = prim.reader(|buf| Some(ctx.buffers[buf.index()].as_slice()));

                // 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.material = Some(mat.clone());
            }

            let handle = ResHandle::with_data("", new_mesh);
            ctx.resource_manager.store_uuid(handle.clone());
            node.mesh = Some(handle);
        }

        for child in gnode.children() {
            let cmesh = ModelLoader::process_node(ctx, materials, child);
            node.children.push(cmesh);
        }
        
        node
    }
}

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 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();

        let mut gltf_out = super::Gltf::default();

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

        let start_inst = Instant::now();
        let materials: Vec<ResHandle<Material>> = gltf.materials()
            .map(|mat| ResHandle::with_data("", Material::from_gltf(&mut context, mat)))
            .collect();
        let mat_time = Instant::now() - start_inst;
        debug!("Loaded {} materials in {}s", materials.len(), mat_time.as_secs_f32());

        for (idx, scene) in gltf.scenes().enumerate() {
            let start_inst = Instant::now();
            let nodes: Vec<GltfNode> = scene.nodes()
                .map(|node| ModelLoader::process_node(&mut context, &materials, node))
                .collect();
            let node_time = Instant::now() - start_inst;
    
            debug!("Loaded {} nodes in the scene in {}s", nodes.len(), node_time.as_secs_f32());

            for mesh in nodes.iter().map(|n| &n.mesh) {
                if let Some(mesh) = mesh {
                    gltf_out.meshes.push(mesh.clone());
                }
            }

            let scene = GltfScene {
                nodes,
            };
            let scene = ResHandle::with_data(&format!("{}:Scene{}", path, idx), scene);
            gltf_out.scenes.push(scene);
        }

        gltf_out.materials = materials;
        
        Ok(Arc::new(ResHandle::with_data(path, gltf_out)))
    }

    #[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::{gltf::Gltf, 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 gltf = loader.load(&mut manager, &path).unwrap();
        let gltf = Arc::downcast::<ResHandle<Gltf>>(gltf.as_arc_any()).unwrap();
        let gltf = gltf.data_ref().unwrap();

        assert_eq!(gltf.scenes.len(), 1);
        let scene = &gltf.scenes[0]
            .data_ref().unwrap();

        assert_eq!(scene.nodes.len(), 1);
        let mnode = &scene.nodes[0];

        assert!(mnode.mesh.is_some());
        assert_eq!(mnode.transform, Transform::from_xyz(0.0, 0.0, 0.0));
        assert_eq!(mnode.children.len(), 0);

        let mesh = mnode.mesh.as_ref().unwrap();
        let mesh = mesh.data_ref().unwrap();
        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.as_ref().unwrap().data_ref().unwrap().base_color_texture.is_some());
    }
}