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

use std::sync::Arc;

use base64::Engine;
use glam::Vec3;

use crate::{ResourceLoader, LoaderError, Mesh, Vertex, Model};

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>>, node: gltf::Node<'_>) -> Vec<Mesh> {
        let mut meshes = vec![];
        if let Some(mesh) = node.mesh() {
            for prim in mesh.primitives() {
                let pos_accessor = prim.get(&gltf::Semantic::Positions).unwrap();
                //assert_eq!(pos_accessor.dimensions(), gltf::accessor::Dimensions::Vec3); // TODO: dont do this
                
                let reader = prim.reader(|buf| Some(buffers[buf.index()].as_slice()));
                let pos: Vec<Vec3> = reader.read_positions()
                    .unwrap()
                    .map(|pos| Vec3::new(pos[0], pos[1], pos[2]))
                    .collect();

                let indices: Option<Vec<u32>> = reader.read_indices()
                    .map(|i| match i {
                        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(),
                    });

                let mut new_mesh = Mesh::default();
                new_mesh.vertices = pos.into_iter().map(|p| Vertex::new(p, glam::Vec2::new(0.0, 0.0))).collect();
                new_mesh.indices = indices;

                meshes.push(new_mesh);
            }
        }

        for child in node.children() {
            let mut child_meshes = self.process_node(buffers, 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<gltf::Buffer<'_>> = gltf.buffers().collect();

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

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

        Ok(Arc::new(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();
        let mesh = &model.meshes[0];
        assert!(mesh.vertices.len() > 0);
        assert!(mesh.indices.clone().unwrap().len() > 0);
    }
}