codecraft 0.1.2

//! Drawing loaded meshes: a camera, a mesh library, and one instanced draw
//! call per mesh.
use std::collections::HashMap;

use bytemuck::{Pod, Zeroable};

use crate::material::Material;
use crate::sceneobjects::lights::{Ambient, Light, Suns};
use bevy_color::Alpha;
use glam::{Mat4, Vec3};
use wgpu::util::DeviceExt;

use bevy_ecs::schedule::IntoScheduleConfigs;

use crate::ecs::{Application, Plugin, Resource};
use crate::materials::openpbr::{
    EnergyTables, Material as OpenPbrMaterial, OpenPbrSurface, material_of,
};
use crate::mesh::{MeshData, Vertex};
use crate::sceneobjects::lights::GpuLight;
use crate::ui::{Color, linear_rgba, wgpu_color};

pub const MAX_MATERIALS: usize = 256;

/// A mesh uploaded to the GPU. Cheap to copy; see [`crate::AppState::mesh`].
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
pub struct MeshId(pub usize);

/// Where the meshes loaded from an asset ended up, by node name.
#[derive(Resource, Clone, Default, Debug)]
pub struct MeshHandles(HashMap<String, MeshId>);

impl MeshHandles {
    pub fn get(&self, name: &str) -> Option<MeshId> {
        self.0.get(name).copied()
    }

    pub fn insert(&mut self, name: impl Into<String>, id: MeshId) {
        self.0.insert(name.into(), id);
    }

    pub fn names(&self) -> impl Iterator<Item = &str> {
        self.0.keys().map(String::as_str)
    }
}

pub use crate::sceneobjects::cameras::{Camera, OrbitCamera, Ray};

/// Placement of a model in the world: position, orientation, and scale.
#[derive(crate::ecs::Component, Clone, Copy, Debug)]
pub struct Transform {
    pub translation: Vec3,
    pub rotation: glam::Quat,
    pub scale: f32,
}

impl Default for Transform {
    fn default() -> Self {
        Self {
            translation: Vec3::ZERO,
            rotation: glam::Quat::IDENTITY,
            scale: 1.0,
        }
    }
}

impl Transform {
    pub fn at(x: f32, y: f32, z: f32) -> Self {
        Self {
            translation: Vec3::new(x, y, z),
            ..Self::default()
        }
    }

    pub fn set_yaw(&mut self, radians: f32) {
        self.rotation = glam::Quat::from_rotation_y(radians);
    }

    /// Turns by `rotation` in world axes (left-multiplied).
    pub fn rotate(&mut self, rotation: glam::Quat) {
        self.rotation = (rotation * self.rotation).normalize();
    }

    pub fn matrix(&self) -> Mat4 {
        Mat4::from_scale_rotation_translation(
            Vec3::splat(self.scale),
            self.rotation,
            self.translation,
        )
    }
}

/// A surface registered with the renderer, as [`MeshId`] is a mesh.
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
pub struct MaterialId(pub u32);

/// The mesh an entity draws, the color to draw it in, and how it takes light.
#[derive(crate::ecs::Component, Clone, Copy, Debug)]
pub struct MeshInstance {
    pub mesh: MeshId,
    /// Multiplies the material's base colour.
    pub color: Color,
    pub material: MaterialId,
}

/// One model as the GPU sees it.
#[repr(C)]
#[derive(Clone, Copy, Pod, Zeroable)]
pub struct ModelInstance {
    model: [f32; 16],
    color: [f32; 4],
    /// x: material index; padded to a whole vector because the attribute must be one.
    material: [u32; 4],
}

impl ModelInstance {
    pub fn model(&self) -> Mat4 {
        Mat4::from_cols_array(&self.model)
    }
}

/// The models to draw this frame, sorted by mesh so each is one instanced draw call.
#[derive(Resource, Default)]
pub struct ModelDrawList {
    pub opaque: Vec<(MeshId, ModelInstance)>,
    pub translucent: Vec<(MeshId, ModelInstance)>,
}

pub fn collect_models_system(
    mut draw_list: bevy_ecs::system::ResMut<ModelDrawList>,
    models: bevy_ecs::system::Query<(&Transform, &MeshInstance)>,
) {
    draw_list.opaque.clear();
    draw_list.translucent.clear();

    for (transform, instance) in &models {
        let draw = (
            instance.mesh,
            ModelInstance {
                model: transform.matrix().to_cols_array(),
                color: linear_rgba(instance.color),
                material: [instance.material.0, 0, 0, 0],
            },
        );
        if instance.color.alpha() < 1.0 {
            draw_list.translucent.push(draw);
        } else {
            draw_list.opaque.push(draw);
        }
    }

    draw_list.opaque.sort_by_key(|(mesh, _)| *mesh);
    draw_list.translucent.sort_by_key(|(mesh, _)| *mesh);
}

