//! Billboard particle renderer.

use super::super::buffers::GpuPhysicsBuffers;
use crate::camera::Camera;

#[repr(C)]
#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
struct RenderUniforms {
    view_proj: [f32; 16],
    camera_right: [f32; 3],
    _pad0: f32,
    camera_up: [f32; 3],
    _pad1: f32,
    camera_pos: [f32; 3],
    _pad2: f32,
}

pub struct BillboardRenderer {
    pipeline: wgpu::RenderPipeline,
    bind_group_layout: wgpu::BindGroupLayout,
    bind_group: Option<wgpu::BindGroup>,
    uniform_buffer: wgpu::Buffer,
    needs_rebind: bool,
}

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

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

        let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("physics_billboard_layout"),
            entries: &[
                wgpu::BindGroupLayoutEntry {
                    binding: 0,
                    visibility: wgpu::ShaderStages::VERTEX,
                    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::VERTEX,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Storage { read_only: true },
                        has_dynamic_offset: false,
                        min_binding_size: None,
                    },
                    count: None,
                },
            ],
        });

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

        let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: Some("physics_billboard_pipeline"),
            layout: Some(&pipeline_layout),
            vertex: wgpu::VertexState {
                module: &module,
                entry_point: Some("vs_main"),
                buffers: &[],
                compilation_options: Default::default(),
            },
            fragment: Some(wgpu::FragmentState {
                module: &module,
                entry_point: Some("fs_main"),
                targets: &[Some(wgpu::ColorTargetState {
                    format: crate::formats::TEXTURE_FORMAT,
                    blend: Some(wgpu::BlendState {
                        color: wgpu::BlendComponent {
                            src_factor: wgpu::BlendFactor::SrcAlpha,
                            dst_factor: wgpu::BlendFactor::OneMinusSrcAlpha,
                            operation: wgpu::BlendOperation::Add,
                        },
                        alpha: wgpu::BlendComponent {
                            src_factor: wgpu::BlendFactor::One,
                            dst_factor: wgpu::BlendFactor::OneMinusSrcAlpha,
                            operation: wgpu::BlendOperation::Add,
                        },
                    }),
                    write_mask: wgpu::ColorWrites::ALL,
                })],
                compilation_options: Default::default(),
            }),
            primitive: wgpu::PrimitiveState {
                topology: wgpu::PrimitiveTopology::TriangleList,
                ..Default::default()
            },
            depth_stencil: Some(wgpu::DepthStencilState {
                format: crate::formats::DEPTH_FORMAT,
                depth_write_enabled: Some(false),
                depth_compare: Some(wgpu::CompareFunction::Less),
                stencil: wgpu::StencilState::default(),
                bias: wgpu::DepthBiasState::default(),
            }),
            multisample: wgpu::MultisampleState::default(),
            multiview_mask: None,
            cache: None,
        });

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

        Self {
            pipeline,
            bind_group_layout,
            bind_group: None,
            uniform_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_billboard_bind_group"),
            layout: &self.bind_group_layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: self.uniform_buffer.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: buffers.particle_state.as_entire_binding(),
                },
            ],
        }));
        self.needs_rebind = false;
    }

    /// Prepare rendering state. Call before the render pass begins.
    pub fn prepare(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        buffers: &GpuPhysicsBuffers,
        camera: &Camera,
    ) {
        let view = camera.view_matrix;
        let proj = camera.projection_matrix;
        let vp = proj * view;

        let right = [view.col(0).x, view.col(0).y, view.col(0).z];
        let up = [view.col(1).x, view.col(1).y, view.col(1).z];
        let pos = [camera.position.x, camera.position.y, camera.position.z];

        let uniforms = RenderUniforms {
            view_proj: vp.to_cols_array(),
            camera_right: right,
            _pad0: 0.0,
            camera_up: up,
            _pad1: 0.0,
            camera_pos: pos,
            _pad2: 0.0,
        };

        queue.write_buffer(&self.uniform_buffer, 0, bytemuck::bytes_of(&uniforms));
        self.ensure_bind_group(device, buffers);
    }

    /// Render particles. Call after prepare, inside a render pass.
    pub fn render<'a>(&'a self, render_pass: &mut wgpu::RenderPass<'a>, max_particles: u32) {
        if let Some(bg) = &self.bind_group {
            render_pass.set_pipeline(&self.pipeline);
            render_pass.set_bind_group(0, bg, &[]);
            render_pass.draw(0..6, 0..max_particles);
        }
    }
}
