//! Particle simulation compute pass.

use super::buffers::GpuPhysicsBuffers;

/// Simulation compute pass — integrates position/velocity, applies forces.
pub struct SimulationPass {
    pipeline: wgpu::ComputePipeline,
    bind_group_layout: wgpu::BindGroupLayout,
    bind_group: Option<wgpu::BindGroup>,
    sim_params_buffer: wgpu::Buffer,
    needs_rebind: bool,
}

#[repr(C)]
#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
struct SimParamsGpu {
    delta_time: f32,
    total_time: f32,
    particle_count: u32,
    gravity_x: f32,
    gravity_y: f32,
    gravity_z: f32,
    drag_coeff: f32,
    damping: f32,
    _pad0: f32,
    _pad1: f32,
    _pad2: f32,
    _pad3: f32,
}

impl SimulationPass {
    pub fn new(device: &wgpu::Device) -> Self {
        let common_src = include_str!("../../shaders/physics/physics_common.wgsl");
        let sim_src = include_str!("../../shaders/physics/particle_simulate.wgsl");
        let full_src = format!("{}\n{}", common_src, sim_src);

        let module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("physics_simulate"),
            source: wgpu::ShaderSource::Wgsl(full_src.into()),
        });

        let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("physics_simulate_layout"),
            entries: &[
                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,
                },
                wgpu::BindGroupLayoutEntry {
                    binding: 1,
                    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,
                },
                wgpu::BindGroupLayoutEntry {
                    binding: 2,
                    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,
                },
                wgpu::BindGroupLayoutEntry {
                    binding: 3,
                    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("physics_simulate_pipeline_layout"),
            bind_group_layouts: &[Some(&bind_group_layout)],
            immediate_size: 0,
        });

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

        let sim_params_buffer = crate::compute::create_buffer_empty(
            device,
            "physics_sim_params",
            std::mem::size_of::<SimParamsGpu>() as u64,
            wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
        );

        Self {
            pipeline,
            bind_group_layout,
            bind_group: None,
            sim_params_buffer,
            needs_rebind: true,
        }
    }

    fn ensure_bind_group(&mut self, device: &wgpu::Device, buffers: &GpuPhysicsBuffers) {
        if !self.needs_rebind && self.bind_group.is_some() {
            return;
        }
        self.bind_group = Some(device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("physics_simulate_bind_group"),
            layout: &self.bind_group_layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: self.sim_params_buffer.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: buffers.particle_state.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 2,
                    resource: buffers.force_field_buffer.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 3,
                    resource: buffers.debug_stats.as_entire_binding(),
                },
            ],
        }));
        self.needs_rebind = false;
    }

    pub fn dispatch(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        encoder: &mut wgpu::CommandEncoder,
        buffers: &GpuPhysicsBuffers,
        dt: f32,
        total_time: f32,
    ) {
        let params = SimParamsGpu {
            delta_time: dt,
            total_time,
            particle_count: buffers.max_particles,
            gravity_x: 0.0,
            gravity_y: -9.81,
            gravity_z: 0.0,
            drag_coeff: 0.05,
            damping: 0.01,
            _pad0: 0.0,
            _pad1: 0.0,
            _pad2: 0.0,
            _pad3: 0.0,
        };

        queue.write_buffer(&self.sim_params_buffer, 0, bytemuck::bytes_of(&params));
        queue.write_buffer(&buffers.debug_stats, 0, &[0u8; 32]);

        self.ensure_bind_group(device, buffers);

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