#[repr(C)]
#[derive(Clone, Copy, Pod, Zeroable)]
struct CameraUniform {
    view_proj: [f32; 16],
    light_view_proj: [f32; 16],
    ao_view_proj: [f32; 16],
    /// xyz towards the key light, w its strength.
    light_dir: [f32; 4],
    key_color: [f32; 4],
    /// xyz towards the rim light, w its strength.
    rim_dir: [f32; 4],
    rim_color: [f32; 4],
    ambient: [f32; 4],
    eye: [f32; 4],
    /// x: penumbra per unit of map depth, y: shadow texel in world units, z: shadow texel in uv.
    shadow: [f32; 4],
    /// x: reach in map uv, y: world height per unit depth, z: strength, w: reach in world units.
    ao: [f32; 4],
    /// x: local light count (not `arrayLength`, which reports the whole fixed-size buffer).
    counts: [f32; 4],
    /// x, y: depth-to-slice scale and bias, matching the culling pass; z, w: frame in pixels.
    cluster: [f32; 4],
}

/// How the shadow map was fitted around the scene this frame.
#[derive(Clone, Copy, Debug)]
pub struct ShadowFit {
    pub light_view_proj: Mat4,
    /// Penumbra radius per unit of map depth, in map uv: `tan(angle) * depth_span / width`.
    pub spread: f32,
    pub texel_world: f32,
    pub depth_span: f32,
    pub width: f32,
}

/// Fits an orthographic light view around `min..max`, sized to the bounding sphere
/// so it does not change as the light swings.
pub fn fit_light(
    min: Vec3,
    max: Vec3,
    light_dir: Vec3,
    angle_degrees: f32,
    map_size: f32,
) -> ShadowFit {
    let center = (min + max) * 0.5;
    let radius = ((max - min).length() * 0.5).max(1e-5);
    let direction = light_dir.normalize_or(Vec3::Y);
    // look_at needs an up that is not the direction itself.
    let up = if direction.y.abs() > 0.99 {
        Vec3::Z
    } else {
        Vec3::Y
    };

    let eye = center + direction * radius * 2.0;
    let view = glam::camera::rh::view::look_at_mat4(eye, center, up);
    let (near, far) = (0.0, radius * 4.0);
    let projection =
        glam::camera::rh::proj::directx::orthographic(-radius, radius, -radius, radius, near, far);

    let width = radius * 2.0;
    ShadowFit {
        light_view_proj: projection * view,
        spread: angle_degrees.to_radians().tan() * (far - near) / width,
        texel_world: width / map_size.max(1.0),
        depth_span: far - near,
        width,
    }
}

/// The world-space bounds of everything being drawn, or `None` if nothing is.
pub fn scene_bounds(
    draws: &[(MeshId, ModelInstance)],
    bounds_of: impl Fn(MeshId) -> Option<(Vec3, Vec3)>,
) -> Option<(Vec3, Vec3)> {
    let mut min = Vec3::splat(f32::MAX);
    let mut max = Vec3::splat(f32::MIN);
    let mut any = false;

    for (mesh, instance) in draws {
        let Some((local_min, local_max)) = bounds_of(*mesh) else {
            continue;
        };
        let model = Mat4::from_cols_array(&instance.model);
        for i in 0..8 {
            let corner = Vec3::new(
                if i & 1 == 0 { local_min.x } else { local_max.x },
                if i & 2 == 0 { local_min.y } else { local_max.y },
                if i & 4 == 0 { local_min.z } else { local_max.z },
            );
            let world = model.transform_point3(corner);
            min = min.min(world);
            max = max.max(world);
            any = true;
        }
    }

    any.then_some((min, max))
}

struct GpuMesh {
    vertices: wgpu::Buffer,
    indices: wgpu::Buffer,
    index_count: u32,
    base_color: Option<Color>,
    min: Vec3,
    max: Vec3,
}

pub const DEPTH_FORMAT: wgpu::TextureFormat = wgpu::TextureFormat::Depth32Float;

const SHADOW_MAP_SIZE: u32 = 2048;

const AO_MAP_SIZE: u32 = 1024;

/// Occlusion hunt reach as a fraction of the fitted box (about 1.5 pieces).
const AO_REACH: f32 = 0.045;

pub const SAMPLES: u32 = 4;

/// The images a frame is drawn into, shared by every 3D pass.
pub struct Targets<'a> {
    pub color: &'a wgpu::TextureView,
    pub resolve: &'a wgpu::TextureView,
    pub depth: &'a wgpu::TextureView,
}

fn storage_entry(binding: u32) -> wgpu::BindGroupLayoutEntry {
    wgpu::BindGroupLayoutEntry {
        binding,
        visibility: wgpu::ShaderStages::FRAGMENT,
        ty: wgpu::BindingType::Buffer {
            ty: wgpu::BufferBindingType::Storage { read_only: true },
            has_dynamic_offset: false,
            min_binding_size: None,
        },
        count: None,
    }
}

