// Quantum Culling — GPU compute pass for per-meshlet visibility.
// Reads observer context + meshlet descriptors, writes indirect draw commands.
// Layer filtering + sphere-vs-FOV distance test.
// Group 1: optional QuantumContext for Born-rule LOD and phi-derived thresholds.

use crate::dreamlet_catalog::GpuQuantumContext;
use crate::shader;

/// Indirect draw arguments (16 bytes, matches wgpu DrawIndirect layout).
#[repr(C)]
#[derive(Debug, Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
pub struct DrawIndirectArgs {
    pub vertex_count: u32,
    pub instance_count: u32,
    pub first_vertex: u32,
    pub first_instance: u32,
}

/// GPU compute pipeline for Quantum Culling.
pub struct QuantumCullPass {
    pipeline: wgpu::ComputePipeline,
    bind_group_layout: wgpu::BindGroupLayout,
    /// Group 1: quantum context bind group layout.
    quantum_bgl: wgpu::BindGroupLayout,
    /// Cached bind group — recreated only when buffers change.
    cached_bind_group: Option<CachedCullBindGroup>,
    /// Fallback quantum context buffer (144 bytes zeroed) for when no bridge is active.
    fallback_quantum_buf: wgpu::Buffer,
}

/// Cached bind group for the cull pass.
struct CachedCullBindGroup {
    bind_group: wgpu::BindGroup,
}

