// Mesh shader pipeline support.
// Wraps wgpu's experimental mesh/task shader API for meshlet-based rendering.
// Requires Features::EXPERIMENTAL_MESH_SHADER.

use crate::formats::DEPTH_FORMAT;

/// Mesh shader context — manages mesh pipelines for meshlet rendering.
#[derive(Default)]
pub struct MeshShaderContext {
    pub enabled: bool,
    pub max_output_layers: u32,
    pub max_task_mesh_workgroup_total_count: u32,
    pub max_task_mesh_workgroups_per_dimension: u32,
}

impl MeshShaderContext {
    /// Initialize mesh shader context from device limits.
    pub fn init_if_supported(features: wgpu::Features, limits: &wgpu::Limits) -> Self {
        let enabled = features.contains(wgpu::Features::EXPERIMENTAL_MESH_SHADER);
        if !enabled {
            return Self::default();
        }

        Self {
            enabled,
            max_output_layers: limits.max_mesh_output_layers,
            max_task_mesh_workgroup_total_count: limits.max_task_mesh_workgroup_total_count,
            max_task_mesh_workgroups_per_dimension: limits.max_task_mesh_workgroups_per_dimension,
        }
    }

    /// Create a mesh pipeline with optional task shader.
    pub fn create_pipeline(&self, device: &wgpu::Device, desc: &MeshPipelineDesc<'_>) -> Option<wgpu::RenderPipeline> {
        if !self.enabled {
            return None;
        }

        let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some(desc.label),
            bind_group_layouts: desc.bind_group_layouts,
            immediate_size: 0,
        });

        let task = desc.task_module.map(|module| wgpu::TaskState {
            module,
            entry_point: Some(desc.task_entry.unwrap_or("task_main")),
            compilation_options: Default::default(),
        });

        let depth_stencil = if desc.depth_test {
            Some(wgpu::DepthStencilState {
                format: DEPTH_FORMAT,
                depth_write_enabled: Some(desc.depth_write),
                depth_compare: Some(wgpu::CompareFunction::Less),
                stencil: wgpu::StencilState::default(),
                bias: wgpu::DepthBiasState::default(),
            })
        } else {
            None
        };

        let pipeline = device.create_mesh_pipeline(&wgpu::MeshPipelineDescriptor {
            label: Some(desc.label),
            layout: Some(&pipeline_layout),
            task,
            mesh: wgpu::MeshState {
                module: desc.mesh_module,
                entry_point: Some(desc.mesh_entry.unwrap_or("mesh_main")),
                compilation_options: Default::default(),
            },
            primitive: wgpu::PrimitiveState {
                topology: desc.topology,
                front_face: wgpu::FrontFace::Ccw,
                cull_mode: desc.cull_mode,
                ..Default::default()
            },
            depth_stencil,
            multisample: wgpu::MultisampleState {
                count: desc.sample_count,
                ..Default::default()
            },
            fragment: desc.fragment_module.map(|module| wgpu::FragmentState {
                module,
                entry_point: Some(desc.fragment_entry.unwrap_or("fs_main")),
                compilation_options: Default::default(),
                targets: desc.color_targets,
            }),
            multiview: None,
            cache: None,
        });

        Some(pipeline)
    }

    /// Recommended workgroup size for mesh shaders on this adapter.
    /// NVIDIA: 32 threads, AMD: 64 threads, fallback: 32.
    pub fn recommended_workgroup_size(&self, adapter_name: &str) -> u32 {
        let name_lower = adapter_name.to_lowercase();
        if name_lower.contains("amd") || name_lower.contains("radeon") {
            64
        } else {
            32
        }
    }
}

