//! [`ForwardRenderer`]: a single-pass forward renderer with a depth buffer and
//! one directional light.
//!
//! Stage 4 scope: draw a list of [`RenderObject`]s (each a [`GpuMesh`] +
//! [`Material`] + [`Transform`]) through the lit shader, into a caller-provided
//! color target, using an owned depth texture. Shadows, multiple lights, and
//! post-processing arrive in later stages.

use std::num::NonZeroU64;

use bytemuck::{Pod, Zeroable};
use glam::Mat3;
use serde::{Deserialize, Serialize};

use super::mesh::{GpuMesh, Vertex};
use super::{Camera, Material};
use crate::math::{Color, Transform, Vec3, Vec4};

/// Depth buffer format used by the forward pass.
pub const DEPTH_FORMAT: wgpu::TextureFormat = wgpu::TextureFormat::Depth32Float;

/// A directional light: parallel rays with a travel `direction`.
///
/// Also a **reflected, addable component**: drop one on an entity to author a
/// sun/key light in the scene, dual-editable from the editor and scripts.
/// (Gathering light entities into the forward pass is a later-stage wiring; the
/// renderer currently takes its [`Lighting`] directly.)
#[derive(Debug, Clone, Copy, Serialize, Deserialize, crate::reflect::Reflect)]
pub struct DirectionalLight {
    /// The direction the light travels (does not need to be normalized).
    pub direction: Vec3,
    /// Light color.
    pub color: Color,
    /// Scalar intensity multiplier.
    pub intensity: f32,
}

impl Default for DirectionalLight {
    fn default() -> Self {
        Self {
            direction: Vec3::new(-0.5, -1.0, -0.35),
            color: Color::WHITE,
            intensity: 1.0,
        }
    }
}

/// Scene lighting for a forward pass: one directional light plus an ambient term.
#[derive(Debug, Clone, Copy)]
pub struct Lighting {
    /// The single directional (sun) light.
    pub light: DirectionalLight,
    /// Flat ambient color added everywhere (cheap fill light).
    pub ambient: Color,
}

impl Default for Lighting {
    fn default() -> Self {
        Self {
            light: DirectionalLight::default(),
            ambient: Color::rgb(0.08, 0.08, 0.10),
        }
    }
}

/// One drawable: a GPU mesh placed by `transform` and shaded with `material`.
pub struct RenderObject<'a> {
    /// The mesh to draw.
    pub mesh: &'a GpuMesh,
    /// Its surface material.
    pub material: Material,
    /// World placement.
    pub transform: Transform,
}

#[repr(C)]
#[derive(Clone, Copy, Pod, Zeroable)]
struct GlobalsUniform {
    view_proj: [[f32; 4]; 4],
    camera_pos: [f32; 4],
    light_dir: [f32; 4],
    light_color: [f32; 4],
    ambient: [f32; 4],
}

#[repr(C)]
#[derive(Clone, Copy, Pod, Zeroable)]
struct ObjectUniform {
    model: [[f32; 4]; 4],
    normal_mtx: [[f32; 4]; 4],
    albedo: [f32; 4],
    mr: [f32; 4],
}

/// A forward renderer owning its pipeline, depth buffer, and uniform storage.
pub struct ForwardRenderer {
    pipeline: wgpu::RenderPipeline,
    globals_buffer: wgpu::Buffer,
    globals_bind_group: wgpu::BindGroup,
    object_layout: wgpu::BindGroupLayout,
    object_buffer: wgpu::Buffer,
    object_bind_group: wgpu::BindGroup,
    /// Per-object stride: `size_of::<ObjectUniform>` rounded up to the device's
    /// minimum dynamic-uniform-buffer offset alignment.
    object_stride: u64,
    object_capacity: u32,
    depth: Option<DepthTarget>,
    color_format: wgpu::TextureFormat,
}

