polyscope-render 0.5.10

//! Surface mesh GPU rendering resources.

use glam::{Vec3, Vec4};
use wgpu::util::DeviceExt;

/// Uniforms for surface mesh rendering.
/// Note: Layout must match WGSL `MeshUniforms` exactly.
/// WGSL `vec3<f32>` has 16-byte alignment, requiring extra padding.
#[repr(C)]
#[derive(Debug, Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
#[allow(clippy::pub_underscore_fields)]
pub struct MeshUniforms {
    /// Model transform matrix (must be first for alignment)
    pub model_matrix: [[f32; 4]; 4],
    /// Shading style: 0 = smooth, 1 = flat, 2 = tri-flat
    pub shade_style: u32,
    /// Show edges: 0 = off, 1 = on
    pub show_edges: u32,
    /// Edge width in pixels
    pub edge_width: f32,
    /// Surface transparency (0.0 = opaque, 1.0 = fully transparent)
    pub transparency: f32,
    /// Surface color (RGBA)
    pub surface_color: [f32; 4],
    /// Edge color (RGBA)
    pub edge_color: [f32; 4],
    /// Backface policy: 0 = identical, 1 = different, 2 = custom, 3 = cull
    pub backface_policy: u32,
    /// Slice plane clipping enable: 0 = off, 1 = on
    pub slice_planes_enabled: u32,
    /// Use per-vertex colors (1) or surface color (0)
    pub use_vertex_color: u32,
    /// Padding to align to 16 bytes
    pub _pad1: f32,
    /// Padding matching WGSL layout (12 bytes)
    pub _pad2: [f32; 3],
    /// Padding to align vec4 to 16 bytes
    pub _pad3: f32,
    /// Backface color (RGBA), used when `backface_policy` is custom
    pub backface_color: [f32; 4],
}

/// Model uniforms for shadow rendering.
/// Only contains the model matrix (64 bytes) - matches `shadow_map.wgsl` `ModelUniforms`.
#[repr(C)]
#[derive(Debug, Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
pub struct ShadowModelUniforms {
    /// Model transform matrix.
    pub model: [[f32; 4]; 4],
}

impl Default for ShadowModelUniforms {
    fn default() -> Self {
        Self {
            model: [
                [1.0, 0.0, 0.0, 0.0],
                [0.0, 1.0, 0.0, 0.0],
                [0.0, 0.0, 1.0, 0.0],
                [0.0, 0.0, 0.0, 1.0],
            ],
        }
    }
}

impl Default for MeshUniforms {
    fn default() -> Self {
        Self {
            model_matrix: [
                [1.0, 0.0, 0.0, 0.0],
                [0.0, 1.0, 0.0, 0.0],
                [0.0, 0.0, 1.0, 0.0],
                [0.0, 0.0, 0.0, 1.0],
            ],
            shade_style: 0,                      // smooth shading
            show_edges: 0,                       // edges off
            edge_width: 1.0,                     // 1 pixel edge
            transparency: 0.0,                   // fully opaque
            surface_color: [0.5, 0.5, 0.5, 1.0], // gray
            edge_color: [0.0, 0.0, 0.0, 1.0],    // black edges
            backface_policy: 0,                  // identical to front
            slice_planes_enabled: 1,
            use_vertex_color: 0,
            _pad1: 0.0,
            _pad2: [0.0; 3],
            _pad3: 0.0,
            backface_color: [0.3, 0.3, 0.3, 1.0], // darker gray
        }
    }
}

