tirbofish/dropbear

use crate::asset::{ASSET_REGISTRY, Handle};
use crate::attenuation::{Attenuation, RANGE_50};
use crate::buffer::{DynamicBuffer, UniformBuffer, WritableBuffer};
use crate::graphics::SharedGraphicsContext;
use crate::pipelines::light_cube::InstanceInput;
use crate::procedural::ProcedurallyGeneratedObject;
use crate::{entity::Transform, model::Model};
use dropbear_utils::Dirty;
use glam::{DMat4, DQuat, DVec3};
use std::fmt::{Display, Formatter};
use std::sync::Arc;

const LIGHT_FORWARD_AXIS: DVec3 = DVec3::new(0.0, -1.0, 0.0);

pub const MAX_LIGHTS: usize = 10;

#[repr(C)]
#[derive(Debug, Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
// ENSURE THAT THE SIZE OF THE UNIFORM IS OF A MULTIPLE OF 16. USE `size_of::<LightUniform>()`
pub struct LightUniform {
    pub position: [f32; 4],
    pub direction: [f32; 4], // outer cutoff is .w value
    pub colour: [f32; 4],    // last value is the light type
    // pub light_type: u32,
    pub constant: f32,
    pub linear: f32,
    pub quadratic: f32,
    pub cutoff: f32,
}

fn dvec3_to_uniform_array(vec: DVec3) -> [f32; 4] {
    [vec.x as f32, vec.y as f32, vec.z as f32, 1.0]
}

fn dvec3_colour_to_uniform_array(vec: DVec3, light_type: LightType) -> [f32; 4] {
    [
        vec.x as f32,
        vec.y as f32,
        vec.z as f32,
        light_type as u32 as f32,
    ]
}

fn dvec3_direction_to_uniform_array(vec: DVec3, outer_cutoff_angle: f32) -> [f32; 4] {
    [
        vec.x as f32,
        vec.y as f32,
        vec.z as f32,
        f32::cos(outer_cutoff_angle.to_radians()),
    ]
}

impl Default for LightUniform {
    fn default() -> Self {
        Self {
            position: [0.0, 0.0, 0.0, 1.0],
            direction: [0.0, -1.0, 0.0, 1.0],
            colour: [1.0, 1.0, 1.0, 1.0],
            // light_type: 0,
            constant: 0.0,
            linear: 0.0,
            quadratic: 0.0,
            cutoff: f32::cos(12.5_f32.to_radians()),
        }
    }
}

#[repr(C)]
#[derive(Debug, Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
pub struct LightArrayUniform {
    pub lights: [LightUniform; MAX_LIGHTS],
    pub light_count: u32,
    pub ambient_strength: f32,
    pub _padding: [u32; 2],
}

impl Default for LightArrayUniform {
    fn default() -> Self {
        Self {
            lights: [LightUniform::default(); MAX_LIGHTS],
            light_count: 0,
            ambient_strength: 0.5,
            _padding: [0; 2],
        }
    }
}

#[derive(
    Default, Debug, Clone, Copy, serde::Serialize, serde::Deserialize, PartialEq, Eq, Hash,
)]
pub enum LightType {
    #[default]
    // Example: Sunlight
    Directional = 0,
    // Example: Lamp
    Point = 1,
    // Example: Flashlight
    Spot = 2,
}

impl Display for LightType {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            LightType::Directional => write!(f, "Directional"),
            LightType::Point => write!(f, "Point"),
            LightType::Spot => write!(f, "Spot"),
        }
    }
}

impl From<LightType> for u32 {
    fn from(val: LightType) -> Self {
        match val {
            LightType::Directional => 0,
            LightType::Point => 1,
            LightType::Spot => 2,
        }
    }
}

#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct LightComponent {
    #[serde(default)]
    pub position: DVec3, // point, spot

    #[serde(default)]
    pub direction: DVec3, // directional, spot

    #[serde(default)]
    pub colour: DVec3, // all

    #[serde(default)]
    pub light_type: LightType, // all

    #[serde(default)]
    pub intensity: f32, // all

    #[serde(default)]
    pub attenuation: Attenuation, // point, spot

    #[serde(default)]
    pub enabled: bool, // all - light

    #[serde(default)]
    pub visible: bool, // all - cube

    #[serde(default)]
    pub cutoff_angle: f32, // spot

    #[serde(default)]
    pub outer_cutoff_angle: f32, // spot

    #[serde(default)]
    pub cast_shadows: bool,

    #[serde(default)]
    pub depth: std::ops::Range<f32>, // all
}

impl Default for LightComponent {
    fn default() -> Self {
        Self {
            position: DVec3::ZERO,
            direction: LIGHT_FORWARD_AXIS,
            colour: DVec3::ONE,
            light_type: LightType::Point,
            intensity: 1.0,
            attenuation: RANGE_50,
            enabled: true,
            cutoff_angle: 12.5,
            outer_cutoff_angle: 17.5,
            visible: true,
            cast_shadows: true,
            depth: 0.1..100.0,
        }
    }
}

impl LightComponent {
    pub fn default_direction() -> DVec3 {
        let dir = DVec3::new(-0.35, -1.0, -0.25);
        dir.normalize()
    }

