crystal_ball 0.3.0

use std::collections::{BTreeSet, HashMap};
use std::convert::{TryFrom, TryInto};
use std::path::Path;
use std::sync::Arc;

use gltf::mesh::Reader;
use gltf::texture::MagFilter;
use gltf::{Buffer, Document, Node};

use crate::color::{Color, ColorFormat, Image, Interpolation, Texture};
use crate::materials::{Material, PbrMaterial};
use crate::math::{Bounds3, Hit, Mat4, Point2, Point3, Ray, Vec3, Vec4};
use crate::shapes::triangle_mesh::TriangleMesh;
use crate::shapes::{Mesh, Shape};
use crate::util::Error;

/// A [`Shape`] paired with a [`Material`].
#[derive(Clone)]
pub struct Object {
    // TODO: Arc ist pointless, cause you need to create a new shape if you want to change it's transform
    pub shape: Arc<dyn Shape>,
    pub material: Arc<dyn Material>,
}

impl Object {
    /// Create a new [`Object`].
    pub fn new(shape: Arc<dyn Shape>, material: Arc<dyn Material>) -> Self {
        Object { shape, material }
    }

    /// Propagate the transform for all children of a parent recursively.
    fn apply_transform_recursively(transform_matrices: &mut [Mat4], node: &Node) {
        node.children().for_each(|c| {
            transform_matrices[c.index()] =
                transform_matrices[node.index()] * transform_matrices[c.index()];

            Self::apply_transform_recursively(transform_matrices, &c);
        });
    }