impl QuantumCullPass {
    /// Create the cull pipeline with embedded shader.
    pub fn new(device: &wgpu::Device) -> Self {
        let module = shader::load_wgsl_embedded(device, "quantum_cull", include_str!("../shaders/quantum_cull.wgsl"));

        let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("quantum_cull_bgl"),
            entries: &[
                // binding 0: observer context (uniform)
                wgpu::BindGroupLayoutEntry {
                    binding: 0,
                    visibility: wgpu::ShaderStages::COMPUTE,
                    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::COMPUTE,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Storage { read_only: true },
                        has_dynamic_offset: false,
                        min_binding_size: None,
                    },
                    count: None,
                },
                // binding 2: indirect draw buffer (storage, read-write)
                wgpu::BindGroupLayoutEntry {
                    binding: 2,
                    visibility: wgpu::ShaderStages::COMPUTE,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Storage { read_only: false },
                        has_dynamic_offset: false,
                        min_binding_size: None,
                    },
                    count: None,
                },
            ],
        });

        // Group 1: quantum context (storage, read-only)
        let quantum_bgl = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("quantum_cull_quantum_bgl"),
            entries: &[wgpu::BindGroupLayoutEntry {
                binding: 0,
                visibility: wgpu::ShaderStages::COMPUTE,
                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("quantum_cull_layout"),
            bind_group_layouts: &[Some(&bind_group_layout), Some(&quantum_bgl)],
            immediate_size: 0,
        });

        let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: Some("quantum_cull"),
            layout: Some(&pipeline_layout),
            module: &module,
            entry_point: Some("cull_main"),
            compilation_options: Default::default(),
            cache: None,
        });

        // Fallback quantum context: 144 bytes zeroed (classical default).
        let fallback_quantum_buf = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("quantum_cull_fallback_ctx"),
            size: std::mem::size_of::<GpuQuantumContext>() as u64,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        Self {
            pipeline,
            bind_group_layout,
            quantum_bgl,
            cached_bind_group: None,
            fallback_quantum_buf,
        }
    }

    /// Dispatch the cull pass. Caches bind group and recreates only when buffers change.
    /// Zeros the indirect buffer before dispatch to prevent stale draw args.
    /// `indirect_capacity` is the maximum number of entries the indirect buffer can hold.
    /// `meshlet_count` is clamped to this capacity to prevent out-of-bounds writes.
    /// `quantum_context_buf`: optional GPU buffer with GpuQuantumContext for Born-rule LOD.
    pub fn dispatch(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        encoder: &mut wgpu::CommandEncoder,
        observer_buf: &wgpu::Buffer,
        meshlet_buf: &wgpu::Buffer,
        indirect_buf: &wgpu::Buffer,
        meshlet_count: u32,
        quantum_context_buf: Option<&wgpu::Buffer>,
    ) {
        // Clamp meshlet_count to indirect buffer capacity.
        // The indirect buffer byte size = capacity * sizeof(DrawIndirectArgs).
        // We derive capacity from the buffer size to avoid requiring a separate parameter.
        let indirect_capacity = (indirect_buf.size() / std::mem::size_of::<DrawIndirectArgs>() as u64) as u32;
        let meshlet_count = meshlet_count.min(indirect_capacity);
        if meshlet_count < indirect_capacity && meshlet_count > 0 {
            // Normal case: meshlet_count fits within buffer. No warning needed.
        }

        // Zero indirect buffer before cull — prevents stale draw args from prior frames.
        let zero_size = (meshlet_count as u64) * std::mem::size_of::<DrawIndirectArgs>() as u64;
        if zero_size > 0 {
            let zeros = vec![0u8; zero_size as usize];
            queue.write_buffer(indirect_buf, 0, &zeros);
        }

        // Recreate bind group each frame. Bind group creation is lightweight
        // compared to compute dispatch, and avoids tracking buffer identity
        // (wgpu 27 removed global_id/Id APIs).
        {
            let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: Some("quantum_cull_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: indirect_buf.as_entire_binding(),
                    },
                ],
            });
            self.cached_bind_group = Some(CachedCullBindGroup { bind_group });
        }

        // Group 1: quantum context (use caller's buffer or fallback zeroed).
        let q_buf = quantum_context_buf.unwrap_or(&self.fallback_quantum_buf);
        let quantum_bg = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("quantum_cull_quantum_bg"),
            layout: &self.quantum_bgl,
            entries: &[wgpu::BindGroupEntry {
                binding: 0,
                resource: q_buf.as_entire_binding(),
            }],
        });

        let bg = &self
            .cached_bind_group
            .as_ref()
            .expect("bind group created above")
            .bind_group;
        let workgroups = meshlet_count.div_ceil(64);
        let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
            label: Some("quantum_cull"),
            timestamp_writes: None,
        });
        pass.set_pipeline(&self.pipeline);
        pass.set_bind_group(0, bg, &[]);
        pass.set_bind_group(1, &quantum_bg, &[]);
        pass.dispatch_workgroups(workgroups, 1, 1);
    }

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

/// Pre-allocated indirect draw buffer.
pub struct IndirectDrawBuffer {
    pub buffer: wgpu::Buffer,
    capacity: u32,
}

impl IndirectDrawBuffer {
    /// Create an indirect draw buffer with capacity for `count` meshlets.
    pub fn new(device: &wgpu::Device, capacity: u32) -> Self {
        let size = (capacity as u64) * std::mem::size_of::<DrawIndirectArgs>() as u64;
        let buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("indirect_draw"),
            size: size.max(16), // minimum 1 entry
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::INDIRECT | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        Self { buffer, capacity }
    }

    /// Current capacity in meshlet count.
    pub fn capacity(&self) -> u32 {
        self.capacity
    }

    /// Explicitly destroy GPU resources.
    pub fn destroy(&mut self) {
        self.buffer.destroy();
    }
}

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

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

    #[test]
    fn draw_indirect_args_zeroed() {
        let args = DrawIndirectArgs::zeroed();
        assert_eq!(args.vertex_count, 0);
        assert_eq!(args.instance_count, 0);
    }

    #[test]
    fn workgroup_ceil_div() {
        assert_eq!(100u32.div_ceil(64), 2);
        assert_eq!(64u32.div_ceil(64), 1);
        assert_eq!(65u32.div_ceil(64), 2);
    }
}
