// DreamMatter — observer-aware GPU particle system.
// Companion to PhysicsGpu. Simulates particles within observer FOV,
// promotes boundary-crossing particles to server authority.
// Three compute passes: simulate → emit → cull.

use crate::particle::Particle;
use crate::shader;

/// DreamMatter simulation parameters (48 bytes).
#[repr(C)]
#[derive(Debug, Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
pub struct DreamMatterParams {
    pub delta_time: f32,
    pub gravity: f32,
    pub damping: f32,
    pub particle_count: u32,
    pub observer_pos: [f32; 3],
    pub fov_radius: f32,
    pub topology_layer: u32,
    pub _pad: [u32; 3],
}

/// Promotion event read back from GPU (32 bytes).
#[repr(C)]
#[derive(Debug, Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
pub struct PromotionEvent {
    pub position: [f32; 3],
    pub age: f32,
    pub kind: u32,
    pub _pad: [u32; 3],
}

/// Maximum promotion events per frame.
pub const MAX_PROMOTION_EVENTS: u32 = 256;

/// Cached bind group for sim/emit passes.
struct CachedSimBindGroup {
    bind_group: wgpu::BindGroup,
}

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

/// DreamMatter GPU system — owns sim/emit/cull pipelines and particle storage.
pub struct DreamMatterSystem {
    sim_pipeline: wgpu::ComputePipeline,
    emit_pipeline: wgpu::ComputePipeline,
    cull_pipeline: wgpu::ComputePipeline,
    sim_bgl: wgpu::BindGroupLayout,
    cull_bgl: wgpu::BindGroupLayout,
    params_buffer: wgpu::Buffer,
    particle_buffer: Option<wgpu::Buffer>,
    promotion_buffer: wgpu::Buffer,
    promotion_count_buffer: wgpu::Buffer,
    promotion_readback: wgpu::Buffer,
    count_readback: wgpu::Buffer,
    capacity: u32,
    cached_promotions: Vec<PromotionEvent>,
    /// Cached bind group for sim pass (shared layout with emit).
    cached_sim_bg: Option<CachedSimBindGroup>,
    /// Cached bind group for emit pass (shared layout with sim).
    cached_emit_bg: Option<CachedSimBindGroup>,
    /// Cached bind group for cull pass.
    cached_cull_bg: Option<CachedCullBindGroup>,
}