    /// Load the index buffer of a mesh.
    ///
    /// Returns [`None`] if there is no index buffer.
    fn load_indices<'a, 's, F: Clone + Fn(Buffer<'a>) -> Option<&'s [u8]>>(
        reader: &Reader<'a, 's, F>,
        node: &Node,
    ) -> Option<Vec<u32>> {
        match reader.read_indices() {
            Some(indices) => Some(indices.into_u32().collect()),
            None => {
                eprint!("ERROR: Indices not found. ");
                match node.name() {
                    Some(name) => eprintln!("Skipping '{}'", name),
                    None => eprintln!("Skipping mesh"),
                }

                None
            }
        }
    }

    /// Load the vertex buffer of a mesh.
    ///
    /// Returns [`None`] if there is no vertex buffer.
    fn load_vertices<'a, 's, F: Clone + Fn(Buffer<'a>) -> Option<&'s [u8]>>(
        reader: &Reader<'a, 's, F>,
        matrix: Mat4,
        node: &Node,
    ) -> Option<Vec<Point3>> {
        let vertices_vec = match reader.read_positions() {
            Some(vertices) => vertices
                .map(|v| matrix * Point3::new(v[0] as f64, v[1] as f64, v[2] as f64))
                .collect(),
            None => {
                eprint!("ERROR: Attribute POSITION not found. ");
                match node.name() {
                    Some(name) => eprintln!("Skipping '{}'", name),
                    None => eprintln!("Skipping mesh"),
                }

                return None;
            }
        };

        Some(vertices_vec)
    }

    /// Load the normal buffer of a mesh.
    ///
    /// Returns [`None`] if there is no normal buffer.
    fn load_normals<'a, 's, F: Clone + Fn(Buffer<'a>) -> Option<&'s [u8]>>(
        reader: &Reader<'a, 's, F>,
        matrix: Mat4,
        node: &Node,
    ) -> Option<Vec<Vec3>> {
        match reader.read_normals() {
            Some(normals) => Some(
                normals
                    .map(|n| {
                        (matrix * Vec3::new(n[0] as f64, n[1] as f64, n[2] as f64)).normalize()
                    })
                    .collect(),
            ),
            None => {
                eprint!("WARNING: Attribute NORMAL not found. ");
                match node.name() {
                    Some(name) => eprintln!("Generating normals for '{}'", name),
                    None => eprintln!("Generating normals"),
                }
                None
            }
        }
    }

    /// Load the tangent buffer of a mesh.
    ///
    /// Returns [`None`] if there is no tangent buffer.
    fn load_tangents<'a, 's, F: Clone + Fn(Buffer<'a>) -> Option<&'s [u8]>>(
        reader: &Reader<'a, 's, F>,
        matrix: Mat4,
        node: &Node,
    ) -> Option<Vec<Vec4>> {
        match reader.read_tangents() {
            Some(tangents) => Some(
                tangents
                    .map(|t| {
                        let tangent =
                            (matrix * Vec3::new(t[0] as f64, t[1] as f64, t[2] as f64)).normalize();

                        Vec4::new(tangent.x, tangent.y, tangent.z, t[3] as f64)
                    })
                    .collect(),
            ),
            None => {
                eprint!("WARNING: Attribute TANGENT not found. ");
                match node.name() {
                    Some(name) => eprintln!("Normal maps won't be supported for '{}'", name),
                    None => eprintln!("Normal maps won't be supported"),
                }
                None
            }
        }
    }

    /// Load the UV buffer of a mesh.
    ///
    /// Returns [`None`] if there is no UV buffer.
    fn load_uvs<'a, 's, F: Clone + Fn(Buffer<'a>) -> Option<&'s [u8]>>(
        reader: &Reader<'a, 's, F>,
        node: &Node,
    ) -> Option<Vec<Point2>> {
        match reader.read_tex_coords(0) {
            Some(tex_coords) => Some(
                tex_coords
                    .into_f32()
                    .map(|t| Point2::new(t[0] as f64, t[1] as f64))
                    .collect(),
            ),
            None => {
                eprint!("WARNING: Attribute TEXCOORD not found. ");
                match node.name() {
                    Some(name) => {
                        eprintln!("'{}' will only support colors, not textures", name)
                    }
                    None => eprintln!("Mesh will only support colors, not textures"),
                }
                None
            }
        }
    }

    /// Load all images from a glTF document and store them in a [`HashMap`].
    ///
    /// See [`Object::load_image`] for information when an image is [`None`].
    fn load_images(
        document: &Document,
        images: &[gltf::image::Data],
    ) -> HashMap<usize, Option<Arc<dyn Texture>>> {
        document
            .textures()
            .map(|t| {
                let index = t.source().index();
                let data = &images[index];

                (index, Self::load_image(data, &t))
            })
            .collect()
    }

    /// Load an image.
    ///
    /// If the image's [`ColorFormat`] is unsupported, [`None`] is returned.
    fn load_image(data: &gltf::image::Data, texture: &gltf::Texture) -> Option<Arc<dyn Texture>> {
        let format = match ColorFormat::try_from(data.format) {
            Ok(f) => f,
            Err(e) => {
                println!("ERROR: {}", e);

                return None;
            }
        };

        let image = Image::from_raw_with_format(
            data.width,
            data.height,
            &data.pixels,
            format,
            match texture.sampler().mag_filter() {
                None => Interpolation::Bilinear,
                Some(filter) => match filter {
                    MagFilter::Nearest => Interpolation::Closest,
                    MagFilter::Linear => Interpolation::Bilinear,
                },
            },
        );

        Some(Arc::new(image))
    }

    /// Load all materials from a glTF document and store them in a [`HashMap`].
    fn load_materials(
        document: &Document,
        images: &HashMap<usize, Option<Arc<dyn Texture>>>,
    ) -> HashMap<usize, Arc<PbrMaterial>> {
        document
            .materials()
            .filter(|m| m.index().is_some())
            .map(|m| (m.index().unwrap(), Arc::new(Self::load_material(images, m))))
            .collect()
    }

    /// Load a material.
    fn load_material(
        images: &HashMap<usize, Option<Arc<dyn Texture>>>,
        material: gltf::Material,
    ) -> PbrMaterial {
        let pbr = material.pbr_metallic_roughness();

        let base_color_factor: [f32; 3] = pbr.base_color_factor()[0..3].try_into().unwrap();
        let base_color = Color::from(base_color_factor);
        let base_color_texture = pbr
            .base_color_texture()
            .and_then(|i| images.get(&i.texture().index())?.as_ref().map(Arc::clone));

        let metallic = pbr.metallic_factor() as f64;
        let roughness = pbr.roughness_factor() as f64;

        let metallic_roughness_texture = pbr
            .metallic_roughness_texture()
            .and_then(|i| images.get(&i.texture().index())?.as_ref().map(Arc::clone));

        let transmission = material
            .transmission()
            .map(|t| t.transmission_factor())
            .unwrap_or_default() as f64;
        let ior = material.ior().unwrap_or(1.0) as f64;

        let emissive_color_factor = material.emissive_factor();
        let emissive_color = Color::from(emissive_color_factor);
        let emissive_texture = material
            .emissive_texture()
            .and_then(|i| images.get(&i.texture().index())?.as_ref().map(Arc::clone));
        let emissive_strength = material.emissive_strength().unwrap_or(1.0) as f64;

        let normal_texture = material
            .normal_texture()
            .and_then(|i| images.get(&i.texture().index())?.as_ref().map(Arc::clone));

        PbrMaterial {
            base_color,
            base_color_texture,
            metallic,
            roughness,
            metallic_roughness_texture,
            transmission,
            ior,
            emissive_color,
            emissive_texture,
            emissive_strength,
            normal_texture,
        }
    }

    /// Load all meshes from a glTF document.
    ///
    /// - If a mesh does not have indices or vertices, it will be skipped.
    /// - If a mesh does not have normals, flat shading normals will be generated.
    /// - If a mesh does not have tangents, normal mapping won't be applied.
    /// - If a mesh does not have UVs, all vertices will have `(0.0, 0.0)` as their UVs.
    ///
    /// For information about supported materials see [`PbrMaterial`].
    ///
    /// Returns an [`Error`] if the document can not be parsed.
    pub fn load_gltf<P: AsRef<Path>>(path: P) -> Result<Vec<Object>, Error> {
        let (document, buffers, images) = gltf::import(path)?;

        Self::load_gltf_document(&document, &buffers, &images)
    }

    /// Load a glTF document.
    pub(crate) fn load_gltf_document(
        document: &Document,
        buffers: &[gltf::buffer::Data],
        images: &[gltf::image::Data],
    ) -> Result<Vec<Object>, Error> {
        let mut nodes = document.nodes().collect::<Vec<Node>>();
        nodes.sort_by_key(|n| n.index());

        let nodes_indices = nodes.iter().map(|n| n.index()).collect::<BTreeSet<usize>>();

        let children_indices = nodes
            .iter()
            .flat_map(|n| n.children().map(|c| c.index()))
            .collect::<BTreeSet<usize>>();

        let mut transform_matrices = nodes
            .iter()
            .map(|n| Mat4::from(n.transform().matrix().map(|r| r.map(f64::from))).transpose())
            .collect::<Vec<Mat4>>();

        nodes_indices.difference(&children_indices).for_each(|i| {
            Self::apply_transform_recursively(&mut transform_matrices, &nodes[*i]);
        });

        let mut objects = vec![];

        let images = Self::load_images(document, images);
        let materials = Self::load_materials(document, &images);

        for node in nodes {
            if let Some(mesh) = node.mesh() {
                for primitive in mesh.primitives() {
                    let reader = primitive.reader(|buffer| Some(&buffers[buffer.index()]));

                    let matrix = transform_matrices[node.index()];

                    let indices = match Self::load_indices(&reader, &node) {
                        Some(indices) => indices,
                        None => continue,
                    };
                    let vertices = match Self::load_vertices(&reader, matrix, &node) {
                        Some(vertices) => vertices,
                        None => continue,
                    };
                    let normals = Self::load_normals(&reader, matrix, &node);
                    let tangents = Self::load_tangents(&reader, matrix, &node);
                    let uvs = Self::load_uvs(&reader, &node);

                    let triangle_mesh = TriangleMesh::new(
                        indices,
                        vertices,
                        normals.unwrap_or_default(),
                        tangents.unwrap_or_default(),
                        uvs.unwrap_or_default(),
                    );
                    let triangles = triangle_mesh.triangles().collect();

                    let material = match primitive.material().index() {
                        None => Arc::new(PbrMaterial::default()),
                        Some(i) => Arc::clone(&materials[&i]),
                    };

                    objects.push(Object::new(Arc::new(Mesh::new(triangles)), material));
                }
            }
        }

        Ok(objects)
    }
}

impl Shape for Object {
    fn intersects(&self, ray: Ray) -> Option<Hit> {
        self.shape.intersects(ray)
    }

    fn bounds(&self) -> Bounds3 {
        self.shape.bounds()
    }
}