use crate::sim::sim_manager::SimManager;

#[repr(C)]
#[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
struct RendererParticle {
    position: [f32; 2],
    data: u32,
    life: f32,
}

fn create_particle_buffer(device: &wgpu::Device, capacity: usize) -> wgpu::Buffer {
    device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("Particle buffer"),
        size: (capacity * size_of::<RendererParticle>()) as u64,
        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
        mapped_at_creation: false,
    })
}

pub struct ParticlesRendererState {
    pub particle_pipeline: wgpu::RenderPipeline,
    pub particle_buffer: wgpu::Buffer,
    pub particle_bind_group: wgpu::BindGroup,
    pub particle_capacity: usize,

    particles_count: usize,
}

impl ParticlesRendererState {
    pub fn update_buffers(&mut self, device: &wgpu::Device, queue: &wgpu::Queue, sim: &SimManager) {
        puffin::profile_function!();
        // TODO could cull this to visible, maybe it's slower?
        let particles: Vec<RendererParticle> = sim
            .particle_manager
            .particles
            .iter()
            .map(|p| RendererParticle {
                position: [p.position.x, p.position.y],
                data: p.material as u32,
                life: p.life,
            })
            .collect();

        if particles.len() > self.particle_capacity {
            self.particle_capacity = particles.len().next_power_of_two();
            self.particle_buffer = create_particle_buffer(&device, self.particle_capacity);
        }

        queue.write_buffer(&self.particle_buffer, 0, bytemuck::cast_slice(&particles));

        self.particles_count = particles.len();
    }

    pub fn render(&self, render_pass: &mut wgpu::RenderPass) {
        render_pass.set_pipeline(&self.particle_pipeline);
        render_pass.set_bind_group(0, &self.particle_bind_group, &[]);
        render_pass.draw(0..4, 0..self.particles_count as u32);
    }

    pub fn new(
        device: &wgpu::Device,
        surface_config: &wgpu::SurfaceConfiguration,
        camera_uniform_buffer: &wgpu::Buffer,
        palette_buffer: &wgpu::Buffer,
    ) -> Self {
        let particle_bind_group_layout =
            device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                label: None,
                entries: &[
                    // camera uniform
                    wgpu::BindGroupLayoutEntry {
                        ty: wgpu::BindingType::Buffer {
                            ty: wgpu::BufferBindingType::Uniform,
                            has_dynamic_offset: false,
                            min_binding_size: None,
                        },
                        binding: 0,
                        count: None,
                        visibility: wgpu::ShaderStages::VERTEX,
                    },
                    // palette
                    wgpu::BindGroupLayoutEntry {
                        ty: wgpu::BindingType::Buffer {
                            ty: wgpu::BufferBindingType::Storage { read_only: true },
                            has_dynamic_offset: false,
                            min_binding_size: None,
                        },
                        binding: 1,
                        count: None,
                        visibility: wgpu::ShaderStages::FRAGMENT,
                    },
                    // instance buffer
                    wgpu::BindGroupLayoutEntry {
                        ty: wgpu::BindingType::Buffer {
                            ty: wgpu::BufferBindingType::Storage { read_only: true },
                            has_dynamic_offset: false,
                            min_binding_size: None,
                        },
                        binding: 2,
                        count: None,
                        visibility: wgpu::ShaderStages::VERTEX | wgpu::ShaderStages::FRAGMENT,
                    },
                ],
            });

        let particle_shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("Particle shader"),
            source: wgpu::ShaderSource::Wgsl(include_str!("../shader/particle.wgsl").into()),
        });

        let particle_pipeline_layout =
            device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                label: None,
                immediate_size: 0,
                bind_group_layouts: &[Some(&particle_bind_group_layout)],
            });

        let particle_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: None,
            layout: Some(&particle_pipeline_layout),
            vertex: wgpu::VertexState {
                module: &particle_shader,
                entry_point: Some("vs_main"),
                buffers: &[],
                compilation_options: wgpu::PipelineCompilationOptions::default(),
            },
            fragment: Some(wgpu::FragmentState {
                module: &particle_shader,
                entry_point: Some("fs_main"),
                targets: &[Some(wgpu::ColorTargetState {
                    format: surface_config.format,
                    blend: Some(wgpu::BlendState::ALPHA_BLENDING),
                    write_mask: wgpu::ColorWrites::ALL,
                })],
                compilation_options: wgpu::PipelineCompilationOptions::default(),
            }),
            primitive: wgpu::PrimitiveState {
                topology: wgpu::PrimitiveTopology::TriangleStrip,
                strip_index_format: None,
                front_face: wgpu::FrontFace::Ccw,
                cull_mode: None,
                polygon_mode: wgpu::PolygonMode::Fill,
                unclipped_depth: false,
                conservative: false,
            },
            depth_stencil: None,
            multisample: wgpu::MultisampleState::default(),
            multiview_mask: None,
            cache: None,
        });

        let particle_capacity = 2048;
        let particle_buffer = create_particle_buffer(&device, particle_capacity);

        let particle_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &particle_bind_group_layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: camera_uniform_buffer.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: palette_buffer.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 2,
                    resource: particle_buffer.as_entire_binding(),
                },
            ],
        });

        ParticlesRendererState {
            particle_pipeline,
            particle_buffer,
            particle_bind_group,
            particle_capacity,

            particles_count: 0,
        }
    }
}