SolarEngine/solar_engine/src/light.rs

use bytemuck::{Pod, Zeroable};
use cgmath::Vector3;

#[repr(u32)]
#[derive(Clone, Copy, Debug)]
pub enum LightType {
    Directional = 0,
    Point = 1,
    Spot = 2
}

#[repr(C)]
#[derive(Clone, Copy, Pod, Zeroable, Debug)]
pub struct GpuLight {
    pub position: [f32; 3],
    pub light_type: u32,
    pub color: [f32; 3],
    pub intensity: f32,
    pub direction: [f32; 3],
    pub range: f32,
    pub inner_cutoff: f32,
    pub outer_cutoff: f32,
}

pub struct Light {
    pub light_type: LightType,
    pub position: Vector3<f32>,
    pub direction: Vector3<f32>,
    pub color: Vector3<f32>,
    pub intensity: f32,
    pub range: f32,
    pub inner_cutoff: f32,
    pub outer_cutoff: f32,
}

impl Light {
    pub fn to_gpu(&self) -> GpuLight {
        GpuLight {
            position: self.position.into(),
            light_type: self.light_type as u32,
            color: self.color.into(),
            intensity: self.intensity,
            direction: self.direction.into(),
            range: self.range,
            inner_cutoff: self.inner_cutoff,
            outer_cutoff: self.outer_cutoff,
        }
    }
}

pub struct LightManager {
    pub lights: Vec<Light>,
    pub buffer: wgpu::Buffer,
    pub bind_group: wgpu::BindGroup,
    pub count_buffer: wgpu::Buffer,
    pub layout: wgpu::BindGroupLayout
}

#[repr(C)]
#[derive(Clone, Copy, Pod, Zeroable, Debug)]
pub struct LightCount {
    pub count: u32,
}

impl LightManager {
    pub fn new(device: &wgpu::Device, max_lights: usize) -> Self {
        let buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("Light Buffer"),
            size: (max_lights * size_of::<GpuLight>()) as wgpu::BufferAddress,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        let layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("Light Bind Group Layout"),
            entries: &[
                // Binding 0: Storage buffer for lights
                wgpu::BindGroupLayoutEntry {
                    binding: 0,
                    visibility: wgpu::ShaderStages::VERTEX_FRAGMENT,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Storage { read_only: true },
                        has_dynamic_offset: false,
                        min_binding_size: None,
                    },
                    count: None,
                },
                // Binding 1: Uniform buffer for light count
                wgpu::BindGroupLayoutEntry {
                    binding: 1,
                    visibility: wgpu::ShaderStages::FRAGMENT,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Uniform,
                        has_dynamic_offset: false,
                        min_binding_size: None,
                    },
                    count: None,
                },
            ],
        });

        let count_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("Light Count Buffer"),
            size: size_of::<LightCount>() as wgpu::BufferAddress,
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("Light Bind Group"),
            layout: &layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: buffer.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: count_buffer.as_entire_binding(),
                },
            ],
        });

        Self {
            lights: Vec::new(),
            buffer,
            bind_group,
            layout,
            count_buffer,
        }
    }

    pub fn add_light(&mut self, light: Light) {
        self.lights.push(light);
    }

    pub fn update_gpu(&self, queue: &wgpu::Queue) {
        let data: Vec<GpuLight> = self.lights.iter().map(|l| l.to_gpu()).collect();
        queue.write_buffer(&self.buffer, 0, bytemuck::cast_slice(&data));

        let count = LightCount {
            count: self.lights.len() as u32,
        };
        queue.write_buffer(&self.count_buffer, 0, bytemuck::bytes_of(&count));
    }

    pub fn bind_group(&self) -> &wgpu::BindGroup {
        &self.bind_group
    }

    pub fn layout(&self) -> &wgpu::BindGroupLayout {
        &self.layout
    }
    
    pub fn clear(&mut self) {
        self.lights.clear();
    }
}