/// GPU resources for rendering a surface mesh.
pub struct SurfaceMeshRenderData {
    /// Position buffer (storage buffer, vec4 for alignment).
    pub vertex_buffer: wgpu::Buffer,
    /// Index buffer (triangle indices).
    pub index_buffer: wgpu::Buffer,
    /// Normal buffer (vertex normals, vec4 for alignment).
    pub normal_buffer: wgpu::Buffer,
    /// Barycentric coordinate buffer for wireframe rendering.
    /// Each triangle vertex gets `[1,0,0]`, `[0,1,0]`, `[0,0,1]`.
    pub barycentric_buffer: wgpu::Buffer,
    /// Color buffer (per-vertex colors for quantities, vec4).
    pub color_buffer: wgpu::Buffer,
    /// Edge is real buffer - marks which edges are real polygon edges vs triangulation internal.
    pub edge_is_real_buffer: wgpu::Buffer,
    /// Uniform buffer for mesh-specific settings.
    pub uniform_buffer: wgpu::Buffer,
    /// Bind group for this surface mesh.
    pub bind_group: wgpu::BindGroup,
    /// Number of triangles.
    pub num_triangles: u32,
    /// Number of indices (`num_triangles` * 3).
    pub num_indices: u32,
    /// Shadow pass bind group (for rendering to shadow map).
    pub shadow_bind_group: Option<wgpu::BindGroup>,
    /// Shadow model uniform buffer (contains only model matrix).
    pub shadow_model_buffer: Option<wgpu::Buffer>,
}

impl SurfaceMeshRenderData {
    /// Creates new render data from mesh geometry.
    ///
    /// # Arguments
    /// * `device` - The wgpu device
    /// * `bind_group_layout` - The bind group layout for surface meshes
    /// * `camera_buffer` - The camera uniform buffer
    /// * `vertices` - Vertex positions
    /// * `triangles` - Triangle indices (each [u32; 3] is one triangle)
    /// * `vertex_normals` - Per-vertex normals
    /// * `edge_is_real` - Per-triangle-vertex flags marking real polygon edges vs internal triangulation edges
    #[must_use]
    pub fn new(
        device: &wgpu::Device,
        bind_group_layout: &wgpu::BindGroupLayout,
        camera_buffer: &wgpu::Buffer,
        vertices: &[Vec3],
        triangles: &[[u32; 3]],
        vertex_normals: &[Vec3],
        edge_is_real: &[Vec3],
    ) -> Self {
        let num_triangles = triangles.len() as u32;
        let num_indices = num_triangles * 3;

        // Expand vertex data to be per-triangle-vertex (not per-original-vertex)
        // This is needed because the shader accesses data by vertex_index (0, 1, 2, ...)
        // and we use non-indexed drawing with storage buffers
        let mut expanded_positions: Vec<f32> = Vec::with_capacity(triangles.len() * 3 * 4);
        let mut expanded_normals: Vec<f32> = Vec::with_capacity(triangles.len() * 3 * 4);
        let mut expanded_colors: Vec<f32> = Vec::with_capacity(triangles.len() * 3 * 4);
        let mut barycentric_data: Vec<f32> = Vec::with_capacity(triangles.len() * 3 * 4);
        let mut edge_is_real_data: Vec<f32> = Vec::with_capacity(triangles.len() * 3 * 4);

        let mut tri_vertex_idx = 0;
        for tri in triangles {
            // Barycentric coordinates for wireframe
            let bary_coords = [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]];

            for (i, &vi) in tri.iter().enumerate() {
                let v = vertices[vi as usize];
                expanded_positions.extend_from_slice(&[v.x, v.y, v.z, 1.0]);

                let n = vertex_normals[vi as usize];
                expanded_normals.extend_from_slice(&[n.x, n.y, n.z, 0.0]);

                // Default color (will be overwritten by surface_color in shader)
                expanded_colors.extend_from_slice(&[0.0, 0.0, 0.0, 0.0]);

                barycentric_data.extend_from_slice(&[
                    bary_coords[i][0],
                    bary_coords[i][1],
                    bary_coords[i][2],
                    0.0,
                ]);

                // Edge is real flags (same for all vertices of triangle)
                let eir = edge_is_real[tri_vertex_idx];
                edge_is_real_data.extend_from_slice(&[eir.x, eir.y, eir.z, 0.0]);

                tri_vertex_idx += 1;
            }
        }

