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

use std::{ffi::OsStr, path::{Path, PathBuf}, sync::Arc};

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

use crate::{loader::{LoaderError, PinedBoxLoaderFuture, ResourceLoader}, util, ResHandle, ResourceData, ResourceManager, ResourceStorage};
use super::{Gltf, 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: 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>>, scene: &mut SceneGraph, scene_parent: &SceneNode, 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);

        let scene_node = scene.add_node_under(scene_parent, node.transform, ());

        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::new_ready(None, new_mesh);
            ctx.resource_manager.store_uuid(handle.clone());
            node.mesh = Some(handle.clone());
            scene.insert(&scene_node, handle);
        }

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

    fn extensions() -> &'static [&'static str] {
        &[
            "gltf", "glb"
        ]
    }

    fn does_support_file(path: &str) -> bool {
        match Path::new(path).extension().and_then(OsStr::to_str) {
            Some(ext) => {
                Self::extensions().contains(&ext)
            },
            _ => false,
        }
    }
}

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

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

    fn load(&self, resource_manager: ResourceManager, path: &str) -> PinedBoxLoaderFuture {
        // cant use &str across async
        let path = path.to_string();

        Box::pin(async move {
            // 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: resource_manager.clone(),
                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::new_ready(None, 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 mut graph = SceneGraph::new();
                let root_node = graph.root_node();

                for node in scene.nodes() {
                    let n = ModelLoader::process_node(&mut context, &materials, &mut graph, &root_node, node);

                    if let Some(mesh) = n.mesh {
                        gltf_out.meshes.push(mesh.clone());
                    }
                }

                let graph = ResHandle::new_ready(Some(path.as_str()), graph);
                gltf_out.scenes.push(graph);

                /* 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::new_ready(Some(path.as_str()), scene);
                gltf_out.scenes.push(scene); */
            }

            gltf_out.materials = materials;
            
            Ok(Box::new(gltf_out) as Box<dyn ResourceData>)
        })
    }

    #[allow(unused_variables)]
    fn load_bytes(&self, resource_manager: ResourceManager, bytes: Vec<u8>, offset: usize, length: usize) -> PinedBoxLoaderFuture {
        todo!()
    }

    fn create_erased_handle(&self) -> Arc<dyn ResourceStorage> {
        Arc::from(ResHandle::<Gltf>::new_loading(None))
    }
}

#[cfg(test)]
mod tests {
    use lyra_ecs::{query::Entities, relation::ChildOf};

    use crate::tests::busy_wait_resource;

    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 manager = ResourceManager::new();
        let gltf = manager.request::<Gltf>(&path).unwrap();
        busy_wait_resource(&gltf, 15.0);
        let gltf = gltf.data_ref().unwrap();

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

        let mut node = None;
        scene.traverse_down(|no, _tran| {
            node = Some(no.clone());
        });

        let world = scene.world();
        let node = node.unwrap();
        
        let data = world.view_one::<(&ResHandle<Mesh>, &Transform)>(node.entity()).get();
        debug_assert!(data.is_some(), "The mesh was not loaded"); // transform will always be there
        let data = data.unwrap();

        // ensure there are no children of the node
        assert_eq!(
            world.view::<Entities>()
                .relates_to::<ChildOf>(node.entity())
                .into_iter()
                .count(),
            0
        );
        
        assert_eq!(*data.1, Transform::from_xyz(0.0, 0.0, 0.0));
        
        let mesh = data.0;
        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());
    }
}