/// Draws [`ModelDrawList`] with one instanced draw call per mesh.
pub struct MeshRenderer {
    pipeline: wgpu::RenderPipeline,
    translucent_pipeline: wgpu::RenderPipeline,
    shadow_pipeline: wgpu::RenderPipeline,
    camera_buffer: wgpu::Buffer,
    bind_group: wgpu::BindGroup,
    shadow_bind_group: wgpu::BindGroup,
    shadow_view: wgpu::TextureView,
    ao_pipeline: wgpu::RenderPipeline,
    ao_view: wgpu::TextureView,
    msaa: Option<(u32, u32, wgpu::TextureView)>,
    resolve: Option<(u32, u32, wgpu::TextureView)>,
    format: wgpu::TextureFormat,
    instance_buffer: wgpu::Buffer,
    instance_capacity: usize,
    meshes: Vec<GpuMesh>,
    depth: Option<(u32, u32, wgpu::TextureView)>,
    materials: Vec<OpenPbrMaterial>,
    energy: Vec<f32>,
    /// How many of `materials` have reached the GPU; the table only grows.
    materials_uploaded: usize,
    material_buffer: wgpu::Buffer,
    energy_buffer: wgpu::Buffer,
}

impl MeshRenderer {
    pub fn new(
        device: &wgpu::Device,
        format: wgpu::TextureFormat,
        clusters: &crate::clustered::Clusters,
    ) -> Self {
        let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("model shader"),
            source: wgpu::ShaderSource::Wgsl(
                format!(
                    "{}\n{}",
                    crate::materials::openpbr::SHADER,
                    include_str!("model.wgsl"),
                )
                .into(),
            ),
        });

        let camera_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("model camera uniform"),
            size: std::mem::size_of::<CameraUniform>() as u64,
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        let shadow_texture = device.create_texture(&wgpu::TextureDescriptor {
            label: Some("shadow map"),
            size: wgpu::Extent3d {
                width: SHADOW_MAP_SIZE,
                height: SHADOW_MAP_SIZE,
                depth_or_array_layers: 1,
            },
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format: DEPTH_FORMAT,
            usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::TEXTURE_BINDING,
            view_formats: &[],
        });
        let shadow_view = shadow_texture.create_view(&wgpu::TextureViewDescriptor::default());

        let ao_texture = device.create_texture(&wgpu::TextureDescriptor {
            label: Some("occlusion map"),
            size: wgpu::Extent3d {
                width: AO_MAP_SIZE,
                height: AO_MAP_SIZE,
                depth_or_array_layers: 1,
            },
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format: DEPTH_FORMAT,
            usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::TEXTURE_BINDING,
            view_formats: &[],
        });
        let ao_view = ao_texture.create_view(&wgpu::TextureViewDescriptor::default());

        let shadow_sampler = device.create_sampler(&wgpu::SamplerDescriptor {
            label: Some("shadow sampler"),
            address_mode_u: wgpu::AddressMode::ClampToEdge,
            address_mode_v: wgpu::AddressMode::ClampToEdge,
            address_mode_w: wgpu::AddressMode::ClampToEdge,
            mag_filter: wgpu::FilterMode::Linear,
            min_filter: wgpu::FilterMode::Linear,
            compare: Some(wgpu::CompareFunction::LessEqual),
            ..Default::default()
        });

        // Fixed-size buffers so the bind group never has to be rebuilt mid-frame.
        let material_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("openpbr materials"),
            size: (MAX_MATERIALS * std::mem::size_of::<OpenPbrMaterial>()) as u64,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        let energy_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("openpbr energy tables"),
            size: (MAX_MATERIALS * std::mem::size_of::<EnergyTables>()) as u64,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("model bind group 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: None,
                    },
                    count: None,
                },
                wgpu::BindGroupLayoutEntry {
                    binding: 1,
                    visibility: wgpu::ShaderStages::FRAGMENT,
                    ty: wgpu::BindingType::Texture {
                        sample_type: wgpu::TextureSampleType::Depth,
                        view_dimension: wgpu::TextureViewDimension::D2,
                        multisampled: false,
                    },
                    count: None,
                },
                wgpu::BindGroupLayoutEntry {
                    binding: 2,
                    visibility: wgpu::ShaderStages::FRAGMENT,
                    ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Comparison),
                    count: None,
                },
                wgpu::BindGroupLayoutEntry {
                    binding: 3,
                    visibility: wgpu::ShaderStages::FRAGMENT,
                    ty: wgpu::BindingType::Texture {
                        sample_type: wgpu::TextureSampleType::Depth,
                        view_dimension: wgpu::TextureViewDimension::D2,
                        multisampled: false,
                    },
                    count: None,
                },
                storage_entry(4),
                storage_entry(5),
                storage_entry(6),
                storage_entry(7),
                storage_entry(8),
            ],
        });

        let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("model bind group"),
            layout: &bind_group_layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: camera_buffer.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: wgpu::BindingResource::TextureView(&shadow_view),
                },
                wgpu::BindGroupEntry {
                    binding: 2,
                    resource: wgpu::BindingResource::Sampler(&shadow_sampler),
                },
                wgpu::BindGroupEntry {
                    binding: 3,
                    resource: wgpu::BindingResource::TextureView(&ao_view),
                },
                wgpu::BindGroupEntry {
                    binding: 4,
                    resource: material_buffer.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 5,
                    resource: energy_buffer.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 6,
                    resource: clusters.lights.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 7,
                    resource: clusters.counts.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 8,
                    resource: clusters.indices.as_entire_binding(),
                },
            ],
        });

        // The shadow pass writes the map, so it cannot also bind it.
        let shadow_bind_group_layout =
            device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                label: Some("shadow bind group 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,
                }],
            });
        let shadow_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("shadow bind group"),
            layout: &shadow_bind_group_layout,
            entries: &[wgpu::BindGroupEntry {
                binding: 0,
                resource: camera_buffer.as_entire_binding(),
            }],
        });

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

        const VERTEX_ATTRS: [wgpu::VertexAttribute; 2] =
            wgpu::vertex_attr_array![0 => Float32x3, 1 => Float32x3];
        const INSTANCE_ATTRS: [wgpu::VertexAttribute; 6] = wgpu::vertex_attr_array![
            2 => Float32x4, 3 => Float32x4, 4 => Float32x4, 5 => Float32x4, 6 => Float32x4,
            7 => Uint32x4
        ];

        let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: Some("model pipeline"),
            layout: Some(&pipeline_layout),
            vertex: wgpu::VertexState {
                module: &shader,
                entry_point: Some("vs_main"),
                buffers: &[
                    Some(wgpu::VertexBufferLayout {
                        array_stride: std::mem::size_of::<Vertex>() as u64,
                        step_mode: wgpu::VertexStepMode::Vertex,
                        attributes: &VERTEX_ATTRS,
                    }),
                    Some(wgpu::VertexBufferLayout {
                        array_stride: std::mem::size_of::<ModelInstance>() as u64,
                        step_mode: wgpu::VertexStepMode::Instance,
                        attributes: &INSTANCE_ATTRS,
                    }),
                ],
                compilation_options: wgpu::PipelineCompilationOptions::default(),
            },
            fragment: Some(wgpu::FragmentState {
                module: &shader,
                entry_point: Some("fs_main"),
                targets: &[Some(wgpu::ColorTargetState {
                    format,
                    blend: Some(wgpu::BlendState::REPLACE),
                    write_mask: wgpu::ColorWrites::ALL,
                })],
                compilation_options: wgpu::PipelineCompilationOptions::default(),
            }),
            primitive: wgpu::PrimitiveState {
                topology: wgpu::PrimitiveTopology::TriangleList,
                strip_index_format: None,
                front_face: wgpu::FrontFace::Ccw,
                cull_mode: Some(wgpu::Face::Back),
                polygon_mode: wgpu::PolygonMode::Fill,
                unclipped_depth: false,
                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 {
                count: SAMPLES,
                ..wgpu::MultisampleState::default()
            },
            multiview_mask: None,
            cache: None,
        });

        let shadow_shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("shadow shader"),
            source: wgpu::ShaderSource::Wgsl(include_str!("shadow.wgsl").into()),
        });
        let shadow_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("shadow pipeline layout"),
            bind_group_layouts: &[Some(&shadow_bind_group_layout)],
            immediate_size: 0,
        });
        const SHADOW_VERTEX_ATTRS: [wgpu::VertexAttribute; 1] =
            wgpu::vertex_attr_array![0 => Float32x3];
        let shadow_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: Some("shadow pipeline"),
            layout: Some(&shadow_layout),
            vertex: wgpu::VertexState {
                module: &shadow_shader,
                entry_point: Some("vs_main"),
                buffers: &[
                    Some(wgpu::VertexBufferLayout {
                        array_stride: std::mem::size_of::<Vertex>() as u64,
                        step_mode: wgpu::VertexStepMode::Vertex,
                        attributes: &SHADOW_VERTEX_ATTRS,
                    }),
                    Some(wgpu::VertexBufferLayout {
                        array_stride: std::mem::size_of::<ModelInstance>() as u64,
                        step_mode: wgpu::VertexStepMode::Instance,
                        attributes: &INSTANCE_ATTRS[..4],
                    }),
                ],
                compilation_options: wgpu::PipelineCompilationOptions::default(),
            },
            fragment: None,
            primitive: wgpu::PrimitiveState {
                topology: wgpu::PrimitiveTopology::TriangleList,
                strip_index_format: None,
                front_face: wgpu::FrontFace::Ccw,
                cull_mode: Some(wgpu::Face::Back),
                polygon_mode: wgpu::PolygonMode::Fill,
                unclipped_depth: false,
                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(),
                // Slope-scaled: a face edge-on to the light spans unbounded depth in one texel.
                bias: wgpu::DepthBiasState {
                    constant: 2,
                    slope_scale: 2.0,
                    clamp: 0.0,
                },
            }),
            multisample: wgpu::MultisampleState::default(),
            multiview_mask: None,
            cache: None,
        });

        let translucent_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: Some("translucent model pipeline"),
            layout: Some(&pipeline_layout),
            vertex: wgpu::VertexState {
                module: &shader,
                entry_point: Some("vs_main"),
                buffers: &[
                    Some(wgpu::VertexBufferLayout {
                        array_stride: std::mem::size_of::<Vertex>() as u64,
                        step_mode: wgpu::VertexStepMode::Vertex,
                        attributes: &VERTEX_ATTRS,
                    }),
                    Some(wgpu::VertexBufferLayout {
                        array_stride: std::mem::size_of::<ModelInstance>() as u64,
                        step_mode: wgpu::VertexStepMode::Instance,
                        attributes: &INSTANCE_ATTRS,
                    }),
                ],
                compilation_options: wgpu::PipelineCompilationOptions::default(),
            },
            fragment: Some(wgpu::FragmentState {
                module: &shader,
                entry_point: Some("fs_main"),
                targets: &[Some(wgpu::ColorTargetState {
                    format,
                    blend: Some(wgpu::BlendState::ALPHA_BLENDING),
                    write_mask: wgpu::ColorWrites::ALL,
                })],
                compilation_options: wgpu::PipelineCompilationOptions::default(),
            }),
            primitive: wgpu::PrimitiveState {
                topology: wgpu::PrimitiveTopology::TriangleList,
                strip_index_format: None,
                front_face: wgpu::FrontFace::Ccw,
                cull_mode: Some(wgpu::Face::Back),
                polygon_mode: wgpu::PolygonMode::Fill,
                unclipped_depth: false,
                conservative: false,
            },
            depth_stencil: Some(wgpu::DepthStencilState {
                format: DEPTH_FORMAT,
                depth_write_enabled: Some(false),
                depth_compare: Some(wgpu::CompareFunction::Less),
                stencil: wgpu::StencilState::default(),
                bias: wgpu::DepthBiasState::default(),
            }),
            multisample: wgpu::MultisampleState {
                count: SAMPLES,
                ..wgpu::MultisampleState::default()
            },
            multiview_mask: None,
            cache: None,
        });

        let ao_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: Some("occlusion pipeline"),
            layout: Some(&shadow_layout),
            vertex: wgpu::VertexState {
                module: &shadow_shader,
                entry_point: Some("vs_ao"),
                buffers: &[
                    Some(wgpu::VertexBufferLayout {
                        array_stride: std::mem::size_of::<Vertex>() as u64,
                        step_mode: wgpu::VertexStepMode::Vertex,
                        attributes: &SHADOW_VERTEX_ATTRS,
                    }),
                    Some(wgpu::VertexBufferLayout {
                        array_stride: std::mem::size_of::<ModelInstance>() as u64,
                        step_mode: wgpu::VertexStepMode::Instance,
                        attributes: &INSTANCE_ATTRS[..4],
                    }),
                ],
                compilation_options: wgpu::PipelineCompilationOptions::default(),
            },
            fragment: None,
            primitive: wgpu::PrimitiveState {
                topology: wgpu::PrimitiveTopology::TriangleList,
                strip_index_format: None,
                front_face: wgpu::FrontFace::Ccw,
                // Both sides: the height field needs the underside of a mesh with no top.
                cull_mode: None,
                polygon_mode: wgpu::PolygonMode::Fill,
                unclipped_depth: false,
                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(),
            multiview_mask: None,
            cache: None,
        });

        let instance_capacity = 256;
        let instance_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("model instances"),
            size: (instance_capacity * std::mem::size_of::<ModelInstance>()) as u64,
            usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        Self {
            pipeline,
            translucent_pipeline,
            shadow_pipeline,
            camera_buffer,
            bind_group,
            shadow_bind_group,
            shadow_view,
            ao_pipeline,
            ao_view,
            msaa: None,
            resolve: None,
            format,
            materials: Vec::new(),
            energy: Vec::new(),
            materials_uploaded: 0,
            material_buffer,
            energy_buffer,
            instance_buffer,
            instance_capacity,
            meshes: Vec::new(),
            depth: None,
        }
    }

    /// Uploads a mesh and returns the handle that draws it.
    pub fn upload(&mut self, device: &wgpu::Device, mesh: &MeshData) -> MeshId {
        let vertices = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some(&format!("{} vertices", mesh.name)),
            contents: bytemuck::cast_slice(&mesh.vertices),
            usage: wgpu::BufferUsages::VERTEX,
        });
        let indices = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some(&format!("{} indices", mesh.name)),
            contents: bytemuck::cast_slice(&mesh.indices),
            usage: wgpu::BufferUsages::INDEX,
        });

        self.meshes.push(GpuMesh {
            vertices,
            indices,
            index_count: mesh.indices.len() as u32,
            base_color: mesh.base_color,
            min: mesh.min,
            max: mesh.max,
        });
        MeshId(self.meshes.len() - 1)
    }

    pub fn base_color(&self, mesh: MeshId) -> Option<Color> {
        self.meshes.get(mesh.0).and_then(|m| m.base_color)
    }

    pub fn is_empty(&self) -> bool {
        self.meshes.is_empty()
    }

    #[allow(clippy::too_many_arguments)]
    pub fn render(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        encoder: &mut wgpu::CommandEncoder,
        width: u32,
        height: u32,
        clear_color: Color,
        view: &crate::views::View,
        clear: bool,
        suns: &Suns,
        ambient: &Ambient,
        draws: &ModelDrawList,
        lights: &[GpuLight],
    ) {
        let camera = &view.camera;
        let (opaque, translucent) = (&draws.opaque, &draws.translucent);
        let bounds = scene_bounds(opaque, |mesh| {
            self.meshes.get(mesh.0).map(|m| (m.min, m.max))
        });
        let fit = bounds
            .zip(suns.shadowed)
            .map(|((min, max), sun)| {
                fit_light(min, max, sun.direction, sun.angle, SHADOW_MAP_SIZE as f32)
            })
            .unwrap_or(ShadowFit {
                light_view_proj: Mat4::IDENTITY,
                spread: 0.0,
                texel_world: 0.0,
                depth_span: 0.0,
                width: 0.0,
            });

        let ao_fit = bounds
            .map(|(min, max)| fit_light(min, max, Vec3::Y, 0.0, AO_MAP_SIZE as f32))
            .unwrap_or(ShadowFit {
                light_view_proj: Mat4::IDENTITY,
                spread: 0.0,
                texel_world: 0.0,
                depth_span: 0.0,
                width: 0.0,
            });

        // An empty slot goes up as a zero-strength light so the shader needs no branch.
        let sun = |slot: Option<Light>| {
            let light = slot.unwrap_or(Light::default().strength(0.0));
            let color = light.color();
            (
                [
                    light.direction.x,
                    light.direction.y,
                    light.direction.z,
                    light.strength,
                ],
                [color.x, color.y, color.z, 0.0],
            )
        };
        let (key_dir, key_color) = sun(suns.shadowed);
        let (rim_dir, rim_color) = sun(suns.unshadowed);
        let ambient_color = ambient.color();
        let uniform = CameraUniform {
            view_proj: camera.view_proj(view.aspect()).to_cols_array(),
            light_view_proj: fit.light_view_proj.to_cols_array(),
            ao_view_proj: ao_fit.light_view_proj.to_cols_array(),
            light_dir: key_dir,
            key_color,
            rim_dir,
            rim_color,
            ambient: [ambient_color.x, ambient_color.y, ambient_color.z, 0.0],
            eye: [camera.eye.x, camera.eye.y, camera.eye.z, 0.0],
            shadow: [
                fit.spread,
                fit.texel_world,
                1.0 / SHADOW_MAP_SIZE as f32,
                0.0,
            ],
            ao: [
                AO_REACH,
                ao_fit.depth_span,
                ambient.occlusion,
                ao_fit.width * AO_REACH,
            ],
            counts: [lights.len() as f32, 0.0, 0.0, 0.0],
            cluster: {
                let (scale, bias) = crate::clustered::slice_scale_bias(camera.near, camera.far);
                [scale, bias, width.max(1) as f32, height.max(1) as f32]
            },
        };
        queue.write_buffer(&self.camera_buffer, 0, bytemuck::bytes_of(&uniform));
        self.upload_shading(queue);

        let instances: Vec<ModelInstance> = opaque
            .iter()
            .chain(translucent.iter())
            .map(|(_, instance)| *instance)
            .collect();
        if !instances.is_empty() {
            self.ensure_capacity(device, instances.len());
            queue.write_buffer(&self.instance_buffer, 0, bytemuck::cast_slice(&instances));
            // Translucent models cast neither shadow nor occlusion.
            self.shadow_pass(encoder, opaque);
            self.ao_pass(encoder, opaque);
        }

        let targets = self.targets(device, width, height);
        let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
            label: Some("model render pass"),
            color_attachments: &[Some(wgpu::RenderPassColorAttachment {
                view: targets.color,
                resolve_target: Some(targets.resolve),
                depth_slice: None,
                // Only the first view clears: a clear takes the whole attachment, not the viewport.
                ops: wgpu::Operations {
                    load: match clear {
                        true => wgpu::LoadOp::Clear(wgpu_color(clear_color)),
                        false => wgpu::LoadOp::Load,
                    },
                    store: wgpu::StoreOp::Store,
                },
            })],
            depth_stencil_attachment: Some(wgpu::RenderPassDepthStencilAttachment {
                view: targets.depth,
                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,
        });

        if instances.is_empty() {
            return;
        }

        set_view(&mut pass, view, (width, height));

        pass.set_bind_group(0, &self.bind_group, &[]);
        pass.set_vertex_buffer(1, self.instance_buffer.slice(..));

        pass.set_pipeline(&self.pipeline);
        self.draw_batches(&mut pass, opaque, 0);

        if !translucent.is_empty() {
            pass.set_pipeline(&self.translucent_pipeline);
            self.draw_batches(&mut pass, translucent, opaque.len());
        }
    }

    fn shadow_pass(&self, encoder: &mut wgpu::CommandEncoder, draws: &[(MeshId, ModelInstance)]) {
        let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
            label: Some("shadow pass"),
            color_attachments: &[],
            depth_stencil_attachment: Some(wgpu::RenderPassDepthStencilAttachment {
                view: &self.shadow_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,
        });

        pass.set_pipeline(&self.shadow_pipeline);
        pass.set_bind_group(0, &self.shadow_bind_group, &[]);
        pass.set_vertex_buffer(1, self.instance_buffer.slice(..));
        self.draw_batches(&mut pass, draws, 0);
    }

    /// Draws the scene from straight above, depth only: a height field.
    fn ao_pass(&self, encoder: &mut wgpu::CommandEncoder, draws: &[(MeshId, ModelInstance)]) {
        let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
            label: Some("occlusion pass"),
            color_attachments: &[],
            depth_stencil_attachment: Some(wgpu::RenderPassDepthStencilAttachment {
                view: &self.ao_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,
        });

        pass.set_pipeline(&self.ao_pipeline);
        pass.set_bind_group(0, &self.shadow_bind_group, &[]);
        pass.set_vertex_buffer(1, self.instance_buffer.slice(..));
        self.draw_batches(&mut pass, draws, 0);
    }

    fn draw_batches<'pass>(
        &'pass self,
        pass: &mut wgpu::RenderPass<'pass>,
        draws: &[(MeshId, ModelInstance)],
        base: usize,
    ) {
        let mut start = 0usize;
        while start < draws.len() {
            let mesh_id = draws[start].0;
            let mut end = start + 1;
            while end < draws.len() && draws[end].0 == mesh_id {
                end += 1;
            }

            if let Some(mesh) = self.meshes.get(mesh_id.0) {
                pass.set_vertex_buffer(0, mesh.vertices.slice(..));
                pass.set_index_buffer(mesh.indices.slice(..), wgpu::IndexFormat::Uint32);
                pass.draw_indexed(
                    0..mesh.index_count,
                    0,
                    (base + start) as u32..(base + end) as u32,
                );
            }
            start = end;
        }
    }

    pub fn targets(&mut self, device: &wgpu::Device, width: u32, height: u32) -> Targets<'_> {
        let size = wgpu::Extent3d {
            width: width.max(1),
            height: height.max(1),
            depth_or_array_layers: 1,
        };

        if !matches!(&self.depth, Some((w, h, _)) if *w == width && *h == height) {
            let texture = device.create_texture(&wgpu::TextureDescriptor {
                label: Some("model depth"),
                size,
                mip_level_count: 1,
                sample_count: SAMPLES,
                dimension: wgpu::TextureDimension::D2,
                format: DEPTH_FORMAT,
                usage: wgpu::TextureUsages::RENDER_ATTACHMENT,
                view_formats: &[],
            });
            let view = texture.create_view(&wgpu::TextureViewDescriptor::default());
            self.depth = Some((width, height, view));
        }

        if !matches!(&self.msaa, Some((w, h, _)) if *w == width && *h == height) {
            let texture = device.create_texture(&wgpu::TextureDescriptor {
                label: Some("model msaa"),
                size,
                mip_level_count: 1,
                sample_count: SAMPLES,
                dimension: wgpu::TextureDimension::D2,
                format: self.format,
                usage: wgpu::TextureUsages::RENDER_ATTACHMENT,
                view_formats: &[],
            });
            let view = texture.create_view(&wgpu::TextureViewDescriptor::default());
            self.msaa = Some((width, height, view));
        }

        if !matches!(&self.resolve, Some((w, h, _)) if *w == width && *h == height) {
            let texture = device.create_texture(&wgpu::TextureDescriptor {
                label: Some("model resolve"),
                size,
                mip_level_count: 1,
                sample_count: 1,
                dimension: wgpu::TextureDimension::D2,
                format: self.format,
                usage: wgpu::TextureUsages::RENDER_ATTACHMENT
                    | wgpu::TextureUsages::TEXTURE_BINDING,
                view_formats: &[],
            });
            let view = texture.create_view(&wgpu::TextureViewDescriptor::default());
            self.resolve = Some((width, height, view));
        }

        Targets {
            color: &self.msaa.as_ref().expect("just created").2,
            resolve: &self.resolve.as_ref().expect("just created").2,
            depth: &self.depth.as_ref().expect("just created").2,
        }
    }

    /// Registers `surface`, or finds it if already registered, and returns its id.
    pub fn register_material(&mut self, surface: &OpenPbrSurface) -> MaterialId {
        let packed = material_of(surface);
        if let Some(index) = self.materials.iter().position(|m| *m == packed) {
            return MaterialId(index as u32);
        }
        if self.materials.len() == MAX_MATERIALS {
            log::error!("more than {MAX_MATERIALS} materials; reusing the first");
            return MaterialId(0);
        }

        let tables = EnergyTables::compute(surface);
        self.energy
            .extend_from_slice(bytemuck::cast_slice(std::slice::from_ref(&tables)));
        self.materials.push(packed);
        MaterialId(self.materials.len() as u32 - 1)
    }

    fn upload_shading(&mut self, queue: &wgpu::Queue) {
        if self.materials_uploaded < self.materials.len() {
            let first = self.materials_uploaded;
            let stride = std::mem::size_of::<OpenPbrMaterial>();
            queue.write_buffer(
                &self.material_buffer,
                (first * stride) as u64,
                bytemuck::cast_slice(&self.materials[first..]),
            );
            let floats = std::mem::size_of::<EnergyTables>() / std::mem::size_of::<f32>();
            queue.write_buffer(
                &self.energy_buffer,
                (first * floats * std::mem::size_of::<f32>()) as u64,
                bytemuck::cast_slice(&self.energy[first * floats..]),
            );
            self.materials_uploaded = self.materials.len();
        }
    }

    fn ensure_capacity(&mut self, device: &wgpu::Device, needed: usize) {
        if needed <= self.instance_capacity {
            return;
        }
        let capacity = needed.next_power_of_two();
        self.instance_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("model instances"),
            size: (capacity * std::mem::size_of::<ModelInstance>()) as u64,
            usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        self.instance_capacity = capacity;
    }
}

