tirbofish/dropbear

use crate::buffer::{StorageBuffer, WritableBuffer};
use crate::graphics::SharedGraphicsContext;
use crate::lighting::{Light, LightArrayUniform, MAX_LIGHTS};
use crate::model::{ModelVertex, Vertex};
use crate::pipelines::DropbearShaderPipeline;
use crate::shader::Shader;
use crate::texture::Texture;
use glam::DMat4;
use slank::include_slang;
use std::mem::size_of;
use std::sync::Arc;
use wgpu::{BufferAddress, CompareFunction, DepthBiasState, StencilState};

pub struct LightCubePipeline {
    shader: Shader,
    pipeline_layout: wgpu::PipelineLayout,
    pipeline: wgpu::RenderPipeline,
    storage_buffer: Option<StorageBuffer<LightArrayUniform>>,
}

impl DropbearShaderPipeline for LightCubePipeline {
    fn new(graphics: Arc<SharedGraphicsContext>) -> Self {
        let shader = Shader::from_slang(
            graphics.clone(),
            &slank::CompiledSlangShader::from_bytes("light cube", include_slang!("light_cube")),
        );

        let pipeline_layout =
            graphics
                .device
                .create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                    label: Some("light cube pipeline layout"),
                    bind_group_layouts: &[
                        Some(&graphics.layouts.camera_bind_group_layout),
                        Some(&graphics.layouts.light_cube_layout),
                    ],
                    immediate_size: 0,
                });

        let hdr_format = graphics.hdr.read().format();
        let sample_count: u32 = (*graphics.antialiasing.read()).into();
        let pipeline = graphics
            .device
            .create_render_pipeline(&wgpu::RenderPipelineDescriptor {
                label: Some("light cube pipeline"),
                layout: Some(&pipeline_layout),
                vertex: wgpu::VertexState {
                    module: &shader,
                    entry_point: Some("vs_main"),
                    compilation_options: Default::default(),
                    buffers: &[
                        // model
                        LightCubeVertex::desc(),
                        // instance
                        InstanceInput::desc(),
                    ],
                },
                fragment: Some(wgpu::FragmentState {
                    module: &shader.module,
                    entry_point: Some("fs_main"),
                    targets: &[Some(wgpu::ColorTargetState {
                        format: hdr_format,
                        blend: Some(wgpu::BlendState {
                            alpha: wgpu::BlendComponent::REPLACE,
                            color: wgpu::BlendComponent::REPLACE,
                        }),
                        write_mask: wgpu::ColorWrites::ALL,
                    })],
                    compilation_options: Default::default(),
                }),
                primitive: wgpu::PrimitiveState {
                    topology: wgpu::PrimitiveTopology::TriangleList,
                    strip_index_format: None,
                    front_face: wgpu::FrontFace::Cw,
                    cull_mode: Some(wgpu::Face::Back),
                    polygon_mode: wgpu::PolygonMode::Fill,
                    unclipped_depth: false,
                    conservative: false,
                },
                depth_stencil: Some(wgpu::DepthStencilState {
                    format: Texture::DEPTH_FORMAT,
                    depth_write_enabled: Some(true),
                    depth_compare: Some(CompareFunction::Greater),
                    stencil: StencilState::default(),
                    bias: DepthBiasState::default(),
                }),
                multisample: wgpu::MultisampleState {
                    count: sample_count,
                    mask: !0,
                    alpha_to_coverage_enabled: false,
                },
                cache: None,
                multiview_mask: None,
            });

        let storage_buffer =
            StorageBuffer::new_read_only(&graphics.device, "light cube pipeline storage buffer");

        Self {
            shader,
            pipeline_layout,
            pipeline,
            storage_buffer: Some(storage_buffer),
        }
    }

    fn shader(&self) -> &Shader {
        &self.shader
    }

    fn pipeline_layout(&self) -> &wgpu::PipelineLayout {
        &self.pipeline_layout
    }

    fn pipeline(&self) -> &wgpu::RenderPipeline {
        &self.pipeline
    }
}

