path: root/world/src/mesh.rs

/*
    wearechat - generic multiplayer game with voip
    Copyright (C) 2025 metamuffin

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU Affero General Public License as published by
    the Free Software Foundation, version 3 of the License only.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU Affero General Public License for more details.

    You should have received a copy of the GNU Affero General Public License
    along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/
use crate::{Args, TextureCache, load_texture};
use anyhow::Result;
use gltf::{Mesh, Node, buffer::Data};
use log::{debug, info, warn};
use std::{
    collections::{BTreeMap, BTreeSet},
    path::Path,
};
use weareshared::{
    Affine3A, Vec3A,
    resources::{MeshPart, Prefab},
    store::ResourceStore,
    vec2, vec3a, vec4,
};

pub fn import_mesh(
    mesh: Mesh,
    trans: Affine3A,
    buffers: &[Data],
    store: &ResourceStore,
    path_base: &Path,
    node: &Node,
    prefab: &mut Prefab,
    args: &Args,
    texture_cache: &TextureCache,
    joint_index_map: &BTreeMap<(usize, u16), u32>,
) -> Result<()> {
    for p in mesh.primitives() {
        let name = mesh.name().or(node.name()).map(|e| e.to_owned());
        if let Some(name) = &name {
            info!("adding mesh {name:?}");
        } else {
            info!("adding mesh");
        }
        let reader = p.reader(|buf| Some(&buffers[buf.index()]));

        let mut num_vertex = 0;
        let va_position = reader
            .read_positions()
            .map(|iter| {
                let a = iter.map(|[x, y, z]| vec3a(x, y, z)).collect::<Vec<_>>();
                debug!("{} vertex positions", a.len());
                num_vertex = a.len();
                store.set(&a)
            })
            .transpose()?;

        let va_normal = reader
            .read_normals()
            .map(|iter| {
                let a = iter.map(|[x, y, z]| vec3a(x, y, z)).collect::<Vec<_>>();
                debug!("{} vertex normals", a.len());
                store.set(&a)
            })
            .transpose()?;

        let va_tangent = reader
            .read_tangents()
            .map(|iter| {
                // TODO dont ignore handedness
                let a = iter
                    .map(|[x, y, z, h]| vec4(x, y, z, h))
                    .collect::<Vec<_>>();
                debug!("{} vertex tangents", a.len());
                store.set(&a)
            })
            .transpose()?;

        let va_joint_index = reader
            .read_joints(0)
            .map(|iter| {
                let si = node.skin().unwrap().index();
                let a = iter
                    .into_u16()
                    .map(|x| x.map(|x| joint_index_map[&(si, x)]))
                    .collect::<Vec<_>>();
                debug!("{} vertex joint indecies", a.len());
                eprintln!(
                    "index {:?} {:?}",
                    node.name(),
                    a.iter().flatten().collect::<BTreeSet<_>>()
                );
                if a.len() != num_vertex {
                    warn!("joint index count does not vertex count")
                }
                store.set(&a)
            })
            .transpose()?;

        let va_joint_weight = reader
            .read_weights(0)
            .map(|iter| {
                let a = iter.into_f32().collect::<Vec<_>>();
                debug!("{} vertex joint weights", a.len());
                if a.len() != num_vertex {
                    warn!("joint weight count does not vertex count")
                }
                store.set(&a)
            })
            .transpose()?;

        let va_texcoord = reader
            .read_tex_coords(0)
            .map(|iter| {
                let a = iter.into_f32().map(|[x, y]| vec2(x, y)).collect::<Vec<_>>();
                debug!("{} vertex texture coordinates", a.len());
                store.set(&a)
            })
            .transpose()?;

        let va_albedo = reader
            .read_colors(0)
            .map(|iter| {
                let a = iter
                    .into_rgb_f32()
                    .map(|[x, y, z]| vec3a(x, y, z))
                    .collect::<Vec<_>>();
                debug!("{} vertex colors", a.len());
                store.set(&a)
            })
            .transpose()?;

        let va_alpha = reader
            .read_colors(0)
            .map(|iter| {
                let mut color_a = vec![];
                for p in iter.into_rgba_f32() {
                    color_a.push(p[3]);
                }
                let o = if color_a.iter().any(|x| *x != 1.) {
                    debug!("{} vertex transmissions", color_a.len());
                    Some(store.set(&color_a)?)
                } else {
                    debug!("vertex transmission pruned");
                    None
                };
                Ok::<_, anyhow::Error>(o)
            })
            .transpose()?
            .flatten();

        let index = reader
            .read_indices()
            .unwrap()
            .into_u32()
            .array_chunks::<3>()
            .collect::<Vec<_>>();

        debug!("{} indecies", index.len() * 3);
        let index = Some(store.set(&index)?);

        let mut tex_albedo = None;
        let mut tex_alpha = None;
        if let Some(tex) = p.material().pbr_metallic_roughness().base_color_texture() {
            let r = load_texture(
                "albedo",
                store,
                path_base,
                buffers,
                &tex.texture().source().source(),
                args.webp,
                texture_cache,
            )?;
            tex_albedo = Some(r.clone());
            tex_alpha = Some(r.clone());
        }
        let mut tex_normal = None;
        if let Some(tex) = p.material().normal_texture() {
            tex_normal = Some(load_texture(
                "normal",
                store,
                path_base,
                buffers,
                &tex.texture().source().source(),
                args.webp,
                texture_cache,
            )?);
        }
        let mut tex_emission = None;
        if let Some(tex) = p.material().emissive_texture() {
            tex_emission = Some(load_texture(
                "emission",
                store,
                path_base,
                buffers,
                &tex.texture().source().source(),
                args.webp,
                texture_cache,
            )?);
        }
        let mut tex_transmission = None;
        if let Some(tex) = p
            .material()
            .transmission()
            .and_then(|t| t.transmission_texture())
        {
            tex_transmission = Some(load_texture(
                "transmission",
                store,
                path_base,
                buffers,
                &tex.texture().source().source(),
                args.webp,
                texture_cache,
            )?);
        }
        let mut tex_thickness = None;
        if let Some(tex) = p.material().volume().and_then(|t| t.thickness_texture()) {
            tex_thickness = Some(load_texture(
                "thickness",
                store,
                path_base,
                buffers,
                &tex.texture().source().source(),
                args.webp,
                texture_cache,
            )?);
        }
        let mut tex_occlusion = None;
        if let Some(tex) = p.material().occlusion_texture() {
            tex_occlusion = Some(load_texture(
                "occlusion",
                store,
                path_base,
                buffers,
                &tex.texture().source().source(),
                args.webp,
                texture_cache,
            )?);
        }
        let mut tex_roughness = None;
        let mut tex_metallic = None;
        if let Some(tex) = p
            .material()
            .pbr_metallic_roughness()
            .metallic_roughness_texture()
        {
            let r = load_texture(
                "metallic+roughness",
                store,
                path_base,
                buffers,
                &tex.texture().source().source(),
                args.webp,
                texture_cache,
            )?;
            tex_roughness = Some(r.clone());
            tex_metallic = Some(r.clone());
        }

        let g_metallic = Some(p.material().pbr_metallic_roughness().metallic_factor());
        let g_roughness = Some(p.material().pbr_metallic_roughness().roughness_factor());

        let base_color = p.material().pbr_metallic_roughness().base_color_factor();

        let g_albedo = if base_color[0] != 1. || base_color[1] != 1. || base_color[2] != 1. {
            debug!(
                "albedo is r={},g={},b={}",
                base_color[0], base_color[1], base_color[2]
            );
            Some(Vec3A::new(base_color[0], base_color[1], base_color[2]))
        } else {
            debug!("albedo pruned");
            None
        };
        let g_alpha = if base_color[3] != 1. {
            debug!("alpha is {}", base_color[3]);
            Some(base_color[3])
        } else {
            debug!("alpha pruned");
            None
        };

        let emission = p.material().emissive_factor();
        let g_emission = if emission[0] != 0. || emission[1] != 0. || emission[2] != 0. {
            debug!(
                "emission is r={},g={},b={}",
                base_color[0], base_color[1], base_color[2]
            );
            Some(Vec3A::new(emission[0], emission[1], emission[2]))
        } else {
            debug!("emission pruned");
            None
        };

        let transmission = p
            .material()
            .transmission()
            .map(|t| t.transmission_factor())
            .unwrap_or(0.);

        let g_transmission = if transmission != 0. {
            debug!("transmission is {transmission}");
            Some(transmission)
        } else {
            debug!("transmission pruned");
            None
        };

        let g_dispersion = p
            .material()
            .extension_value("KHR_materials_dispersion")
            .and_then(|e| e.get("dispersion"))
            .and_then(|e| e.as_f64())
            .map(|e| e as f32);
        if let Some(d) = g_dispersion {
            debug!("dispersion is {d}");
        }

        let g_attenuation = p.material().volume().map(|v| {
            let ref_dist = v.attenuation_distance();
            let att = Vec3A::from_array(v.attenuation_color().map(
                // manually derived from attenuation coefficient formula. i hope this is correct.
                |factor| -(factor.powf(1. / ref_dist)).ln(),
            ));
            debug!("attenuation is {att}");
            att
        });
        let g_refractive_index = p.material().ior();
        if let Some(i) = g_refractive_index {
            debug!("refractive index is {i}");
        }
        let g_thickness = p.material().volume().map(|v| v.thickness_factor());

        let g_unlit = if p
            .material()
            .extension_value("KHR_materials_unlit")
            .is_some()
        {
            debug!("unlit");
            Some(())
        } else {
            None
        };

        let g_double_sided = if p.material().double_sided() {
            debug!("double sided");
            Some(())
        } else {
            None
        };

        let armature = node.skin().map(|s| 0);

        let mesh = store.set(&MeshPart {
            name,
            index,
            armature,
            g_albedo,
            g_alpha,
            g_metallic,
            g_roughness,
            g_emission,
            g_transmission,
            g_attenuation,
            g_thickness,
            g_refractive_index,
            g_dispersion,
            g_unlit,
            g_double_sided,
            va_position,
            va_normal,
            va_tangent,
            va_texcoord,
            va_albedo,
            va_alpha,
            va_joint_index,
            va_joint_weight,
            tex_albedo,
            tex_normal,
            tex_roughness,
            tex_metallic,
            tex_alpha,
            tex_emission,
            tex_transmission,
            tex_thickness,
            tex_occlusion,
            // not supported by gltf
            hint_mirror: None, // TODO
            hint_static: None, // TODO Set when instancing
            va_transmission: None,
            va_emission: None,
            va_metallic: None,
            va_roughness: None,
        })?;

        prefab.mesh.push((trans, mesh))
    }
    Ok(())
}