// Compute shader abstraction.
// Wraps wgpu compute pipeline creation and dispatch.

/// Buffer binding type for compute kernels.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BufferKind {
    Storage,
    StorageReadOnly,
    Uniform,
}

impl BufferKind {
    fn to_binding_type(self) -> wgpu::BufferBindingType {
        match self {
            Self::Storage => wgpu::BufferBindingType::Storage { read_only: false },
            Self::StorageReadOnly => wgpu::BufferBindingType::Storage { read_only: true },
            Self::Uniform => wgpu::BufferBindingType::Uniform,
        }
    }
}

/// Descriptor for a single buffer binding in a compute kernel.
pub struct ComputeBinding<'a> {
    pub binding: u32,
    pub kind: BufferKind,
    pub resource: wgpu::BindingResource<'a>,
}

/// A compute kernel: pipeline + bind group, ready to dispatch.
///
/// This struct is intended for one-time creation, not per-frame use.
/// Creating a `ComputeKernel` involves pipeline compilation and layout creation,
/// which are expensive GPU operations. Callers that need to rebind buffers
/// should cache the pipeline and layout separately and only recreate the bind group.
pub struct ComputeKernel {
    pub pipeline: wgpu::ComputePipeline,
    pub bind_group: wgpu::BindGroup,
}

impl ComputeKernel {
    /// Create a compute kernel from a shader module and entry point.
    ///
    /// This performs pipeline compilation — call once at initialization time,
    /// not per-frame. For pipelines that need buffer rebinding, cache the
    /// pipeline and layout separately (see `GpuTweenContext` for an example).
    pub fn new(
        device: &wgpu::Device,
        module: &wgpu::ShaderModule,
        entry_point: &str,
        bindings: &[ComputeBinding<'_>],
    ) -> Self {
        let layout_entries: Vec<wgpu::BindGroupLayoutEntry> = bindings
            .iter()
            .map(|b| wgpu::BindGroupLayoutEntry {
                binding: b.binding,
                visibility: wgpu::ShaderStages::COMPUTE,
                ty: wgpu::BindingType::Buffer {
                    ty: b.kind.to_binding_type(),
                    has_dynamic_offset: false,
                    min_binding_size: None,
                },
                count: None,
            })
            .collect();

        let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some(entry_point),
            entries: &layout_entries,
        });

        let bind_group_entries: Vec<wgpu::BindGroupEntry<'_>> = bindings
            .iter()
            .map(|b| wgpu::BindGroupEntry {
                binding: b.binding,
                resource: b.resource.clone(),
            })
            .collect();

        let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some(entry_point),
            layout: &bind_group_layout,
            entries: &bind_group_entries,
        });

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

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

        Self { pipeline, bind_group }
    }

    /// Dispatch this kernel with the given workgroup counts.
    pub fn dispatch(&self, encoder: &mut wgpu::CommandEncoder, x: u32, y: u32, z: u32) {
        let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
            label: None,
            timestamp_writes: None,
        });
        pass.set_pipeline(&self.pipeline);
        pass.set_bind_group(0, &self.bind_group, &[]);
        pass.dispatch_workgroups(x, y, z);
    }
}

/// Helper to create a GPU buffer with initial data.
pub fn create_buffer_init(device: &wgpu::Device, label: &str, data: &[u8], usage: wgpu::BufferUsages) -> wgpu::Buffer {
    use wgpu::util::DeviceExt;
    device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
        label: Some(label),
        contents: data,
        usage,
    })
}

/// Helper to create an empty GPU buffer of given size.
pub fn create_buffer_empty(device: &wgpu::Device, label: &str, size: u64, usage: wgpu::BufferUsages) -> wgpu::Buffer {
    device.create_buffer(&wgpu::BufferDescriptor {
        label: Some(label),
        size,
        usage,
        mapped_at_creation: false,
    })
}

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

    #[test]
    fn buffer_kind_variants() {
        assert_eq!(
            BufferKind::Storage.to_binding_type(),
            wgpu::BufferBindingType::Storage { read_only: false }
        );
        assert_eq!(
            BufferKind::StorageReadOnly.to_binding_type(),
            wgpu::BufferBindingType::Storage { read_only: true }
        );
        assert_eq!(BufferKind::Uniform.to_binding_type(), wgpu::BufferBindingType::Uniform);
    }
}