/// Builder for a model in the world, spawned with [`crate::AppState::spawn_model`].
///
/// ```no_run
/// # use codecraft::{AppState, render3d::Model, ui::Color};
/// # fn demo(app: &mut AppState) {
/// app.spawn_model(Model::new("chess_piece_king").at(0.0, 0.0, 0.0).color(Color::WHITE));
/// # }
/// ```
pub struct Model {
    pub mesh: String,
    pub transform: Transform,
    /// `None` keeps the color the mesh's own material came with.
    pub color: Option<Color>,
    pub material: Material,
    /// The full OpenPBR surface; wins over `material` when set.
    pub surface: Option<OpenPbrSurface>,
}

impl Model {
    pub fn new(mesh: impl Into<String>) -> Self {
        Self {
            mesh: mesh.into(),
            transform: Transform::default(),
            color: None,
            material: Material::default(),
            surface: None,
        }
    }

    /// The surface this model is drawn with.
    pub fn surface(&self) -> OpenPbrSurface {
        self.surface
            .clone()
            .unwrap_or_else(|| self.material.to_openpbr())
    }

    pub fn at(mut self, x: f32, y: f32, z: f32) -> Self {
        self.transform.translation = Vec3::new(x, y, z);
        self
    }

    /// Rotation about the up axis, in radians.
    pub fn yaw(mut self, radians: f32) -> Self {
        self.transform.set_yaw(radians);
        self
    }

