// MaterializePass — skeletal animation compute shader.
// Reads bone matrices + bind-pose vertices, writes skinned vertex positions.
// Called once per skinned entity per frame before QuantumCullPass.

use crate::shader;

/// Maximum joints supported per skeleton.
pub const MAX_JOINTS: usize = 256;

/// Bone matrix for GPU upload (64 bytes).
#[repr(C)]
#[derive(Debug, Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
pub struct BoneMatrix {
    pub transform: [f32; 16],
}

/// GPU buffer for skeleton bone matrices.
pub struct SkeletonBuffer {
    pub buffer: wgpu::Buffer,
}

impl SkeletonBuffer {
    /// Create a skeleton buffer for up to MAX_JOINTS bones.
    pub fn new(device: &wgpu::Device) -> Self {
        let size = (MAX_JOINTS as u64) * std::mem::size_of::<BoneMatrix>() as u64;
        let buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("skeleton_bones"),
            size,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        Self { buffer }
    }

    /// Upload bone matrices. No allocation.
    pub fn upload(&self, queue: &wgpu::Queue, matrices: &[BoneMatrix]) {
        let count = matrices.len().min(MAX_JOINTS);
        queue.write_buffer(&self.buffer, 0, bytemuck::cast_slice(&matrices[..count]));
    }

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

/// Materialize compute pass — applies skeletal bone transforms to mesh vertices.
pub struct MaterializePass {
    pipeline: wgpu::ComputePipeline,
    bind_group_layout: wgpu::BindGroupLayout,
}

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

        let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("materialize_bgl"),
            entries: &[
                // binding 0: bone matrices (storage, read)
                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,
                },
                // binding 1: input vertices (storage, read)
                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: output vertices (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,
                },
            ],
        });

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

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

        Self {
            pipeline,
            bind_group_layout,
        }
    }

    /// Dispatch the materialize compute pass for skinned vertices.
    pub fn dispatch(
        &self,
        device: &wgpu::Device,
        encoder: &mut wgpu::CommandEncoder,
        skeleton: &SkeletonBuffer,
        input_vertices: &wgpu::Buffer,
        output_vertices: &wgpu::Buffer,
        vertex_count: u32,
    ) {
        let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("materialize_bg"),
            layout: &self.bind_group_layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: skeleton.buffer.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: input_vertices.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 2,
                    resource: output_vertices.as_entire_binding(),
                },
            ],
        });

        let workgroups = vertex_count.div_ceil(64);
        let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
            label: Some("materialize"),
            timestamp_writes: None,
        });
        pass.set_pipeline(&self.pipeline);
        pass.set_bind_group(0, &bind_group, &[]);
        pass.dispatch_workgroups(workgroups, 1, 1);
    }
}

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

    #[test]
    fn bone_matrix_size() {
        assert_eq!(std::mem::size_of::<BoneMatrix>(), 64);
    }

    #[test]
    fn max_joints_256() {
        assert_eq!(MAX_JOINTS, 256);
    }
}