impl DreamMatterSystem {
    /// Create the DreamMatter system.
    pub fn new(device: &wgpu::Device) -> Self {
        let sim_module = shader::load_wgsl_embedded(
            device,
            "dreammatter_simulate",
            include_str!("../shaders/dreammatter_simulate.wgsl"),
        );
        let emit_module = shader::load_wgsl_embedded(
            device,
            "dreammatter_emit",
            include_str!("../shaders/dreammatter_emit.wgsl"),
        );
        let cull_module = shader::load_wgsl_embedded(
            device,
            "dreammatter_cull",
            include_str!("../shaders/dreammatter_cull.wgsl"),
        );

        // Sim/emit bind group layout: uniform params + storage particles
        let sim_bgl = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("dreammatter_sim_bgl"),
            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,
                },
            ],
        });

        // Cull bind group layout: params + particles + promotions + count
        let cull_bgl = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("dreammatter_cull_bgl"),
            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: false },
                        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 sim_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("dreammatter_sim"),
            bind_group_layouts: &[Some(&sim_bgl)],
            immediate_size: 0,
        });
        let cull_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("dreammatter_cull"),
            bind_group_layouts: &[Some(&cull_bgl)],
            immediate_size: 0,
        });

        let sim_pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: Some("dreammatter_simulate"),
            layout: Some(&sim_layout),
            module: &sim_module,
            entry_point: Some("simulate"),
            compilation_options: Default::default(),
            cache: None,
        });
        let emit_pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: Some("dreammatter_emit"),
            layout: Some(&sim_layout),
            module: &emit_module,
            entry_point: Some("emit"),
            compilation_options: Default::default(),
            cache: None,
        });
        let cull_pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: Some("dreammatter_cull"),
            layout: Some(&cull_layout),
            module: &cull_module,
            entry_point: Some("cull"),
            compilation_options: Default::default(),
            cache: None,
        });

        let params_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("dreammatter_params"),
            size: std::mem::size_of::<DreamMatterParams>() as u64,
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        let promotion_size = (MAX_PROMOTION_EVENTS as u64) * std::mem::size_of::<PromotionEvent>() as u64;
        let promotion_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("dreammatter_promotions"),
            size: promotion_size,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
            mapped_at_creation: false,
        });

        let promotion_count_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("dreammatter_promotion_count"),
            size: 4,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        let promotion_readback = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("dreammatter_promotion_readback"),
            size: promotion_size,
            usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
            mapped_at_creation: false,
        });

        let count_readback = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("dreammatter_count_readback"),
            size: 4,
            usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
            mapped_at_creation: false,
        });

        Self {
            sim_pipeline,
            emit_pipeline,
            cull_pipeline,
            sim_bgl,
            cull_bgl,
            params_buffer,
            particle_buffer: None,
            promotion_buffer,
            promotion_count_buffer,
            promotion_readback,
            count_readback,
            capacity: 0,
            cached_promotions: Vec::new(),
            cached_sim_bg: None,
            cached_emit_bg: None,
            cached_cull_bg: None,
        }
    }

    /// Ensure particle buffer has capacity for `count` particles.
    pub fn ensure_capacity(&mut self, device: &wgpu::Device, count: u32) {
        if count <= self.capacity && self.particle_buffer.is_some() {
            return;
        }
        let new_cap = count.next_power_of_two().max(64);
        let size = (new_cap as u64) * std::mem::size_of::<Particle>() as u64;
        self.particle_buffer = Some(device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("dreammatter_particles"),
            size,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        }));
        self.capacity = new_cap;
        // Invalidate cached bind groups since particle buffer changed.
        self.cached_sim_bg = None;
        self.cached_emit_bg = None;
        self.cached_cull_bg = None;
    }

    /// Dispatch the simulate compute pass.
    pub fn dispatch_simulate(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        encoder: &mut wgpu::CommandEncoder,
        params: &DreamMatterParams,
    ) {
        let Some(particle_buf) = &self.particle_buffer else {
            return;
        };
        queue.write_buffer(&self.params_buffer, 0, bytemuck::bytes_of(params));

        if self.cached_sim_bg.is_none() {
            let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: Some("dreammatter_sim_bg"),
                layout: &self.sim_bgl,
                entries: &[
                    wgpu::BindGroupEntry {
                        binding: 0,
                        resource: self.params_buffer.as_entire_binding(),
                    },
                    wgpu::BindGroupEntry {
                        binding: 1,
                        resource: particle_buf.as_entire_binding(),
                    },
                ],
            });
            self.cached_sim_bg = Some(CachedSimBindGroup { bind_group });
        }

        let bg = &self
            .cached_sim_bg
            .as_ref()
            .expect("sim bind group created above")
            .bind_group;
        let wg = params.particle_count.div_ceil(64);
        let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
            label: Some("dreammatter_simulate"),
            timestamp_writes: None,
        });
        pass.set_pipeline(&self.sim_pipeline);
        pass.set_bind_group(0, bg, &[]);
        pass.dispatch_workgroups(wg, 1, 1);
    }

    /// Dispatch the emit compute pass.
    pub fn dispatch_emit(&mut self, device: &wgpu::Device, encoder: &mut wgpu::CommandEncoder, particle_count: u32) {
        let Some(particle_buf) = &self.particle_buffer else {
            return;
        };

        if self.cached_emit_bg.is_none() {
            let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: Some("dreammatter_emit_bg"),
                layout: &self.sim_bgl,
                entries: &[
                    wgpu::BindGroupEntry {
                        binding: 0,
                        resource: self.params_buffer.as_entire_binding(),
                    },
                    wgpu::BindGroupEntry {
                        binding: 1,
                        resource: particle_buf.as_entire_binding(),
                    },
                ],
            });
            self.cached_emit_bg = Some(CachedSimBindGroup { bind_group });
        }

        let bg = &self
            .cached_emit_bg
            .as_ref()
            .expect("emit bind group created above")
            .bind_group;
        let wg = particle_count.div_ceil(64);
        let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
            label: Some("dreammatter_emit"),
            timestamp_writes: None,
        });
        pass.set_pipeline(&self.emit_pipeline);
        pass.set_bind_group(0, bg, &[]);
        pass.dispatch_workgroups(wg, 1, 1);
    }

    /// Dispatch the cull compute pass. Resets promotion count first.
    pub fn dispatch_cull(
        &mut self,
        device: &wgpu::Device,
        _queue: &wgpu::Queue,
        encoder: &mut wgpu::CommandEncoder,
        particle_count: u32,
    ) {
        let Some(particle_buf) = &self.particle_buffer else {
            return;
        };
        // Reset promotion count to 0 via encoder (non-blocking, avoids queue.write_buffer stall).
        encoder.clear_buffer(&self.promotion_count_buffer, 0, None);

        if self.cached_cull_bg.is_none() {
            let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: Some("dreammatter_cull_bg"),
                layout: &self.cull_bgl,
                entries: &[
                    wgpu::BindGroupEntry {
                        binding: 0,
                        resource: self.params_buffer.as_entire_binding(),
                    },
                    wgpu::BindGroupEntry {
                        binding: 1,
                        resource: particle_buf.as_entire_binding(),
                    },
                    wgpu::BindGroupEntry {
                        binding: 2,
                        resource: self.promotion_buffer.as_entire_binding(),
                    },
                    wgpu::BindGroupEntry {
                        binding: 3,
                        resource: self.promotion_count_buffer.as_entire_binding(),
                    },
                ],
            });
            self.cached_cull_bg = Some(CachedCullBindGroup { bind_group });
        }

        let bg = &self
            .cached_cull_bg
            .as_ref()
            .expect("cull bind group created above")
            .bind_group;
        let wg = particle_count.div_ceil(64);
        let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
            label: Some("dreammatter_cull"),
            timestamp_writes: None,
        });
        pass.set_pipeline(&self.cull_pipeline);
        pass.set_bind_group(0, bg, &[]);
        pass.dispatch_workgroups(wg, 1, 1);
    }

    /// Copy promotion events from GPU to readback buffer.
    ///
    /// SAFETY: wgpu inserts an implicit pipeline barrier between the previous
    /// compute pass (dispatch_cull writes via storage atomics) and this copy
    /// operation. All atomic writes are visible before the copy executes.
    /// Call after cull dispatch, before submit.
    pub fn copy_promotion_events(&self, encoder: &mut wgpu::CommandEncoder) {
        let size = (MAX_PROMOTION_EVENTS as u64) * std::mem::size_of::<PromotionEvent>() as u64;
        encoder.copy_buffer_to_buffer(&self.promotion_buffer, 0, &self.promotion_readback, 0, size);
        encoder.copy_buffer_to_buffer(&self.promotion_count_buffer, 0, &self.count_readback, 0, 4);
    }

    /// Get cached promotion events from the last readback.
    pub fn promotion_events(&self) -> &[PromotionEvent] {
        &self.cached_promotions
    }

    /// Current particle buffer capacity.
    pub fn capacity(&self) -> u32 {
        self.capacity
    }

    /// Access the particle buffer for external use.
    pub fn particle_buffer(&self) -> Option<&wgpu::Buffer> {
        self.particle_buffer.as_ref()
    }

    /// Explicitly destroy GPU resources.
    pub fn destroy(&mut self) {
        if let Some(buf) = self.particle_buffer.take() {
            buf.destroy();
        }
        self.params_buffer.destroy();
        self.promotion_buffer.destroy();
        self.promotion_count_buffer.destroy();
        self.promotion_readback.destroy();
        self.count_readback.destroy();
    }
}

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

    #[test]
    fn dreammatter_params_size() {
        assert_eq!(std::mem::size_of::<DreamMatterParams>(), 48);
    }

    #[test]
    fn promotion_event_size() {
        assert_eq!(std::mem::size_of::<PromotionEvent>(), 32);
    }

    #[test]
    fn promotion_event_zeroed() {
        let e = PromotionEvent::zeroed();
        assert_eq!(e.kind, 0);
        assert_eq!(e.position, [0.0; 3]);
    }

    #[test]
    fn max_promotion_events_bounded() {
        assert_eq!(MAX_PROMOTION_EVENTS, 256);
    }
}