    pub fn scale(mut self, scale: f32) -> Self {
        self.transform.scale = scale;
        self
    }

    /// How bright and how tight a highlight the surface takes.
    pub fn gloss(mut self, specular: f32, shininess: f32) -> Self {
        self.material = Material::gloss(specular, shininess);
        self
    }

    pub fn material(mut self, material: Material) -> Self {
        self.material = material;
        self
    }

    /// Sets a full OpenPBR surface.
    ///
    /// ```no_run
    /// # use codecraft::{AppState, render3d::Model};
    /// # use codecraft::materials::openpbr::{Color3, OpenPbrSurface};
    /// # fn demo(app: &mut AppState) {
    /// app.spawn_model(Model::new("chess_piece_king").of(OpenPbrSurface {
    ///     base_color: Color3::new(0.9, 0.87, 0.8),
    ///     specular_roughness: 0.25,
    ///     coat_weight: 1.0,
    ///     coat_roughness: 0.05,
    ///     ..OpenPbrSurface::default()
    /// }));
    /// # }
    /// ```
    pub fn of(mut self, surface: OpenPbrSurface) -> Self {
        self.surface = Some(surface);
        self
    }

    pub fn color(mut self, color: Color) -> Self {
        self.color = Some(color);
        self
    }
}

/// Registers the 3D resources and the per-frame draw-list system.
pub struct Render3dPlugin;

