// Meshlet render pass — vertex fallback + mesh shader paths.
// Renders meshlet geometry using indirect draw commands from the cull pass.

use crate::formats::DEPTH_FORMAT;
use crate::shader;

/// Render path selection.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MeshletRenderPath {
    /// Always available. Reads meshlet data from storage buffers.
    VertexFallback,
    /// Requires hardware mesh shader support. Currently stub.
    MeshShader,
}

/// Meshlet render pipeline. Selects vertex fallback or mesh shader based on GPU caps.
pub struct MeshletRenderPass {
    pub path: MeshletRenderPath,
    fallback_pipeline: wgpu::RenderPipeline,
    bind_group_layout: wgpu::BindGroupLayout,
    /// Mesh shader pipeline (created when hardware supports EXPERIMENTAL_MESH_SHADER).
    mesh_pipeline: Option<wgpu::RenderPipeline>,
    /// Cached bind group — recreated only when buffers change.
    cached_bind_group: Option<CachedMeshletBindGroup>,
}

/// Cached bind group for meshlet rendering.
struct CachedMeshletBindGroup {
    bind_group: wgpu::BindGroup,
}

impl MeshletRenderPass {
    /// Create the meshlet render pass. Selects path based on adapter capabilities.
    pub fn new(device: &wgpu::Device, surface_format: wgpu::TextureFormat, mesh_shader_enabled: bool) -> Self {
        // Always build fallback; mesh shader path is a future addition.
        let path = if mesh_shader_enabled {
            MeshletRenderPath::MeshShader
        } else {
            MeshletRenderPath::VertexFallback
        };

        let (fallback_pipeline, bind_group_layout) = Self::build_fallback_pipeline(device, surface_format);

        // Attempt to build mesh shader pipeline if hardware supports it.
        let mesh_pipeline = if mesh_shader_enabled {
            Self::build_mesh_pipeline(device, surface_format, &bind_group_layout)
        } else {
            None
        };

        if mesh_pipeline.is_some() {
            log::info!("Meshlet render: mesh shader pipeline active");
        }

        Self {
            path,
            fallback_pipeline,
            bind_group_layout,
            mesh_pipeline,
            cached_bind_group: None,
        }
    }