struct DepthTarget {
    view: wgpu::TextureView,
    width: u32,
    height: u32,
}

impl ForwardRenderer {
    /// Builds the renderer for a given color target format (e.g. the surface
    /// format for a window, or `Rgba8Unorm` for offscreen rendering).
    pub fn new(device: &wgpu::Device, color_format: wgpu::TextureFormat) -> Self {
        let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("oxide.forward.lit"),
            source: wgpu::ShaderSource::Wgsl(include_str!("shaders/lit.wgsl").into()),
        });

        let globals_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("oxide.forward.globals_layout"),
            entries: &[wgpu::BindGroupLayoutEntry {
                binding: 0,
                visibility: wgpu::ShaderStages::VERTEX_FRAGMENT,
                ty: wgpu::BindingType::Buffer {
                    ty: wgpu::BufferBindingType::Uniform,
                    has_dynamic_offset: false,
                    min_binding_size: NonZeroU64::new(std::mem::size_of::<GlobalsUniform>() as u64),
                },
                count: None,
            }],
        });

        let object_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("oxide.forward.object_layout"),
            entries: &[wgpu::BindGroupLayoutEntry {
                binding: 0,
                visibility: wgpu::ShaderStages::VERTEX_FRAGMENT,
                ty: wgpu::BindingType::Buffer {
                    ty: wgpu::BufferBindingType::Uniform,
                    has_dynamic_offset: true,
                    min_binding_size: NonZeroU64::new(std::mem::size_of::<ObjectUniform>() as u64),
                },
                count: None,
            }],
        });

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

        let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: Some("oxide.forward.pipeline"),
            layout: Some(&pipeline_layout),
            vertex: wgpu::VertexState {
                module: &shader,
                entry_point: Some("vs_main"),
                compilation_options: Default::default(),
                buffers: &[Vertex::LAYOUT],
            },
            primitive: wgpu::PrimitiveState {
                topology: wgpu::PrimitiveTopology::TriangleList,
                strip_index_format: None,
                front_face: wgpu::FrontFace::Ccw,
                cull_mode: Some(wgpu::Face::Back),
                unclipped_depth: false,
                polygon_mode: wgpu::PolygonMode::Fill,
                conservative: false,
            },
            depth_stencil: Some(wgpu::DepthStencilState {
                format: DEPTH_FORMAT,
                depth_write_enabled: Some(true),
                depth_compare: Some(wgpu::CompareFunction::Less),
                stencil: wgpu::StencilState::default(),
                bias: wgpu::DepthBiasState::default(),
            }),
            multisample: wgpu::MultisampleState::default(),
            fragment: Some(wgpu::FragmentState {
                module: &shader,
                entry_point: Some("fs_main"),
                compilation_options: Default::default(),
                targets: &[Some(wgpu::ColorTargetState {
                    format: color_format,
                    blend: Some(wgpu::BlendState::ALPHA_BLENDING),
                    write_mask: wgpu::ColorWrites::ALL,
                })],
            }),
            multiview_mask: None,
            cache: None,
        });

        let globals_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("oxide.forward.globals"),
            size: std::mem::size_of::<GlobalsUniform>() as u64,
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        let globals_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("oxide.forward.globals_bg"),
            layout: &globals_layout,
            entries: &[wgpu::BindGroupEntry {
                binding: 0,
                resource: globals_buffer.as_entire_binding(),
            }],
        });

        let object_stride = align_up(
            std::mem::size_of::<ObjectUniform>() as u64,
            device.limits().min_uniform_buffer_offset_alignment as u64,
        );
        let object_capacity = 16;
        let (object_buffer, object_bind_group) =
            create_object_storage(device, &object_layout, object_stride, object_capacity);

        Self {
            pipeline,
            globals_buffer,
            globals_bind_group,
            object_layout,
            object_buffer,
            object_bind_group,
            object_stride,
            object_capacity,
            depth: None,
            color_format,
        }
    }

    /// The color target format this renderer was built for.
    pub fn color_format(&self) -> wgpu::TextureFormat {
        self.color_format
    }

    /// Renders `objects` into `target` (whose full physical size is
    /// `width`×`height`) as seen by `camera` placed at `view_transform`, lit
    /// by `lighting`. Drawing is restricted to `viewport_rect` (a sub-
    /// rectangle of the target), and the projection uses that rect's aspect
    /// ratio.
    ///
    /// The color target is *loaded* (not cleared) so a clear pass run before
    /// this — e.g. the window's clear color — shows through as the background;
    /// the depth buffer is cleared to 1.0 each call.
    #[allow(clippy::too_many_arguments)]
    pub fn render(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        target: &wgpu::TextureView,
        (width, height): (u32, u32),
        viewport_rect: crate::math::Rect,
        camera: &Camera,
        view_transform: &Transform,
        lighting: &Lighting,
        objects: &[RenderObject<'_>],
    ) {
        let (width, height) = (width.max(1), height.max(1));
        // Clamp the viewport rect to the target so wgpu doesn't complain.
        let vp_w = viewport_rect.width().max(1.0).min(width as f32);
        let vp_h = viewport_rect.height().max(1.0).min(height as f32);
        let vp_x = viewport_rect.min.x.max(0.0).min(width as f32 - vp_w);
        let vp_y = viewport_rect.min.y.max(0.0).min(height as f32 - vp_h);

        // Depth must match the full color target's dimensions (the
        // attachment binding requires that). Pixels outside `set_viewport`
        // are never written, so the extra depth is wasted memory but never
        // incorrect.
        self.ensure_depth(device, width, height);
        self.ensure_object_capacity(device, objects.len() as u32);

        // Globals — aspect comes from the viewport rect, not the target.
        let aspect = vp_w / vp_h;
        let view_proj = camera.view_projection(aspect, view_transform);
        let to_light = (-lighting.light.direction).normalize_or_zero();
        let lc = lighting.light.color;
        let amb = lighting.ambient;
        let globals = GlobalsUniform {
            view_proj: view_proj.to_cols_array_2d(),
            camera_pos: view_transform.translation.extend(1.0).to_array(),
            light_dir: to_light.extend(0.0).to_array(),
            light_color: (Vec4::new(lc.r, lc.g, lc.b, 1.0) * lighting.light.intensity).to_array(),
            ambient: Vec4::new(amb.r, amb.g, amb.b, 1.0).to_array(),
        };
        queue.write_buffer(&self.globals_buffer, 0, bytemuck::bytes_of(&globals));

        // Per-object uniforms.
        for (i, obj) in objects.iter().enumerate() {
            let model = obj.transform.to_matrix();
            let normal_mtx = Mat3::from_mat4(model).inverse().transpose();
            let normal_mtx4 = [
                normal_mtx.x_axis.extend(0.0).to_array(),
                normal_mtx.y_axis.extend(0.0).to_array(),
                normal_mtx.z_axis.extend(0.0).to_array(),
                [0.0, 0.0, 0.0, 1.0],
            ];
            let a = obj.material.albedo;
            let uniform = ObjectUniform {
                model: model.to_cols_array_2d(),
                normal_mtx: normal_mtx4,
                albedo: [a.r, a.g, a.b, a.a],
                mr: [obj.material.metallic, obj.material.roughness, 0.0, 0.0],
            };
            queue.write_buffer(
                &self.object_buffer,
                i as u64 * self.object_stride,
                bytemuck::bytes_of(&uniform),
            );
        }

        let depth_view = &self.depth.as_ref().expect("depth ensured above").view;
        let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("oxide.forward.encoder"),
        });
        {
            let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
                label: Some("oxide.forward.pass"),
                color_attachments: &[Some(wgpu::RenderPassColorAttachment {
                    view: target,
                    resolve_target: None,
                    depth_slice: None,
                    ops: wgpu::Operations {
                        load: wgpu::LoadOp::Load,
                        store: wgpu::StoreOp::Store,
                    },
                })],
                depth_stencil_attachment: Some(wgpu::RenderPassDepthStencilAttachment {
                    view: depth_view,
                    depth_ops: Some(wgpu::Operations {
                        load: wgpu::LoadOp::Clear(1.0),
                        store: wgpu::StoreOp::Store,
                    }),
                    stencil_ops: None,
                }),
                timestamp_writes: None,
                occlusion_query_set: None,
                multiview_mask: None,
            });
            // Restrict drawing to the host's viewport sub-rect. Pixels
            // outside this rectangle keep whatever the prior pass (e.g.
            // ClearPass or the window clear) wrote there.
            pass.set_viewport(vp_x, vp_y, vp_w, vp_h, 0.0, 1.0);
            pass.set_pipeline(&self.pipeline);
            pass.set_bind_group(0, &self.globals_bind_group, &[]);
            for (i, obj) in objects.iter().enumerate() {
                let offset = (i as u64 * self.object_stride) as u32;
                pass.set_bind_group(1, &self.object_bind_group, &[offset]);
                pass.set_vertex_buffer(0, obj.mesh.vertex_buffer.slice(..));
                pass.set_index_buffer(obj.mesh.index_buffer.slice(..), wgpu::IndexFormat::Uint32);
                pass.draw_indexed(0..obj.mesh.index_count, 0, 0..1);
            }
        }
        queue.submit([encoder.finish()]);
    }

    fn ensure_depth(&mut self, device: &wgpu::Device, width: u32, height: u32) {
        let stale = match &self.depth {
            Some(d) => d.width != width || d.height != height,
            None => true,
        };
        if stale {
            let texture = device.create_texture(&wgpu::TextureDescriptor {
                label: Some("oxide.forward.depth"),
                size: wgpu::Extent3d {
                    width,
                    height,
                    depth_or_array_layers: 1,
                },
                mip_level_count: 1,
                sample_count: 1,
                dimension: wgpu::TextureDimension::D2,
                format: DEPTH_FORMAT,
                usage: wgpu::TextureUsages::RENDER_ATTACHMENT,
                view_formats: &[],
            });
            self.depth = Some(DepthTarget {
                view: texture.create_view(&wgpu::TextureViewDescriptor::default()),
                width,
                height,
            });
        }
    }

    fn ensure_object_capacity(&mut self, device: &wgpu::Device, needed: u32) {
        if needed > self.object_capacity {
            let capacity = needed.next_power_of_two();
            let (buffer, bind_group) =
                create_object_storage(device, &self.object_layout, self.object_stride, capacity);
            self.object_buffer = buffer;
            self.object_bind_group = bind_group;
            self.object_capacity = capacity;
        }
    }
}

/// Allocates the per-object uniform buffer (`capacity` slots of `stride` bytes)
/// and a dynamic-offset bind group over it.
fn create_object_storage(
    device: &wgpu::Device,
    layout: &wgpu::BindGroupLayout,
    stride: u64,
    capacity: u32,
) -> (wgpu::Buffer, wgpu::BindGroup) {
    let buffer = device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("oxide.forward.objects"),
        size: stride * capacity as u64,
        usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
        mapped_at_creation: false,
    });
    let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("oxide.forward.object_bg"),
        layout,
        entries: &[wgpu::BindGroupEntry {
            binding: 0,
            resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
                buffer: &buffer,
                offset: 0,
                size: NonZeroU64::new(std::mem::size_of::<ObjectUniform>() as u64),
            }),
        }],
    });
    (buffer, bind_group)
}

/// Rounds `value` up to the next multiple of `align` (a power of two).
fn align_up(value: u64, align: u64) -> u64 {
    let align = align.max(1);
    value.div_ceil(align) * align
}