impl Plugin for Render3dPlugin {
    fn build(&self, app: &mut Application) {
        app.insert_resource(Camera::default());
        app.init_resource::<crate::views::Views>();
        app.init_resource::<crate::ui::Profiler>();
        app.insert_resource(crate::sceneobjects::lights::Suns::default());
        app.insert_resource(crate::sceneobjects::lights::Ambient::default());
        app.insert_resource(MeshHandles::default());
        app.insert_resource(ModelDrawList::default());
        app.insert_resource(crate::sceneobjects::lights::LightDrawList::default());
        // init_resource so a host/UI plugin that already inserted these is not overwritten.
        app.init_resource::<crate::input::Keys>();
        app.init_resource::<crate::time::Time>();
        app.init_resource::<crate::ui::MouseInput>();
        app.init_resource::<crate::ui::CursorPosition>();
        app.init_resource::<crate::ui::PointerCapture>();
        app.init_resource::<crate::hid::GamepadState>();
        app.init_resource::<crate::hid::Gamepads>();
        // The rig reads this frame's wheel and click edges before the UI schedule spends them.
        app.add_update_systems(
            crate::sceneobjects::cameras::orbit_camera_system
                .before(crate::ui::clear_input_edge_system),
        );
        app.add_update_systems(collect_models_system);
        app.add_update_systems(crate::sceneobjects::lights::collect_lights_system);
        app.add_update_systems(crate::sceneobjects::lights::collect_suns_system);
    }
}

/// Confines a pass to one view's corner of the frame; clamped because a viewport outside the attachment is a validation error.
pub(crate) fn set_view(
    pass: &mut wgpu::RenderPass<'_>,
    view: &crate::views::View,
    frame: (u32, u32),
) {
    let (frame_width, frame_height) = (frame.0 as f32, frame.1 as f32);
    let x = view.rect.x.clamp(0.0, frame_width);
    let y = view.rect.y.clamp(0.0, frame_height);
    let width = view.rect.width.clamp(0.0, frame_width - x);
    let height = view.rect.height.clamp(0.0, frame_height - y);
    if width <= 0.0 || height <= 0.0 {
        return;
    }
    pass.set_viewport(x, y, width, height, 0.0, 1.0);
    pass.set_scissor_rect(x as u32, y as u32, width as u32, height as u32);
}