/// Descriptor for mesh pipeline creation.
pub struct MeshPipelineDesc<'a> {
    pub label: &'a str,
    pub bind_group_layouts: &'a [Option<&'a wgpu::BindGroupLayout>],

    // Shader modules
    pub task_module: Option<&'a wgpu::ShaderModule>,
    pub task_entry: Option<&'a str>,
    pub mesh_module: &'a wgpu::ShaderModule,
    pub mesh_entry: Option<&'a str>,
    pub fragment_module: Option<&'a wgpu::ShaderModule>,
    pub fragment_entry: Option<&'a str>,

    // Pipeline state
    pub topology: wgpu::PrimitiveTopology,
    pub cull_mode: Option<wgpu::Face>,
    pub depth_test: bool,
    pub depth_write: bool,
    pub sample_count: u32,
    pub color_targets: &'a [Option<wgpu::ColorTargetState>],
}

/// Meshlet data structure for CPU-side meshlet generation.
#[derive(Debug, Clone)]
#[repr(C)]
pub struct Meshlet {
    pub vertex_offset: u32,
    pub vertex_count: u32,
    pub triangle_offset: u32,
    pub triangle_count: u32,
}

/// Build meshlets from an indexed triangle mesh.
/// max_vertices and max_triangles define per-meshlet limits.
pub fn build_meshlets(indices: &[u32], max_vertices: u32, max_triangles: u32) -> Vec<Meshlet> {
    assert!(
        indices.len().is_multiple_of(3),
        "meshlet: index count not divisible by 3"
    );
    let max_t = max_triangles as usize;
    let mut meshlets = Vec::new();
    let triangle_count = indices.len() / 3;
    let mut tri_idx = 0;

    while tri_idx < triangle_count {
        let mut unique_verts: Vec<u32> = Vec::with_capacity(max_vertices as usize);
        let start_tri = tri_idx;
        let mut tris_in_meshlet = 0;

        while tri_idx < triangle_count && tris_in_meshlet < max_t {
            let i0 = indices[tri_idx * 3];
            let i1 = indices[tri_idx * 3 + 1];
            let i2 = indices[tri_idx * 3 + 2];

            let mut new_verts = 0u32;
            if !unique_verts.contains(&i0) {
                new_verts += 1;
            }
            if !unique_verts.contains(&i1) {
                new_verts += 1;
            }
            if !unique_verts.contains(&i2) {
                new_verts += 1;
            }

            if unique_verts.len() as u32 + new_verts > max_vertices {
                break;
            }

            if !unique_verts.contains(&i0) {
                unique_verts.push(i0);
            }
            if !unique_verts.contains(&i1) {
                unique_verts.push(i1);
            }
            if !unique_verts.contains(&i2) {
                unique_verts.push(i2);
            }

            tri_idx += 1;
            tris_in_meshlet += 1;
        }

        meshlets.push(Meshlet {
            vertex_offset: 0, // resolved by caller
            vertex_count: unique_verts.len() as u32,
            triangle_offset: (start_tri * 3) as u32,
            triangle_count: tris_in_meshlet as u32,
        });
    }

    meshlets
}

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

    #[test]
    fn default_disabled() {
        let ctx = MeshShaderContext::default();
        assert!(!ctx.enabled);
    }

    #[test]
    fn workgroup_size_amd() {
        let ctx = MeshShaderContext::default();
        assert_eq!(ctx.recommended_workgroup_size("AMD Radeon RX 7900"), 64);
    }

    #[test]
    fn workgroup_size_nvidia() {
        let ctx = MeshShaderContext::default();
        assert_eq!(ctx.recommended_workgroup_size("NVIDIA RTX 4090"), 32);
    }

    #[test]
    fn build_meshlets_simple() {
        let indices = vec![0, 1, 2, 1, 3, 2];
        let meshlets = build_meshlets(&indices, 64, 126);
        assert_eq!(meshlets.len(), 1);
        assert_eq!(meshlets[0].triangle_count, 2);
        assert_eq!(meshlets[0].vertex_count, 4);
    }

    #[test]
    fn build_meshlets_splits() {
        let indices = vec![0, 1, 2, 3, 4, 5];
        let meshlets = build_meshlets(&indices, 3, 126);
        assert_eq!(meshlets.len(), 2);
    }
}