    pub fn new(
        colour: DVec3,
        light_type: LightType,
        intensity: f32,
        attenuation: Option<Attenuation>,
    ) -> Self {
        let direction = match light_type {
            LightType::Directional | LightType::Spot => Self::default_direction(),
            LightType::Point => DVec3::ZERO,
        };

        Self {
            position: Default::default(),
            direction,
            colour,
            light_type,
            intensity,
            attenuation: attenuation.unwrap_or(RANGE_50),
            enabled: true,
            cutoff_angle: 12.5,
            outer_cutoff_angle: 17.5,
            cast_shadows: true,
            visible: true,
            depth: 0.1..100.0,
        }
    }

    pub fn to_transform(&self) -> Transform {
        Transform {
            position: self.position,
            rotation: Self::direction_to_quaternion(self.direction),
            scale: DVec3::ONE,
        }
    }

    /// Converts a direction vector to a quaternion rotation
    fn direction_to_quaternion(direction: DVec3) -> DQuat {
        if direction.length_squared() < 1e-6 {
            return DQuat::IDENTITY;
        }

        let forward = LIGHT_FORWARD_AXIS;
        let normalized_direction = direction.normalize();

        let dot = forward.dot(normalized_direction);

        if dot > 0.9999 {
            DQuat::IDENTITY
        } else if dot < -0.9999 {
            DQuat::from_axis_angle(DVec3::Y, std::f64::consts::PI)
        } else {
            let axis = forward.cross(normalized_direction).normalize();
            let angle = dot.acos();
            DQuat::from_axis_angle(axis, angle)
        }
    }

    pub fn directional(colour: DVec3, intensity: f32) -> Self {
        Self::new(colour, LightType::Directional, intensity, None)
    }

    pub fn point(colour: DVec3, intensity: f32, attenuation: Attenuation) -> Self {
        Self::new(colour, LightType::Point, intensity, Some(attenuation))
    }

    pub fn spot(colour: DVec3, intensity: f32) -> Self {
        Self::new(colour, LightType::Spot, intensity, None)
    }

    pub fn hide_cube(&mut self) {
        self.visible = false;
    }

    pub fn show_cube(&mut self) {
        self.visible = true;
    }

    pub fn disable_light(&mut self) {
        self.enabled = false;
    }

    pub fn enable_light(&mut self) {
        self.enabled = true;
    }
}

#[derive(Clone)]
pub struct Light {
    pub uniform: Dirty<LightUniform>,
    pub cube_model: Handle<Model>,
    pub label: String,
    pub buffer: UniformBuffer<LightUniform>,
    pub bind_group: wgpu::BindGroup,
    pub instance_buffer: DynamicBuffer<InstanceInput>,
    pub component: LightComponent,
}

impl Light {
    pub async fn new(
        graphics: Arc<SharedGraphicsContext>,
        light: LightComponent,
        label: Option<&str>,
    ) -> Self {
        puffin::profile_function!();

        let uniform = Dirty::new(LightUniform {
            position: dvec3_to_uniform_array(light.position),
            direction: dvec3_direction_to_uniform_array(light.direction, light.outer_cutoff_angle),
            colour: dvec3_colour_to_uniform_array(
                light.colour * light.intensity as f64,
                light.light_type,
            ),
            constant: light.attenuation.constant,
            linear: light.attenuation.linear,
            quadratic: light.attenuation.quadratic,
            cutoff: f32::cos(light.cutoff_angle.to_radians()),
        });

        log::trace!("Created new light uniform");

        let cube_model = ProcedurallyGeneratedObject::cuboid(DVec3::ONE).build_model(
            graphics.clone(),
            None,
            Some("light cube"),
            ASSET_REGISTRY.clone(),
        );

        let label_str = label.unwrap_or("Light").to_string();

        let buffer = UniformBuffer::new(&graphics.device, &label_str);
        let bind_group = graphics
            .device
            .create_bind_group(&wgpu::BindGroupDescriptor {
                label: Some(&format!("{} light bind group", label_str)),
                layout: &graphics.layouts.light_cube_layout,
                entries: &[wgpu::BindGroupEntry {
                    binding: 0,
                    resource: buffer.buffer().as_entire_binding(),
                }],
            });

        let transform = light.to_transform();
        let instance: InstanceInput = DMat4::from_scale_rotation_translation(
            transform.scale,
            transform.rotation,
            transform.position,
        )
        .into();

        let mut instance_buffer = DynamicBuffer::new(
            &graphics.device,
            1,
            wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
            "Light Instance Buffer",
        );

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

        log::debug!("Created new light [{}]", label_str);

        Self {
            uniform,
            cube_model,
            label: label_str,
            buffer,
            bind_group,
            instance_buffer,
            component: light.clone(),
        }
    }

    pub fn update(&mut self, graphics: &SharedGraphicsContext) {
        puffin::profile_function!();

        let light = &mut self.component;

        self.uniform.position = dvec3_to_uniform_array(light.position);

        self.uniform.direction =
            dvec3_direction_to_uniform_array(light.direction, light.outer_cutoff_angle);

        self.uniform.colour =
            dvec3_colour_to_uniform_array(light.colour * light.intensity as f64, light.light_type);
        self.uniform.constant = light.attenuation.constant;
        self.uniform.linear = light.attenuation.linear;
        self.uniform.quadratic = light.attenuation.quadratic;

        self.uniform.cutoff = f32::cos(light.cutoff_angle.to_radians());

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

    pub fn uniform(&self) -> &Dirty<LightUniform> {
        &self.uniform
    }

    pub fn mark_clean(&mut self) {
        self.uniform.mark_clean();
    }

    pub fn model(&self) -> Handle<Model> {
        self.cube_model
    }

    pub fn label(&self) -> &str {
        &self.label
    }
}