roast2d_internal 0.4.0

//! Core 3D mesh type with primitive generators and loaders.

use std::path::Path;

use wgpu::util::DeviceExt;

use crate::engine::Engine;

use super::vertex::Mesh3DVertex;

#[derive(Clone)]
pub struct Mesh3D {
    pub vertex: wgpu::Buffer,
    pub index: wgpu::Buffer,
    pub index_count: u32,
}

impl Mesh3D {
    pub fn new(g: &Engine, vertices: &[Mesh3DVertex], indices: &[u16]) -> Self {
        let state = g.backend_state();
        let vertex = state
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("mesh3d.vertices"),
                contents: bytemuck::cast_slice(vertices),
                usage: wgpu::BufferUsages::VERTEX,
            });
        let index = state
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("mesh3d.indices"),
                contents: bytemuck::cast_slice(indices),
                usage: wgpu::BufferUsages::INDEX,
            });
        Self {
            vertex,
            index,
            index_count: indices.len() as u32,
        }
    }

    /// Generate a cube mesh centered at origin.
    ///
    /// # Arguments
    /// * `g` - Engine reference
    /// * `size` - Size of the cube (edge length)
    pub fn cube(g: &Engine, size: f32) -> Self {
        Self::cuboid(g, size, size, size)
    }

    /// Generate a cuboid (box) mesh centered at origin.
    ///
    /// # Arguments
    /// * `g` - Engine reference
    /// * `width` - Size along X axis
    /// * `height` - Size along Y axis
    /// * `depth` - Size along Z axis
    pub fn cuboid(g: &Engine, width: f32, height: f32, depth: f32) -> Self {
        let hw = width / 2.0;
        let hh = height / 2.0;
        let hd = depth / 2.0;

        let mut vertices = Vec::with_capacity(24);
        let mut indices = Vec::with_capacity(36);

        // Helper to add a face
        let mut add_face = |positions: [[f32; 3]; 4], normal: [f32; 3]| {
            let base = vertices.len() as u16;
            let n = glam::Vec3::from(normal);
            for (i, pos) in positions.iter().enumerate() {
                let uv = match i {
                    0 => glam::Vec2::new(0.0, 0.0),
                    1 => glam::Vec2::new(1.0, 0.0),
                    2 => glam::Vec2::new(1.0, 1.0),
                    _ => glam::Vec2::new(0.0, 1.0),
                };
                vertices.push(Mesh3DVertex::new(glam::Vec3::from(*pos), n, uv));
            }
            indices.extend_from_slice(&[base, base + 1, base + 2, base + 2, base + 3, base]);
        };

        // +Z face (front)
        add_face(
            [[-hw, -hh, hd], [hw, -hh, hd], [hw, hh, hd], [-hw, hh, hd]],
            [0.0, 0.0, 1.0],
        );
        // -Z face (back)
        add_face(
            [
                [hw, -hh, -hd],
                [-hw, -hh, -hd],
                [-hw, hh, -hd],
                [hw, hh, -hd],
            ],
            [0.0, 0.0, -1.0],
        );
        // +X face (right)
        add_face(
            [[hw, -hh, hd], [hw, -hh, -hd], [hw, hh, -hd], [hw, hh, hd]],
            [1.0, 0.0, 0.0],
        );
        // -X face (left)
        add_face(
            [
                [-hw, -hh, -hd],
                [-hw, -hh, hd],
                [-hw, hh, hd],
                [-hw, hh, -hd],
            ],
            [-1.0, 0.0, 0.0],
        );
        // +Y face (top)
        add_face(
            [[-hw, hh, hd], [hw, hh, hd], [hw, hh, -hd], [-hw, hh, -hd]],
            [0.0, 1.0, 0.0],
        );
        // -Y face (bottom)
        add_face(
            [
                [-hw, -hh, -hd],
                [hw, -hh, -hd],
                [hw, -hh, hd],
                [-hw, -hh, hd],
            ],
            [0.0, -1.0, 0.0],
        );

        Self::new(g, &vertices, &indices)
    }

    /// Generate a plane mesh on the XZ plane with normal pointing up (+Y).
    ///
    /// # Arguments
    /// * `g` - Engine reference
    /// * `width` - Size along X axis
    /// * `depth` - Size along Z axis
    pub fn plane(g: &Engine, width: f32, depth: f32) -> Self {
        let hw = width / 2.0;
        let hd = depth / 2.0;
        let normal = glam::Vec3::Y;

        let vertices = vec![
            Mesh3DVertex::new(
                glam::Vec3::new(-hw, 0.0, -hd),
                normal,
                glam::Vec2::new(0.0, 0.0),
            ),
            Mesh3DVertex::new(
                glam::Vec3::new(hw, 0.0, -hd),
                normal,
                glam::Vec2::new(1.0, 0.0),
            ),
            Mesh3DVertex::new(
                glam::Vec3::new(hw, 0.0, hd),
                normal,
                glam::Vec2::new(1.0, 1.0),
            ),
            Mesh3DVertex::new(
                glam::Vec3::new(-hw, 0.0, hd),
                normal,
                glam::Vec2::new(0.0, 1.0),
            ),
        ];
        // CCW winding order for +Y facing normal (viewed from above)
        let indices = vec![0, 2, 1, 0, 3, 2];

        Self::new(g, &vertices, &indices)
    }

    /// Generate a UV sphere mesh centered at origin.
    ///
    /// # Arguments
    /// * `g` - Engine reference
    /// * `radius` - Sphere radius
    /// * `rings` - Number of horizontal rings (latitude divisions)
    /// * `slices` - Number of vertical slices (longitude divisions)
    pub fn sphere(g: &Engine, radius: f32, rings: u32, slices: u32) -> Self {
        let rings = rings.max(2);
        let slices = slices.max(3);

        let mut vertices = Vec::new();
        let mut indices = Vec::new();

        // Generate vertices
        for ring in 0..=rings {
            let theta = std::f32::consts::PI * ring as f32 / rings as f32;
            let sin_theta = theta.sin();
            let cos_theta = theta.cos();

            for slice in 0..=slices {
                let phi = 2.0 * std::f32::consts::PI * slice as f32 / slices as f32;
                let sin_phi = phi.sin();
                let cos_phi = phi.cos();

                let x = sin_theta * cos_phi;
                let y = cos_theta;
                let z = sin_theta * sin_phi;

                let position = glam::Vec3::new(x * radius, y * radius, z * radius);
                let normal = glam::Vec3::new(x, y, z);
                let uv = glam::Vec2::new(slice as f32 / slices as f32, ring as f32 / rings as f32);

                vertices.push(Mesh3DVertex::new(position, normal, uv));
            }
        }

        // Generate indices
        for ring in 0..rings {
            for slice in 0..slices {
                let current = ring * (slices + 1) + slice;
                let next = current + slices + 1;

                // Two triangles per quad
                indices.push(current as u16);
                indices.push(next as u16);
                indices.push((current + 1) as u16);

                indices.push((current + 1) as u16);
                indices.push(next as u16);
                indices.push((next + 1) as u16);
            }
        }

        Self::new(g, &vertices, &indices)
    }

    /// Generate a cylinder mesh centered at origin, extending along Y axis.
    ///
    /// # Arguments
    /// * `g` - Engine reference
    /// * `radius` - Cylinder radius
    /// * `height` - Cylinder height
    /// * `slices` - Number of slices around the circumference
    pub fn cylinder(g: &Engine, radius: f32, height: f32, slices: u32) -> Self {
        let slices = slices.max(3);
        let half_height = height / 2.0;

        let mut vertices = Vec::new();
        let mut indices = Vec::new();

        // Generate side vertices
        for i in 0..=slices {
            let angle = 2.0 * std::f32::consts::PI * i as f32 / slices as f32;
            let x = angle.cos();
            let z = angle.sin();
            let u = i as f32 / slices as f32;

            // Bottom vertex
            vertices.push(Mesh3DVertex::new(
                glam::Vec3::new(x * radius, -half_height, z * radius),
                glam::Vec3::new(x, 0.0, z),
                glam::Vec2::new(u, 0.0),
            ));
            // Top vertex
            vertices.push(Mesh3DVertex::new(
                glam::Vec3::new(x * radius, half_height, z * radius),
                glam::Vec3::new(x, 0.0, z),
                glam::Vec2::new(u, 1.0),
            ));
        }

        // Side indices
        for i in 0..slices {
            let base = i * 2;
            indices.push(base as u16);
            indices.push((base + 2) as u16);
            indices.push((base + 1) as u16);
            indices.push((base + 1) as u16);
            indices.push((base + 2) as u16);
            indices.push((base + 3) as u16);
        }

        // Top cap
        let top_center_idx = vertices.len() as u16;
        vertices.push(Mesh3DVertex::new(
            glam::Vec3::new(0.0, half_height, 0.0),
            glam::Vec3::Y,
            glam::Vec2::new(0.5, 0.5),
        ));
        for i in 0..=slices {
            let angle = 2.0 * std::f32::consts::PI * i as f32 / slices as f32;
            let x = angle.cos();
            let z = angle.sin();
            vertices.push(Mesh3DVertex::new(
                glam::Vec3::new(x * radius, half_height, z * radius),
                glam::Vec3::Y,
                glam::Vec2::new(x * 0.5 + 0.5, z * 0.5 + 0.5),
            ));
        }
        for i in 0..slices {
            indices.push(top_center_idx);
            indices.push(top_center_idx + 1 + i as u16);
            indices.push(top_center_idx + 2 + i as u16);
        }

        // Bottom cap
        let bottom_center_idx = vertices.len() as u16;
        vertices.push(Mesh3DVertex::new(
            glam::Vec3::new(0.0, -half_height, 0.0),
            -glam::Vec3::Y,
            glam::Vec2::new(0.5, 0.5),
        ));
        for i in 0..=slices {
            let angle = 2.0 * std::f32::consts::PI * i as f32 / slices as f32;
            let x = angle.cos();
            let z = angle.sin();
            vertices.push(Mesh3DVertex::new(
                glam::Vec3::new(x * radius, -half_height, z * radius),
                -glam::Vec3::Y,
                glam::Vec2::new(x * 0.5 + 0.5, z * 0.5 + 0.5),
            ));
        }
        for i in 0..slices {
            indices.push(bottom_center_idx);
            indices.push(bottom_center_idx + 2 + i as u16);
            indices.push(bottom_center_idx + 1 + i as u16);
        }

        Self::new(g, &vertices, &indices)
    }

    /// Generate a cone mesh with base centered at origin, pointing up along Y axis.
    ///
    /// # Arguments
    /// * `g` - Engine reference
    /// * `radius` - Base radius
    /// * `height` - Cone height
    /// * `slices` - Number of slices around the circumference
    pub fn cone(g: &Engine, radius: f32, height: f32, slices: u32) -> Self {
        let slices = slices.max(3);

        let mut vertices = Vec::new();
        let mut indices = Vec::new();

        // The normal for cone sides needs to account for the slope
        let slope = radius / height;
        let normal_y = slope / (1.0 + slope * slope).sqrt();
        let normal_xz = 1.0 / (1.0 + slope * slope).sqrt();

        // Apex vertex
        let apex_idx = 0u16;
        vertices.push(Mesh3DVertex::new(
            glam::Vec3::new(0.0, height, 0.0),
            glam::Vec3::Y,
            glam::Vec2::new(0.5, 1.0),
        ));

        // Side vertices (base ring)
        for i in 0..=slices {
            let angle = 2.0 * std::f32::consts::PI * i as f32 / slices as f32;
            let x = angle.cos();
            let z = angle.sin();

            // Normal pointing outward and slightly up
            let normal = glam::Vec3::new(x * normal_xz, normal_y, z * normal_xz).normalize();

            vertices.push(Mesh3DVertex::new(
                glam::Vec3::new(x * radius, 0.0, z * radius),
                normal,
                glam::Vec2::new(i as f32 / slices as f32, 0.0),
            ));
        }

        // Side indices
        for i in 0..slices {
            indices.push(apex_idx);
            indices.push(1 + i as u16);
            indices.push(2 + i as u16);
        }

        // Base cap
        let base_center_idx = vertices.len() as u16;
        vertices.push(Mesh3DVertex::new(
            glam::Vec3::new(0.0, 0.0, 0.0),
            -glam::Vec3::Y,
            glam::Vec2::new(0.5, 0.5),
        ));
        for i in 0..=slices {
            let angle = 2.0 * std::f32::consts::PI * i as f32 / slices as f32;
            let x = angle.cos();
            let z = angle.sin();
            vertices.push(Mesh3DVertex::new(
                glam::Vec3::new(x * radius, 0.0, z * radius),
                -glam::Vec3::Y,
                glam::Vec2::new(x * 0.5 + 0.5, z * 0.5 + 0.5),
            ));
        }
        for i in 0..slices {
            indices.push(base_center_idx);
            indices.push(base_center_idx + 2 + i as u16);
            indices.push(base_center_idx + 1 + i as u16);
        }

        Self::new(g, &vertices, &indices)
    }

    /// Generate a torus (donut) mesh centered at origin on the XZ plane.
    ///
    /// # Arguments
    /// * `g` - Engine reference
    /// * `radius` - Distance from center to tube center
    /// * `tube_radius` - Radius of the tube
    /// * `rings` - Number of rings around the torus
    /// * `slices` - Number of slices per ring
    pub fn torus(g: &Engine, radius: f32, tube_radius: f32, rings: u32, slices: u32) -> Self {
        let rings = rings.max(3);
        let slices = slices.max(3);

        let mut vertices = Vec::new();
        let mut indices = Vec::new();

        for ring in 0..=rings {
            let theta = 2.0 * std::f32::consts::PI * ring as f32 / rings as f32;
            let cos_theta = theta.cos();
            let sin_theta = theta.sin();

            for slice in 0..=slices {
                let phi = 2.0 * std::f32::consts::PI * slice as f32 / slices as f32;
                let cos_phi = phi.cos();
                let sin_phi = phi.sin();

                // Position on the torus
                let x = (radius + tube_radius * cos_phi) * cos_theta;
                let y = tube_radius * sin_phi;
                let z = (radius + tube_radius * cos_phi) * sin_theta;

                // Normal points from ring center to vertex
                let nx = cos_phi * cos_theta;
                let ny = sin_phi;
                let nz = cos_phi * sin_theta;

                let position = glam::Vec3::new(x, y, z);
                let normal = glam::Vec3::new(nx, ny, nz);
                let uv = glam::Vec2::new(ring as f32 / rings as f32, slice as f32 / slices as f32);

                vertices.push(Mesh3DVertex::new(position, normal, uv));
            }
        }

        // Generate indices
        for ring in 0..rings {
            for slice in 0..slices {
                let current = ring * (slices + 1) + slice;
                let next = current + slices + 1;

                indices.push(current as u16);
                indices.push(next as u16);
                indices.push((current + 1) as u16);

                indices.push((current + 1) as u16);
                indices.push(next as u16);
                indices.push((next + 1) as u16);
            }
        }

        Self::new(g, &vertices, &indices)
    }

    /// Load a 3D mesh from an OBJ file.
    ///
    /// Returns a vector of meshes, one for each object/group in the OBJ file.
    /// If the OBJ file contains multiple objects, use `from_obj_single` to get only the first one.
    ///
    /// # Arguments
    /// * `g` - Engine reference
    /// * `path` - Path to the OBJ file
    ///
    /// # Example
    /// ```ignore
    /// let meshes = Mesh3D::from_obj(g, "assets/models/teapot.obj")?;
    /// ```
    pub fn from_obj<P: AsRef<Path>>(g: &Engine, path: P) -> anyhow::Result<Vec<Self>> {
        let (models, _materials) = tobj::load_obj(
            path.as_ref(),
            &tobj::LoadOptions {
                triangulate: true,
                single_index: true,
                ..Default::default()
            },
        )?;

        let mut meshes = Vec::with_capacity(models.len());

        for model in models {
            let mesh = &model.mesh;

            // Build vertices from OBJ data
            let vertex_count = mesh.positions.len() / 3;
            let mut vertices = Vec::with_capacity(vertex_count);

            let has_normals = !mesh.normals.is_empty();
            let has_texcoords = !mesh.texcoords.is_empty();

            for i in 0..vertex_count {
                let position = glam::Vec3::new(
                    mesh.positions[i * 3],
                    mesh.positions[i * 3 + 1],
                    mesh.positions[i * 3 + 2],
                );

                let normal = if has_normals {
                    glam::Vec3::new(
                        mesh.normals[i * 3],
                        mesh.normals[i * 3 + 1],
                        mesh.normals[i * 3 + 2],
                    )
                } else {
                    // Default normal pointing up if not provided
                    glam::Vec3::Y
                };

                let uv = if has_texcoords {
                    glam::Vec2::new(mesh.texcoords[i * 2], mesh.texcoords[i * 2 + 1])
                } else {
                    glam::Vec2::ZERO
                };

                vertices.push(Mesh3DVertex::new(position, normal, uv));
            }

            // Convert indices to u16
            let indices: Vec<u16> = mesh.indices.iter().map(|&i| i as u16).collect();

            meshes.push(Self::new(g, &vertices, &indices));
        }

        Ok(meshes)
    }

    /// Load a single 3D mesh from an OBJ file.
    ///
    /// Returns only the first mesh in the OBJ file.
    /// Use `from_obj` if the file contains multiple objects and you need all of them.
    ///
    /// # Arguments
    /// * `g` - Engine reference
    /// * `path` - Path to the OBJ file
    pub fn from_obj_single<P: AsRef<Path>>(g: &Engine, path: P) -> anyhow::Result<Self> {
        let meshes = Self::from_obj(g, path)?;
        meshes
            .into_iter()
            .next()
            .ok_or_else(|| anyhow::anyhow!("OBJ file contains no meshes"))
    }

    /// Load a 3D mesh from OBJ data in memory.
    ///
    /// # Arguments
    /// * `g` - Engine reference
    /// * `obj_data` - OBJ file contents as bytes
    /// * `name` - Optional name for error messages
    pub fn from_obj_bytes(g: &Engine, obj_data: &[u8], name: &str) -> anyhow::Result<Vec<Self>> {
        use std::io::BufReader;

        let (models, _materials) = tobj::load_obj_buf(
            &mut BufReader::new(obj_data),
            &tobj::LoadOptions {
                triangulate: true,
                single_index: true,
                ..Default::default()
            },
            |_mtl_path| {
                // We don't load materials from memory
                Err(tobj::LoadError::OpenFileFailed)
            },
        )?;

        let mut meshes = Vec::with_capacity(models.len());

        for model in models {
            let mesh = &model.mesh;

            let vertex_count = mesh.positions.len() / 3;
            let mut vertices = Vec::with_capacity(vertex_count);

            let has_normals = !mesh.normals.is_empty();
            let has_texcoords = !mesh.texcoords.is_empty();

            for i in 0..vertex_count {
                let position = glam::Vec3::new(
                    mesh.positions[i * 3],
                    mesh.positions[i * 3 + 1],
                    mesh.positions[i * 3 + 2],
                );

                let normal = if has_normals {
                    glam::Vec3::new(
                        mesh.normals[i * 3],
                        mesh.normals[i * 3 + 1],
                        mesh.normals[i * 3 + 2],
                    )
                } else {
                    glam::Vec3::Y
                };

                let uv = if has_texcoords {
                    glam::Vec2::new(mesh.texcoords[i * 2], mesh.texcoords[i * 2 + 1])
                } else {
                    glam::Vec2::ZERO
                };

                vertices.push(Mesh3DVertex::new(position, normal, uv));
            }

            let indices: Vec<u16> = mesh.indices.iter().map(|&i| i as u16).collect();
            meshes.push(Self::new(g, &vertices, &indices));
        }

        if meshes.is_empty() {
            anyhow::bail!("OBJ data '{}' contains no meshes", name);
        }

        Ok(meshes)
    }

    /// Load 3D meshes from a glTF file.
    ///
    /// Returns all meshes found in the glTF file.
    /// glTF is the modern standard format for 3D models.
    ///
    /// # Arguments
    /// * `g` - Engine reference
    /// * `path` - Path to the glTF or GLB file
    ///
    /// # Example
    /// ```ignore
    /// let meshes = Mesh3D::from_gltf(g, "assets/models/character.gltf")?;
    /// // or for binary glTF
    /// let meshes = Mesh3D::from_gltf(g, "assets/models/character.glb")?;
    /// ```
    pub fn from_gltf<P: AsRef<Path>>(g: &Engine, path: P) -> anyhow::Result<Vec<Self>> {
        let (document, buffers, _images) = gltf::import(path.as_ref())?;
        Self::from_gltf_document(g, &document, &buffers)
    }

    /// Load a single mesh from a glTF file.
    ///
    /// Convenience method when you know the file contains exactly one mesh.
    pub fn from_gltf_single<P: AsRef<Path>>(g: &Engine, path: P) -> anyhow::Result<Self> {
        let meshes = Self::from_gltf(g, path)?;
        meshes
            .into_iter()
            .next()
            .ok_or_else(|| anyhow::anyhow!("glTF file contains no meshes"))
    }

    /// Load 3D meshes from glTF/GLB data in memory.
    ///
    /// # Arguments
    /// * `g` - Engine reference
    /// * `data` - glTF JSON or GLB binary data
    pub fn from_gltf_bytes(g: &Engine, data: &[u8]) -> anyhow::Result<Vec<Self>> {
        let (document, buffers, _images) = gltf::import_slice(data)?;
        Self::from_gltf_document(g, &document, &buffers)
    }

    /// Load a single mesh from glTF/GLB data in memory.
    pub fn from_gltf_bytes_single(g: &Engine, data: &[u8]) -> anyhow::Result<Self> {
        let meshes = Self::from_gltf_bytes(g, data)?;
        meshes
            .into_iter()
            .next()
            .ok_or_else(|| anyhow::anyhow!("glTF data contains no meshes"))
    }

    /// Internal helper to convert glTF document to meshes
    fn from_gltf_document(
        g: &Engine,
        document: &gltf::Document,
        buffers: &[gltf::buffer::Data],
    ) -> anyhow::Result<Vec<Self>> {
        let mut meshes = Vec::new();

        for mesh in document.meshes() {
            for primitive in mesh.primitives() {
                // Only support triangle primitives
                if primitive.mode() != gltf::mesh::Mode::Triangles {
                    continue;
                }

                let reader = primitive.reader(|buffer| Some(&buffers[buffer.index()]));

                // Get positions (required)
                let positions: Vec<[f32; 3]> = reader
                    .read_positions()
                    .ok_or_else(|| anyhow::anyhow!("glTF primitive has no positions"))?
                    .collect();

                // Get normals (optional, generate default if missing)
                let normals: Vec<[f32; 3]> = reader
                    .read_normals()
                    .map(|iter| iter.collect())
                    .unwrap_or_else(|| vec![[0.0, 1.0, 0.0]; positions.len()]);

                // Get UVs (optional, use zero if missing)
                let uvs: Vec<[f32; 2]> = reader
                    .read_tex_coords(0)
                    .map(|iter| iter.into_f32().collect())
                    .unwrap_or_else(|| vec![[0.0, 0.0]; positions.len()]);

                // Get indices (optional, generate sequential if missing)
                let indices: Vec<u16> = if let Some(indices_reader) = reader.read_indices() {
                    indices_reader.into_u32().map(|i| i as u16).collect()
                } else {
                    (0..positions.len() as u16).collect()
                };

                // Build vertices
                let vertices: Vec<Mesh3DVertex> = positions
                    .iter()
                    .zip(normals.iter())
                    .zip(uvs.iter())
                    .map(|((pos, normal), uv)| {
                        Mesh3DVertex::new(
                            glam::Vec3::from_array(*pos),
                            glam::Vec3::from_array(*normal),
                            glam::Vec2::from_array(*uv),
                        )
                    })
                    .collect();

                meshes.push(Self::new(g, &vertices, &indices));
            }
        }

        if meshes.is_empty() {
            anyhow::bail!("glTF document contains no triangle meshes");
        }

        Ok(meshes)
    }
}