impl LightCubePipeline {
    pub fn light_buffer(&self) -> &wgpu::Buffer {
        self.storage_buffer
            .as_ref()
            .expect("Light cube storage buffer missing")
            .buffer()
    }

    pub fn update(&mut self, graphics: Arc<SharedGraphicsContext>, world: &hecs::World) {
        let mut light_array = LightArrayUniform::default();

        let mut light_index: usize = 0;

        for light in world.query::<&mut Light>().iter() {
            light.update(graphics.as_ref());


            let instance: InstanceInput = light.component.to_transform().matrix().into();

            light
                .instance_buffer
                .write(&graphics.device, &graphics.queue, &[instance]);

            if light.component.enabled && light_index < MAX_LIGHTS {
                let uniform = light.uniform();

                if uniform.is_dirty() {
                    light.buffer.write(&graphics.queue, &uniform);
                }

                light_array.lights[light_index] = *uniform.get();
                light_index += 1;
            }
        }

        light_array.light_count = light_index as u32;

        if let Some(buf) = &self.storage_buffer {
            buf.write(&graphics.queue, &light_array);
        } else {
            panic!("A storage buffer should have been created");
        }
    }

    pub fn buffer(&self) -> &wgpu::Buffer {
        if let Some(s) = &self.storage_buffer {
            s.buffer()
        } else {
            panic!("A storage buffer should have been created");
        }
    }
}

pub struct LightCubeVertex;

impl LightCubeVertex {
    fn desc() -> wgpu::VertexBufferLayout<'static> {
        wgpu::VertexBufferLayout {
            array_stride: size_of::<ModelVertex>() as BufferAddress,
            step_mode: wgpu::VertexStepMode::Vertex,
            attributes: &[wgpu::VertexAttribute {
                offset: 0,
                shader_location: 0,
                format: wgpu::VertexFormat::Float32x3,
            }],
        }
    }
}

/// As mapped in `shaders/light.slang` as
/// ```wgsl
/// struct InstanceInput {
///     @location(5) model_matrix_0: vec4<f32>,
///     @location(6) model_matrix_1: vec4<f32>,
///     @location(7) model_matrix_2: vec4<f32>,
///     @location(8) model_matrix_3: vec4<f32>,
/// }
/// ```
#[repr(C)]
#[derive(Debug, Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
pub struct InstanceInput {
    pub model_matrix: [[f32; 4]; 4],
}

impl Vertex for InstanceInput {
    fn desc() -> wgpu::VertexBufferLayout<'static> {
        wgpu::VertexBufferLayout {
            array_stride: size_of::<InstanceInput>() as BufferAddress,
            step_mode: wgpu::VertexStepMode::Instance,
            attributes: &[
                // model_matrix_0
                wgpu::VertexAttribute {
                    offset: 0,
                    shader_location: 5,
                    format: wgpu::VertexFormat::Float32x4,
                },
                // model_matrix_1
                wgpu::VertexAttribute {
                    offset: size_of::<[f32; 4]>() as wgpu::BufferAddress,
                    shader_location: 6,
                    format: wgpu::VertexFormat::Float32x4,
                },
                // model_matrix_2
                wgpu::VertexAttribute {
                    offset: size_of::<[f32; 8]>() as wgpu::BufferAddress,
                    shader_location: 7,
                    format: wgpu::VertexFormat::Float32x4,
                },
                // model_matrix_3
                wgpu::VertexAttribute {
                    offset: size_of::<[f32; 12]>() as wgpu::BufferAddress,
                    shader_location: 8,
                    format: wgpu::VertexFormat::Float32x4,
                },
            ],
        }
    }
}

impl Into<InstanceInput> for DMat4 {
    fn into(self) -> InstanceInput {
        InstanceInput {
            model_matrix: self.as_mat4().to_cols_array_2d(),
        }
    }
}