        // Create vertex position buffer (vec4 for alignment)
        let vertex_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("mesh vertices"),
            contents: bytemuck::cast_slice(&expanded_positions),
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
        });

        // Create index buffer (sequential indices since data is expanded)
        let index_data: Vec<u32> = (0..num_indices).collect();
        let index_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("mesh indices"),
            contents: bytemuck::cast_slice(&index_data),
            usage: wgpu::BufferUsages::INDEX | wgpu::BufferUsages::COPY_DST,
        });

        // Create normal buffer (vec4 for alignment)
        let normal_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("mesh normals"),
            contents: bytemuck::cast_slice(&expanded_normals),
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
        });

        // Create barycentric coordinate buffer
        let barycentric_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("mesh barycentrics"),
            contents: bytemuck::cast_slice(&barycentric_data),
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
        });

        // Create color buffer (expanded, default to zero - shader uses surface_color when zero)
        let color_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("mesh colors"),
            contents: bytemuck::cast_slice(&expanded_colors),
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
        });

        // Create edge_is_real buffer
        let edge_is_real_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("mesh edge_is_real"),
            contents: bytemuck::cast_slice(&edge_is_real_data),
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
        });

        // Create uniform buffer
        let uniforms = MeshUniforms::default();
        let uniform_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("mesh uniforms"),
            contents: bytemuck::cast_slice(&[uniforms]),
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
        });

        // Create bind group
        // Bindings:
        // 0: camera uniforms (uniform)
        // 1: mesh uniforms (uniform)
        // 2: positions (storage)
        // 3: normals (storage)
        // 4: barycentrics (storage)
        // 5: colors (storage)
        // 6: edge_is_real (storage)
        let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("surface mesh bind group"),
            layout: bind_group_layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: camera_buffer.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: uniform_buffer.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 2,
                    resource: vertex_buffer.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 3,
                    resource: normal_buffer.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 4,
                    resource: barycentric_buffer.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 5,
                    resource: color_buffer.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 6,
                    resource: edge_is_real_buffer.as_entire_binding(),
                },
            ],
        });

        Self {
            vertex_buffer,
            index_buffer,
            normal_buffer,
            barycentric_buffer,
            color_buffer,
            edge_is_real_buffer,
            uniform_buffer,
            bind_group,
            num_triangles,
            num_indices,
            shadow_bind_group: None,
            shadow_model_buffer: None,
        }
    }

    /// Updates the mesh uniform buffer.
    pub fn update_uniforms(&self, queue: &wgpu::Queue, uniforms: &MeshUniforms) {
        queue.write_buffer(&self.uniform_buffer, 0, bytemuck::cast_slice(&[*uniforms]));
    }

    /// Updates the per-vertex color buffer with per-original-vertex colors.
    /// This expands the colors to match the per-triangle-vertex buffer layout.
    pub fn update_colors(&self, queue: &wgpu::Queue, colors: &[Vec4], triangles: &[[u32; 3]]) {
        // Expand per-vertex colors to per-triangle-vertex
        let mut expanded_colors: Vec<f32> = Vec::with_capacity(triangles.len() * 3 * 4);
        for tri in triangles {
            for &vi in tri {
                let c = colors[vi as usize];
                expanded_colors.extend_from_slice(&[c.x, c.y, c.z, c.w]);
            }
        }
        queue.write_buffer(
            &self.color_buffer,
            0,
            bytemuck::cast_slice(&expanded_colors),
        );
    }

    /// Clears the color buffer (sets all colors to zero, which means use `surface_color`).
    pub fn clear_colors(&self, queue: &wgpu::Queue) {
        let zero_colors: Vec<f32> = vec![0.0; self.num_indices as usize * 4];
        queue.write_buffer(&self.color_buffer, 0, bytemuck::cast_slice(&zero_colors));
    }

    /// Initializes shadow rendering resources.
    ///
    /// Creates the bind group and model buffer needed for the shadow pass.
    /// The shadow pass renders this mesh from the light's perspective to create
    /// shadows on the ground plane.
    pub fn init_shadow_resources(
        &mut self,
        device: &wgpu::Device,
        shadow_bind_group_layout: &wgpu::BindGroupLayout,
        light_buffer: &wgpu::Buffer,
    ) {
        // Create model uniform buffer for shadow pass
        let shadow_model_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("Shadow Model Uniform Buffer"),
            contents: bytemuck::cast_slice(&[ShadowModelUniforms::default()]),
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
        });

        // Create bind group matching shadow_map.wgsl:
        // binding 0: light uniforms (view_proj, light_dir)
        // binding 1: model uniforms (model matrix)
        // binding 2: vertex positions (storage buffer)
        let shadow_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("Surface Mesh Shadow Bind Group"),
            layout: shadow_bind_group_layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: light_buffer.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: shadow_model_buffer.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 2,
                    resource: self.vertex_buffer.as_entire_binding(),
                },
            ],
        });

        self.shadow_model_buffer = Some(shadow_model_buffer);
        self.shadow_bind_group = Some(shadow_bind_group);
    }

    /// Returns whether shadow resources have been initialized.
    #[must_use]
    pub fn has_shadow_resources(&self) -> bool {
        self.shadow_bind_group.is_some()
    }

    /// Returns the number of vertices for shadow pass rendering.
    ///
    /// The shadow pass uses non-indexed drawing with storage buffers,
    /// so this returns the number of expanded triangle vertices.
    #[must_use]
    pub fn num_vertices(&self) -> u32 {
        self.num_indices // Same as expanded triangle vertex count
    }

    /// Updates the shadow model uniform buffer with the current transform.
    ///
    /// This should be called whenever the mesh's transform changes to ensure
    /// shadows are rendered at the correct position.
    pub fn update_shadow_model(&self, queue: &wgpu::Queue, model_matrix: [[f32; 4]; 4]) {
        if let Some(buffer) = &self.shadow_model_buffer {
            let uniforms = ShadowModelUniforms {
                model: model_matrix,
            };
            queue.write_buffer(buffer, 0, bytemuck::cast_slice(&[uniforms]));
        }
    }

    // Buffer getters for reflection rendering

    /// Returns the uniform buffer.
    #[must_use]
    pub fn uniform_buffer(&self) -> &wgpu::Buffer {
        &self.uniform_buffer
    }

    /// Returns the position buffer.
    #[must_use]
    pub fn position_buffer(&self) -> &wgpu::Buffer {
        &self.vertex_buffer
    }

    /// Returns the normal buffer.
    #[must_use]
    pub fn normal_buffer(&self) -> &wgpu::Buffer {
        &self.normal_buffer
    }

    /// Returns the barycentric buffer.
    #[must_use]
    pub fn barycentric_buffer(&self) -> &wgpu::Buffer {
        &self.barycentric_buffer
    }

    /// Returns the color buffer.
    #[must_use]
    pub fn color_buffer(&self) -> &wgpu::Buffer {
        &self.color_buffer
    }

    /// Returns the edge is real buffer.
    #[must_use]
    pub fn edge_is_real_buffer(&self) -> &wgpu::Buffer {
        &self.edge_is_real_buffer
    }

    /// Returns the number of vertices (for non-indexed drawing).
    #[must_use]
    pub fn vertex_count(&self) -> u32 {
        self.num_indices
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_mesh_uniforms_size() {
        let size = std::mem::size_of::<MeshUniforms>();

        // Verify size is 16-byte aligned for GPU uniform buffers
        assert_eq!(
            size % 16,
            0,
            "MeshUniforms size ({} bytes) must be 16-byte aligned",
            size
        );

        // Expected size breakdown:
        // model_matrix: 64 bytes ([[f32; 4]; 4])
        // shade_style: 4 bytes (u32)
        // show_edges: 4 bytes (u32)
        // edge_width: 4 bytes (f32)
        // transparency: 4 bytes (f32)
        // surface_color: 16 bytes ([f32; 4])
        // edge_color: 16 bytes ([f32; 4])
        // backface_policy: 4 bytes (u32)
        // slice_planes_enabled: 4 bytes (u32)
        // use_vertex_color: 4 bytes (u32)
        // _pad1: 4 bytes (f32)
        // _pad2: 12 bytes ([f32; 3])
        // _pad3: 4 bytes (f32)
        // backface_color: 16 bytes ([f32; 4])
        // Total: 160 bytes (matches WGSL layout with vec3 alignment)
        assert_eq!(
            size, 160,
            "MeshUniforms should be 160 bytes, got {} bytes",
            size
        );
    }
}