    /// Render meshlets into the given render pass using indirect draw.
    pub fn render(
        &mut self,
        device: &wgpu::Device,
        pass: &mut wgpu::RenderPass<'_>,
        observer_buf: &wgpu::Buffer,
        meshlet_buf: &wgpu::Buffer,
        vertex_buf: &wgpu::Buffer,
        triangle_buf: &wgpu::Buffer,
        indirect_buf: &wgpu::Buffer,
        meshlet_count: u32,
    ) {
        if meshlet_count == 0 {
            return;
        }

        // Recreate bind group each frame. Bind group creation is lightweight
        // compared to actual draw calls, and avoids tracking buffer identity
        // (wgpu 27 removed global_id/Id APIs).
        {
            let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: Some("meshlet_render_bg"),
                layout: &self.bind_group_layout,
                entries: &[
                    wgpu::BindGroupEntry {
                        binding: 0,
                        resource: observer_buf.as_entire_binding(),
                    },
                    wgpu::BindGroupEntry {
                        binding: 1,
                        resource: meshlet_buf.as_entire_binding(),
                    },
                    wgpu::BindGroupEntry {
                        binding: 2,
                        resource: vertex_buf.as_entire_binding(),
                    },
                    wgpu::BindGroupEntry {
                        binding: 3,
                        resource: triangle_buf.as_entire_binding(),
                    },
                ],
            });
            self.cached_bind_group = Some(CachedMeshletBindGroup { bind_group });
        }

        let bg = &self.cached_bind_group.as_ref().unwrap().bind_group;
        pass.set_bind_group(0, bg, &[]);

        // Use mesh shader pipeline if available, otherwise vertex fallback.
        if let Some(ref mesh_pipe) = self.mesh_pipeline {
            pass.set_pipeline(mesh_pipe);
            // Mesh shaders dispatch workgroups instead of draw calls.
            // One workgroup per meshlet — the task shader culls, mesh shader emits geometry.
            pass.draw_mesh_tasks(meshlet_count, 1, 1);
        } else {
            pass.set_pipeline(&self.fallback_pipeline);
            // Vertex fallback: one indirect draw per meshlet.
            // The cull pass zeros instance_count for culled meshlets.
            let stride = std::mem::size_of::<crate::quantum_cull::DrawIndirectArgs>() as u64;
            for i in 0..meshlet_count {
                pass.draw_indirect(indirect_buf, i as u64 * stride);
            }
        }
    }

    /// Access the bind group layout.
    pub fn bind_group_layout(&self) -> &wgpu::BindGroupLayout {
        &self.bind_group_layout
    }

    /// Build mesh shader pipeline using the experimental mesh shader API.
    /// Returns None if shader compilation fails (expected on most current drivers).
    fn build_mesh_pipeline(
        device: &wgpu::Device,
        surface_format: wgpu::TextureFormat,
        bind_group_layout: &wgpu::BindGroupLayout,
    ) -> Option<wgpu::RenderPipeline> {
        // The mesh shader API is experimental in wgpu. We attempt to create the
        // pipeline and return None if it fails (shader compilation, missing features, etc.).
        // This is a best-effort path — the vertex fallback always works.

        let mesh_shader_src = r#"
            // Minimal mesh shader stub — outputs nothing when the experimental API
            // is available but the full shader isn't loaded.
            // Real implementation will read from storage buffers.
            @mesh(max_vertices = 3, max_primitives = 1) @workgroup_size(1)
            fn mesh_main() {
                // Stub: no output. Will be replaced with full meshlet emission.
            }

            @fragment
            fn fs_main() -> @location(0) vec4<f32> {
                return vec4<f32>(0.5, 0.5, 0.5, 1.0);
            }
        "#;

        // Try to create the shader module — may fail on drivers without mesh shader support
        let mesh_module = match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
            device.create_shader_module(wgpu::ShaderModuleDescriptor {
                label: Some("meshlet_mesh_shader"),
                source: wgpu::ShaderSource::Wgsl(mesh_shader_src.into()),
            })
        })) {
            Ok(module) => module,
            Err(_) => {
                log::info!("Mesh shader compilation not available on this driver");
                return None;
            }
        };

        let ctx = crate::mesh_shader::MeshShaderContext {
            enabled: true,
            ..Default::default()
        };
        ctx.create_pipeline(
            device,
            &crate::mesh_shader::MeshPipelineDesc {
                label: "meshlet_mesh_pipeline",
                bind_group_layouts: &[Some(bind_group_layout)],
                task_module: None,
                task_entry: None,
                mesh_module: &mesh_module,
                mesh_entry: Some("mesh_main"),
                fragment_module: Some(&mesh_module),
                fragment_entry: Some("fs_main"),
                topology: wgpu::PrimitiveTopology::TriangleList,
                cull_mode: Some(wgpu::Face::Back),
                depth_test: true,
                depth_write: true,
                sample_count: 1,
                color_targets: &[Some(wgpu::ColorTargetState {
                    format: surface_format,
                    blend: Some(wgpu::BlendState::REPLACE),
                    write_mask: wgpu::ColorWrites::ALL,
                })],
            },
        )
    }

    fn build_fallback_pipeline(
        device: &wgpu::Device,
        surface_format: wgpu::TextureFormat,
    ) -> (wgpu::RenderPipeline, wgpu::BindGroupLayout) {
        let module = shader::load_wgsl_embedded(
            device,
            "meshlet_vertex_fallback",
            include_str!("../shaders/meshlet_vertex_fallback.wgsl"),
        );

        let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("meshlet_render_bgl"),
            entries: &[
                // binding 0: uniforms (uniform)
                wgpu::BindGroupLayoutEntry {
                    binding: 0,
                    visibility: wgpu::ShaderStages::VERTEX | wgpu::ShaderStages::FRAGMENT,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Uniform,
                        has_dynamic_offset: false,
                        min_binding_size: None,
                    },
                    count: None,
                },
                // binding 1: meshlets (storage, read-only)
                wgpu::BindGroupLayoutEntry {
                    binding: 1,
                    visibility: wgpu::ShaderStages::VERTEX,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Storage { read_only: true },
                        has_dynamic_offset: false,
                        min_binding_size: None,
                    },
                    count: None,
                },
                // binding 2: vertices (storage, read-only)
                wgpu::BindGroupLayoutEntry {
                    binding: 2,
                    visibility: wgpu::ShaderStages::VERTEX,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Storage { read_only: true },
                        has_dynamic_offset: false,
                        min_binding_size: None,
                    },
                    count: None,
                },
                // binding 3: triangles (storage, read-only)
                wgpu::BindGroupLayoutEntry {
                    binding: 3,
                    visibility: wgpu::ShaderStages::VERTEX,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Storage { read_only: true },
                        has_dynamic_offset: false,
                        min_binding_size: None,
                    },
                    count: None,
                },
            ],
        });

        let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("meshlet_render_layout"),
            bind_group_layouts: &[Some(&bind_group_layout)],
            immediate_size: 0,
        });

        let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: Some("meshlet_vertex_fallback"),
            layout: Some(&pipeline_layout),
            vertex: wgpu::VertexState {
                module: &module,
                entry_point: Some("vs_main"),
                compilation_options: Default::default(),
                buffers: &[], // All data from storage buffers
            },
            fragment: Some(wgpu::FragmentState {
                module: &module,
                entry_point: Some("fs_main"),
                compilation_options: Default::default(),
                targets: &[Some(wgpu::ColorTargetState {
                    format: surface_format,
                    blend: Some(wgpu::BlendState::REPLACE),
                    write_mask: wgpu::ColorWrites::ALL,
                })],
            }),
            primitive: wgpu::PrimitiveState {
                topology: wgpu::PrimitiveTopology::TriangleList,
                front_face: wgpu::FrontFace::Ccw,
                cull_mode: Some(wgpu::Face::Back),
                ..Default::default()
            },
            depth_stencil: Some(wgpu::DepthStencilState {
                format: DEPTH_FORMAT,
                depth_write_enabled: Some(true),
                depth_compare: Some(wgpu::CompareFunction::Less),
                stencil: wgpu::StencilState::default(),
                bias: wgpu::DepthBiasState::default(),
            }),
            multisample: wgpu::MultisampleState::default(),
            multiview_mask: None,
            cache: None,
        });

        (pipeline, bind_group_layout)
    }
}

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

    #[test]
    fn render_path_default_is_fallback() {
        // Without mesh shader support, should select fallback.
        assert_ne!(MeshletRenderPath::VertexFallback, MeshletRenderPath::MeshShader);
    }

    #[test]
    fn indirect_stride_is_16() {
        assert_eq!(std::mem::size_of::<crate::quantum_cull::DrawIndirectArgs>(), 16);
    }
}
