use anyhow::Context;
use bytemuck::NoUninit;
use derive_builder::Builder;
use std::borrow::Cow;
use std::num::NonZeroU64;
#[cfg(feature = "window")]
use std::sync::Arc;
use std::sync::mpsc::{Receiver, Sender};
use wgpu::naga;
use wgpu::util::DeviceExt;

#[macro_export]
macro_rules! include_wgsl {
    ($filename:literal) => {
        include_str!(concat!(env!("OUT_DIR"), "/shaders/", $filename))
    };
}

// Magic to allow converting a primitive type, a Vec, or a slice into a slice of bytes.
pub trait Buffer<T> {
    fn as_bytes(&self) -> &[T];
}
impl<T: bytemuck::Pod> Buffer<T> for T {
    fn as_bytes(&self) -> &[T] {
        std::slice::from_ref(self)
    }
}
impl<T> Buffer<T> for Vec<T> {
    fn as_bytes(&self) -> &[T] {
        self.as_slice()
    }
}
impl<T> Buffer<T> for &[T] {
    fn as_bytes(&self) -> &[T] {
        self
    }
}

pub struct BufferBuilder<'a> {
    label: Option<String>,
    size: u64,
    usage: wgpu::BufferUsages,
    mapped_at_creation: bool,
    gpu: &'a Gpu,
}

impl<'a> BufferBuilder<'a> {
    fn new<T: bytemuck::Pod>(data: &impl Buffer<T>, gpu: &'a Gpu) -> Self {
        Self::new_by_sz(std::mem::size_of_val(data.as_bytes()), gpu)
    }

    fn new_by_sz(sz: usize, gpu: &'a Gpu) -> Self {
        Self {
            label: None,
            size: wgpu::util::align_to(sz, 16) as u64,
            usage: wgpu::BufferUsages::empty(),
            mapped_at_creation: false,
            gpu: gpu,
        }
    }

    pub fn copy_dst(&mut self) -> &mut Self {
        self.usage |= wgpu::BufferUsages::COPY_DST;
        self
    }

    pub fn copy_src(&mut self) -> &mut Self {
        self.usage |= wgpu::BufferUsages::COPY_SRC;
        self
    }

    pub fn map_read(&mut self) -> &mut Self {
        self.usage |= wgpu::BufferUsages::MAP_READ;
        self
    }

    pub fn query_resolve(&mut self) -> &mut Self {
        self.usage |= wgpu::BufferUsages::QUERY_RESOLVE;
        self
    }

    pub fn storage(&mut self) -> &mut Self {
        self.usage |= wgpu::BufferUsages::STORAGE;
        self
    }

    pub fn uniform(&mut self) -> &mut Self {
        self.usage |= wgpu::BufferUsages::UNIFORM;
        self
    }

    pub fn build(&self) -> wgpu::Buffer {
        self.gpu.device.create_buffer(&wgpu::BufferDescriptor {
            label: self.label.as_deref(),
            size: self.size,
            usage: self.usage,
            mapped_at_creation: self.mapped_at_creation,
        })
    }
}

pub struct GpuBuilder {
    required_features: wgpu::Features,
    #[cfg(feature = "window")]
    window: Option<Arc<winit::window::Window>>,
}

impl GpuBuilder {
    pub fn new() -> Self {
        Self {
            required_features: wgpu::Features::empty(),
            #[cfg(feature = "window")]
            window: None,
        }
    }

    pub fn enable_profiling(mut self) -> Self {
        self.required_features |= wgpu::Features::TIMESTAMP_QUERY;
        self
    }

    #[cfg(feature = "window")]
    pub fn window(mut self, window: winit::window::Window) -> Self {
        self.window = Some(Arc::new(window));
        self
    }

    pub async fn build(self) -> anyhow::Result<Gpu> {
        #[cfg(feature = "window")]
        let instance = wgpu::Instance::new(wgpu::InstanceDescriptor {
            backends: wgpu::Backends::PRIMARY,
            flags: Default::default(),
            memory_budget_thresholds: Default::default(),
            backend_options: Default::default(),
            display: None,
        });
        #[cfg(not(feature = "window"))]
        let instance = wgpu::Instance::new(wgpu::InstanceDescriptor::new_without_display_handle());

        #[cfg(feature = "window")]
        let surface = self
            .window
            .as_ref()
            .map(|window| instance.create_surface(window.clone()))
            .transpose()?;
        #[cfg(feature = "window")]
        let adapter_req_opt = wgpu::RequestAdapterOptions {
            power_preference: wgpu::PowerPreference::default(),
            compatible_surface: surface.as_ref(),
            force_fallback_adapter: false,
            apply_limit_buckets: true,
        };
        #[cfg(not(feature = "window"))]
        let adapter_req_opt = wgpu::RequestAdapterOptions::default();
        let adapter = instance.request_adapter(&adapter_req_opt).await?;
        println!("Got adapter: {:#?}", adapter.get_info());

        let supported = adapter.limits();
        let mut limits = wgpu::Limits::downlevel_defaults();
        limits.max_compute_invocations_per_workgroup =
            supported.max_compute_invocations_per_workgroup;
        limits.max_compute_workgroup_storage_size = supported.max_compute_workgroup_storage_size;
        println!(
            "Max compute invocations per workgroup: {}",
            limits.max_compute_invocations_per_workgroup
        );

        let (device, queue) = adapter
            .request_device(&wgpu::DeviceDescriptor {
                label: None,
                required_features: self.required_features,
                required_limits: limits,
                experimental_features: wgpu::ExperimentalFeatures::disabled(),
                memory_hints: wgpu::MemoryHints::MemoryUsage,
                trace: wgpu::Trace::Off,
            })
            .await?;

        let mut timestamps = None;
        if self
            .required_features
            .contains(wgpu::Features::TIMESTAMP_QUERY)
        {
            timestamps = Some(device.create_query_set(&wgpu::QuerySetDescriptor {
                label: None,
                ty: wgpu::QueryType::Timestamp,
                count: 2,
            }));
        }

        Ok(Gpu {
            instance,
            adapter,
            device,
            queue,
            #[cfg(feature = "window")]
            window: self.window,
            #[cfg(feature = "window")]
            surface: surface,
            timestamps,
        })
    }
}

pub struct Gpu {
    pub instance: wgpu::Instance,
    pub adapter: wgpu::Adapter,
    pub device: wgpu::Device,
    pub queue: wgpu::Queue,
    #[cfg(feature = "window")]
    pub window: Option<Arc<winit::window::Window>>,
    #[cfg(feature = "window")]
    pub surface: Option<wgpu::Surface<'static>>,
    pub timestamps: Option<wgpu::QuerySet>,
}

impl Gpu {
    pub fn buffer<T: bytemuck::Pod>(&self, data: &impl Buffer<T>) -> BufferBuilder<'_> {
        BufferBuilder::new(data, self)
    }

    pub fn buffer_by_sz(&self, sz: usize) -> BufferBuilder<'_> {
        BufferBuilder::new_by_sz(sz, self)
    }
}

fn get_entry_point_names(module: &naga::Module) -> anyhow::Result<Vec<String>> {
    if module.entry_points.is_empty() {
        anyhow::bail!("Expected at least one entry point, found none.");
    }
    Ok(module.entry_points.iter().map(|e| e.name.clone()).collect())
}

fn get_entry_point_workgroup_size(module: &naga::Module, name: &str) -> anyhow::Result<[u32; 3]> {
    for entry_point in module.entry_points.iter() {
        if entry_point.name != name {
            continue;
        }
        return Ok(entry_point.workgroup_size);
    }
    anyhow::bail!("Could not find entry point named {}", name);
}

/*
fn for_each_global_variable<T>(
    module: &naga::Module,
    module_info: &naga::valid::ModuleInfo,
    mut f: impl FnMut(
        naga::Handle<GlobalVariable>,
        &naga::GlobalVariable,
        &naga::ResourceBinding,
        &wgpu::ShaderStages,
        bool,
    ) -> anyhow::Result<Option<T>>,
) -> anyhow::Result<Vec<T>> {
    let mut result = Vec::new();
    for (global_handle, global) in module.global_variables.iter() {
        let Some(binding) = &global.binding else {
            continue;
        };

        let can_write = match global.space {
            naga::AddressSpace::Storage { access } => {
                access.contains(naga::StorageAccess::STORE)
                    || access.contains(naga::StorageAccess::ATOMIC)
            }
            naga::AddressSpace::Uniform => false,
            _ => {
                eprintln!(
                    "Skipping global {:?} with address space {:?}",
                    global.name, global.space
                );
                continue;
            }
        };
        let mut stages = wgpu::ShaderStages::empty();
        for (entry_idx, entry) in module.entry_points.iter().enumerate() {
            let usage = module_info.get_entry_point(entry_idx)[global_handle];
            if usage.is_empty() {
                continue;
            }

            stages |= match entry.stage {
                naga::ShaderStage::Vertex => wgpu::ShaderStages::VERTEX,
                naga::ShaderStage::Fragment => wgpu::ShaderStages::FRAGMENT,
                naga::ShaderStage::Compute => wgpu::ShaderStages::COMPUTE,
                _ => continue,
            }
        }
        if stages.is_empty() {
            eprintln!(
                "Skipping bound global {:?} which is not referenced by any stage",
                global.name
            );
            continue;
        }
        if let Some(value) = f(global_handle, global, binding, &stages, can_write)? {
            result.push(value);
        }
    }
    Ok(result)
}
*/

fn get_bind_group_layout(
    gpu: &Gpu,
    module: &naga::Module,
    module_info: &naga::valid::ModuleInfo,
    module_layout: &naga::proc::Layouter,
) -> anyhow::Result<wgpu::BindGroupLayout> {
    let mut entries = Vec::<wgpu::BindGroupLayoutEntry>::new();
    for (global_handle, global) in module.global_variables.iter() {
        let Some(binding) = &global.binding else {
            continue;
        };

        let can_write = match global.space {
            naga::AddressSpace::Storage { access } => {
                access.contains(naga::StorageAccess::STORE)
                    || access.contains(naga::StorageAccess::ATOMIC)
            }
            naga::AddressSpace::Uniform => false,
            _ => {
                eprintln!(
                    "Skipping global {:?} with address space {:?}",
                    global.name, global.space
                );
                continue;
            }
        };

        let mut stages = wgpu::ShaderStages::empty();
        for (entry_idx, entry) in module.entry_points.iter().enumerate() {
            let usage = module_info.get_entry_point(entry_idx)[global_handle];
            if usage.is_empty() {
                continue;
            }

            stages |= match entry.stage {
                naga::ShaderStage::Vertex => wgpu::ShaderStages::VERTEX,
                naga::ShaderStage::Fragment => wgpu::ShaderStages::FRAGMENT,
                naga::ShaderStage::Compute => wgpu::ShaderStages::COMPUTE,
                _ => continue,
            }
        }
        if stages.is_empty() {
            eprintln!(
                "Skipping bound global {:?} which is not referenced by any stage",
                global.name
            );
            continue;
        }

        let global_ty = &module.types[global.ty];
        let entry = match &module.types[global.ty].inner {
            naga::ir::TypeInner::Array { .. } => {
                let min_size = module_layout[global.ty].size;
                wgpu::BindGroupLayoutEntry {
                    binding: binding.binding,
                    visibility: stages,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Storage {
                            read_only: !can_write,
                        },
                        min_binding_size: NonZeroU64::new(min_size as u64),
                        has_dynamic_offset: false,
                    },
                    count: None,
                }
            }
            naga::ir::TypeInner::Struct { .. } => {
                let min_size = module_layout[global.ty].size;
                wgpu::BindGroupLayoutEntry {
                    binding: binding.binding,
                    visibility: stages,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Uniform,
                        min_binding_size: NonZeroU64::new(min_size as u64),
                        has_dynamic_offset: false,
                    },
                    count: None,
                }
            }
            _ => {
                anyhow::bail!(format!(
                    "Unsupported global type: {:?} (name={:?})",
                    global_ty.name, global_ty.inner
                ))
            }
        };
        entries.push(entry);
    }
    let layout = wgpu::BindGroupLayoutDescriptor {
        label: None,
        entries: &entries,
    };
    Ok(gpu.device.create_bind_group_layout(&layout))
}

pub struct ShaderEntryPoint {
    pub name: String,
    pub workgroup_size: [u32; 3],
}

#[derive(PartialEq, Eq)]
pub enum VertexSemantic {
    Position,
    Normal,
    Uv0,
}

pub struct VertexLayout {
    pub semantics: Vec<VertexSemantic>,
    pub attributes: Vec<wgpu::VertexAttribute>,
    pub stride: u64,
}

// Vertex data is conceived of as zero or more streams of data. We need a way
// to translate various higher-level types (f32, i32, etc) into these untyped
// streams, and finally at the end, get an array of them. This builder handles
// doing exactly that.
// Example usage:
//  let positions = [[-0.5_f32, -0.5, 0.5], [0.5, -0.5, 0.5], [0.0, 0.5, 0.5]];
//  let normals = [[0.0_f32, 0.0, 1.0]; 3];
//  let uvs = [[0.0_f32, 0.0], [1.0, 0.0], [0.5, 1.0]];
//  let data_streams = VertexDataStreamBuilder::new()
//      .stream(VertexSemantic::Position, &positions)
//      .stream(VertexSemantic::Normal, &normals)
//      .stream(VertexSemantic::Uv0, &uvs)
//      .build();

pub struct VertexDataStreamBuilder<'a> {
    streams: Vec<(VertexSemantic, &'a [u8])>,
}

impl<'a> VertexDataStreamBuilder<'a> {
    pub fn new() -> Self {
        Self {
            streams: Vec::new(),
        }
    }

    pub fn stream<T: NoUninit>(mut self, semantic: VertexSemantic, data: &'a [T]) -> Self {
        let raw_data: &[u8] = bytemuck::cast_slice(&data);
        self.streams.push((semantic, raw_data));
        self
    }

    pub fn build(self) -> Vec<(VertexSemantic, &'a [u8])> {
        self.streams
    }
}

impl VertexLayout {
    pub fn vertex_buffer_layout(&self) -> wgpu::VertexBufferLayout<'_> {
        wgpu::VertexBufferLayout {
            array_stride: self.stride,
            step_mode: wgpu::VertexStepMode::Vertex,
            attributes: &self.attributes,
        }
    }

    // Assembles streams of vertex data into an interleaved stream.
    fn interleave_data(
        &self,
        n_vertices: usize,
        data_streams: Vec<(VertexSemantic, &[u8])>,
    ) -> anyhow::Result<Vec<u8>> {
        if data_streams.len() != self.attributes.len() {
            anyhow::bail!(
                "Expected {} vertex data streams, got {}",
                self.attributes.len(),
                data_streams.len()
            );
        }
        let mut streams: Vec<_> = self
            .semantics
            .iter()
            .zip(&self.attributes)
            .map(|(semantic, attr)| {
                let (_, data) = data_streams
                    .iter()
                    .find(|(s, _)| s == semantic)
                    .expect("Vertex shader given unrecognized input semantic");
                let size = attr.format.size() as usize;
                let start = attr.offset as usize;
                (start..start + size, data.chunks_exact(size))
            })
            .collect();

        let stride = self.stride as usize;
        let mut data = vec![0; n_vertices * stride];

        for vertex in data.chunks_exact_mut(stride) {
            for (range, source) in &mut streams {
                let value = source.next().expect("Not enough vertex data");
                vertex[range.clone()].copy_from_slice(value);
            }
        }

        Ok(data)
    }

    pub fn vertex_buffer(
        &self,
        gpu: &Gpu,
        data_streams: Vec<(VertexSemantic, &[u8])>,
        n_vertices: usize,
    ) -> anyhow::Result<wgpu::Buffer> {
        let vertex_data_bytes = self.interleave_data(n_vertices, data_streams)?;
        Ok(gpu
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: None,
                contents: &vertex_data_bytes,
                usage: wgpu::BufferUsages::VERTEX,
            }))
    }
}

pub struct Shader {
    pub bind_group_layout: wgpu::BindGroupLayout,
    pub naga_module: naga::Module,
    pub shader_module: wgpu::ShaderModule,
    pub entry_points: Vec<ShaderEntryPoint>,
    uses_camera: bool,
}

impl Shader {
    pub fn new(gpu: &Gpu, shader: &str) -> anyhow::Result<Self> {
        let naga_module = naga::front::wgsl::parse_str(shader)?;
        let module_info = naga::valid::Validator::new(
            naga::valid::ValidationFlags::all(),
            naga::valid::Capabilities::all(),
        )
        .validate(&naga_module)?;
        let mut module_layout = naga::proc::Layouter::default();
        module_layout.update(naga_module.to_ctx())?;

        Self::ensure_compute_compatibility(gpu, &naga_module)?;
        let bind_group_layout =
            get_bind_group_layout(gpu, &naga_module, &module_info, &module_layout)?;
        let uses_camera = naga_module.global_variables.iter().any(|(_, global)| {
            global.space == naga::AddressSpace::Uniform
                && global.binding.as_ref().is_some_and(|binding| {
                    binding.group == 0 && binding.binding == 0
                })
        });
        let entry_point_names = get_entry_point_names(&naga_module)?;
        let entry_points = entry_point_names
            .iter()
            .map(|name| {
                let workgroup_size = get_entry_point_workgroup_size(&naga_module, name)?;
                Ok(ShaderEntryPoint {
                    name: name.clone(),
                    workgroup_size,
                })
            })
            .collect::<anyhow::Result<Vec<_>>>()?;

        let module_descriptor = wgpu::ShaderModuleDescriptor {
            label: None,
            source: wgpu::ShaderSource::Naga(Cow::Owned(naga_module.clone())),
        };
        let shader_module = gpu.device.create_shader_module(module_descriptor.clone());

        Ok(Self {
            bind_group_layout,
            naga_module: naga_module,
            shader_module,
            entry_points,
            uses_camera,
        })
    }

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

    fn ensure_compute_compatibility(gpu: &Gpu, module: &naga::Module) -> anyhow::Result<()> {
        let mut has_compute_stage = false;
        for entry_point in module.entry_points.iter() {
            if entry_point.stage == naga::ShaderStage::Compute {
                has_compute_stage = true;
            }
        }
        if !has_compute_stage {
            return Ok(());
        }
        let downlevel_caps = gpu.adapter.get_downlevel_capabilities();
        if !downlevel_caps
            .flags
            .contains(wgpu::DownlevelFlags::COMPUTE_SHADERS)
        {
            anyhow::bail!(
                "GPU {} does not support compute shaders.",
                gpu.adapter.get_info().name
            )
        }
        Ok(())
    }

    // Returns the first shader entry point for the given stage.
    pub fn get_shader_by_stage(
        &self,
        stage: naga::ir::ShaderStage,
    ) -> anyhow::Result<&naga::EntryPoint> {
        let Some(shader) = self
            .naga_module
            .entry_points
            .iter()
            .find(|ep| ep.stage == stage)
        else {
            anyhow::bail!("Shader not found for stage {:?}", stage);
        };
        Ok(shader)
    }

    fn vertex_member_to_semantic(&self, name: &str) -> anyhow::Result<VertexSemantic> {
        // Vertex members might be suffixed with _0 by slangc. Since none of our reserved names have underscores, we
        // can safely split by underscore.
        let name = name.split('_').collect::<Vec<_>>()[0];
        match name {
            "objPos" => Ok(VertexSemantic::Position),
            "objNormal" => Ok(VertexSemantic::Normal),
            "uv0" => Ok(VertexSemantic::Uv0),
            _ => anyhow::bail!("Unknown vertex member: {}", name),
        }
    }

    fn struct_member_to_vertex_format(
        &self,
        member: &naga::ir::StructMember,
    ) -> anyhow::Result<wgpu::VertexFormat> {
        let inner_type = &self.naga_module.types[member.ty].inner;
        match inner_type {
            naga::TypeInner::Vector { size, scalar } => match (scalar.kind, scalar.width) {
                (naga::ScalarKind::Float, 4) => match size {
                    naga::VectorSize::Bi => Ok(wgpu::VertexFormat::Float32x2),
                    naga::VectorSize::Tri => Ok(wgpu::VertexFormat::Float32x3),
                    naga::VectorSize::Quad => Ok(wgpu::VertexFormat::Float32x4),
                },
                _ => anyhow::bail!("Unsupported vertex member vector type: {:?}", inner_type),
            },
            _ => anyhow::bail!("Unsupported vertex member type: {:?}", inner_type),
        }
    }

    fn struct_member_shader_location(
        &self,
        member: &naga::ir::StructMember,
    ) -> anyhow::Result<wgpu::ShaderLocation> {
        match &member.binding {
            Some(naga::Binding::Location { location, .. }) => Ok(*location),
            // TODO should handle built-in bindings
            _ => anyhow::bail!("Unsupported vertex member binding: {:?}", member.binding),
        }
    }

    pub fn get_vertex_layout(&self) -> anyhow::Result<VertexLayout> {
        let vertex_shader = self.get_shader_by_stage(naga::ir::ShaderStage::Vertex)?;
        if vertex_shader.function.arguments.len() != 1 {
            anyhow::bail!(
                "Vertex shader must have exactly one argument; got {}",
                vertex_shader.function.arguments.len()
            );
        }
        let arg = &vertex_shader.function.arguments[0];
        let arg_ty = &self.naga_module.types[arg.ty];
        match &arg_ty.inner {
            naga::ir::TypeInner::Struct { members, .. } => {
                let semantics = members
                    .iter()
                    .map(|m| {
                        let Some(name) = &m.name else {
                            // This should never happen
                            anyhow::bail!("Vertex shader input member missing name");
                        };
                        self.vertex_member_to_semantic(name.as_str())
                    })
                    .collect::<anyhow::Result<Vec<_>>>()?;
                let mut acc_offset = 0;
                let attributes = members
                    .iter()
                    .map(|m| -> anyhow::Result<wgpu::VertexAttribute> {
                        let vertex_fmt = self.struct_member_to_vertex_format(m)?;
                        let location = self.struct_member_shader_location(m)?;
                        let offset = acc_offset;
                        acc_offset += vertex_fmt.size();
                        Ok(wgpu::VertexAttribute {
                            format: vertex_fmt,
                            offset,
                            shader_location: location,
                        })
                    })
                    .collect::<anyhow::Result<Vec<_>>>()?;
                let stride = attributes.iter().fold(0, |acc, a| acc + a.format.size());
                Ok(VertexLayout {
                    semantics,
                    attributes,
                    stride,
                })
            }
            _ => {
                anyhow::bail!("Argument type must be a struct; got {:?}", arg_ty.inner);
            }
        }
    }
}

pub struct MaterialPass {
    pub buffers: Vec<wgpu::Buffer>,
    pub bind_group: wgpu::BindGroup,
}

pub struct MaterialPassBuilder {
    buffers: Vec<wgpu::Buffer>,
    bind_group: Option<wgpu::BindGroup>,
}

impl MaterialPassBuilder {
    pub fn new() -> Self {
        Self {
            buffers: Vec::new(),
            bind_group: None,
        }
    }

    pub fn buffers(mut self, buffers: Vec<wgpu::Buffer>) -> Self {
        self.buffers = buffers;
        self
    }

    // Must be run after buffers(), etc.
    pub fn bind_group(mut self, gpu: &Gpu, shader: &Shader) -> Self {
        let mut descriptor_entries = Vec::<wgpu::BindGroupEntry>::new();
        let entries =
            self.buffers
                .iter()
                .enumerate()
                .fold(Vec::new(), |mut entries, (index, buf)| {
                    entries.push(wgpu::BindGroupEntry {
                        binding: index as u32,
                        resource: buf.as_entire_binding(),
                    });
                    entries
                });
        descriptor_entries.extend(entries);
        let descriptor = wgpu::BindGroupDescriptor {
            label: None,
            layout: &shader.bind_group_layout,
            entries: &descriptor_entries,
        };
        let bind_group = gpu.device.create_bind_group(&descriptor);
        self.bind_group = Some(bind_group);
        self
    }

    pub fn build(self) -> anyhow::Result<MaterialPass> {
        if self.bind_group.is_none() {
            anyhow::bail!("bind_group not set");
        }
        Ok(MaterialPass {
            buffers: self.buffers,
            bind_group: self.bind_group.unwrap(),
        })
    }
}

pub struct MaterialInstance {
    pub passes: Vec<MaterialPass>,
    pub readback: Option<wgpu::Buffer>,
}

impl MaterialInstance {
    pub fn new(passes: Vec<MaterialPass>, readback: Option<wgpu::Buffer>) -> Self {
        Self {
            passes: passes,
            readback: readback,
        }
    }
}

// A batch of Material work to be submitted together. This allows us to send a
// batch of work while another is still processing on the GPU.
pub struct MaterialBatch {
    pub instances: Vec<MaterialInstance>,
    pub ready_tx: Sender<(usize, Result<(), wgpu::BufferAsyncError>)>,
    pub ready_rx: Receiver<(usize, Result<(), wgpu::BufferAsyncError>)>,
    // Optional timestamp (ts) buffers for profiling.
    pub timing: Option<MaterialTiming>,
    // Whether the batch is currently processing. This is set true as soon as a batch is reserved with get_batch().
    // The caller is responsible for setting this false
    pub busy: bool,
    // How many instances in the batch are currently processing. Set to
    // batch_size (n_instancces) when the batch is reserved The caller may use
    // this to track completion as each instance in the batch completes.
    pub n_busy: usize,
}

impl MaterialBatch {
    // Use MaterialBuilder instead. It handles creating batches.
    fn new(
        gpu: &Gpu,
        mut passes_generator: impl FnMut() -> (Vec<MaterialPass>, Option<wgpu::Buffer>),
        n_instances: usize,
        enable_timing: bool,
    ) -> Self {
        let mut instances = Vec::<MaterialInstance>::new();
        instances.reserve(n_instances);
        for _ in 0..n_instances {
            let (passes, readback) = passes_generator();
            instances.push(MaterialInstance::new(passes, readback));
        }
        let (ready_tx, ready_rx) = std::sync::mpsc::channel();
        let mut timing = None;
        if enable_timing {
            std::hint::cold_path();
            timing = Some(MaterialTiming::new(gpu));
        }
        Self {
            instances: instances,
            ready_tx: ready_tx,
            ready_rx: ready_rx,
            timing: timing,
            busy: false,
            n_busy: 0,
        }
    }

    pub fn enable_profiling(mut self, gpu: &Gpu) -> Self {
        self.timing = Some(MaterialTiming::new(gpu));
        self
    }
}

pub struct MaterialTiming {
    pub ts_buf: wgpu::Buffer,
    pub ts_readback: wgpu::Buffer,
    pub ready_tx: Sender<Result<(), wgpu::BufferAsyncError>>,
    pub ready_rx: Receiver<Result<(), wgpu::BufferAsyncError>>,
    pub busy: bool,
}

impl MaterialTiming {
    pub fn new(gpu: &Gpu) -> Self {
        let ts_buf = gpu.buffer_by_sz(16).query_resolve().copy_src().build();
        let ts_readback = gpu.buffer_by_sz(16).copy_dst().map_read().build();
        let (ready_tx, ready_rx) = std::sync::mpsc::channel();
        Self {
            ts_buf: ts_buf,
            ts_readback: ts_readback,
            ready_tx: ready_tx,
            ready_rx: ready_rx,
            busy: false,
        }
    }
}

pub struct Material {
    pub batches: Vec<MaterialBatch>,
    pub batch_size: usize,
}

impl Material {
    pub fn get_batch(&mut self) -> Option<&mut MaterialBatch> {
        let batch = self.batches.iter_mut().find(|batch| {
            if batch.busy {
                return false;
            }
            if let Some(timing) = &batch.timing {
                std::hint::cold_path();
                if timing.busy {
                    return false;
                }
            }
            true
        })?;
        // Found a batch.
        batch.busy = true;
        if let Some(timing) = &mut batch.timing {
            std::hint::cold_path();
            timing.busy = true;
        }
        batch.n_busy = self.batch_size;
        Some(batch)
    }
}

pub struct MaterialBuilder {
    batch_size: usize, // instances per batch
    n_batches: usize,
    enable_timing: bool,
}

impl MaterialBuilder {
    pub fn new() -> Self {
        Self {
            batch_size: 1,
            n_batches: 1,
            enable_timing: false,
        }
    }

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

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

    pub fn enable_timing(mut self) -> Self {
        self.enable_timing = true;
        self
    }

    pub fn build(
        self,
        gpu: &Gpu,
        mut passes_generator: impl FnMut() -> (Vec<MaterialPass>, Option<wgpu::Buffer>),
    ) -> Material {
        let mut batches = Vec::<MaterialBatch>::new();
        batches.reserve(self.n_batches);
        for _ in 0..self.n_batches {
            batches.push(MaterialBatch::new(
                gpu,
                &mut passes_generator,
                self.batch_size,
                self.enable_timing,
            ));
        }
        Material {
            batches: batches,
            batch_size: self.batch_size,
        }
    }
}

pub struct PipelineBuilder<'a> {
    shader: &'a Shader,
    zero_initialize: bool,
    color_target: Option<wgpu::ColorTargetState>,
    vertex_buffer_layout: Option<wgpu::VertexBufferLayout<'a>>,
}

impl<'a> PipelineBuilder<'a> {
    pub fn new(shader: &'a Shader) -> Self {
        Self {
            shader,
            zero_initialize: true,
            color_target: None,
            vertex_buffer_layout: None,
        }
    }

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

    pub fn color_target(mut self, color_target: wgpu::ColorTargetState) -> Self {
        self.color_target = Some(color_target);
        self
    }

    pub fn vertex_buffer_layout(
        mut self,
        vertex_buffer_layout: wgpu::VertexBufferLayout<'a>,
    ) -> Self {
        self.vertex_buffer_layout = Some(vertex_buffer_layout);
        self
    }

    // Build a ComputePipeline. Assumes the shader has exactly one compute shader.
    pub fn build_compute(
        self,
        gpu: &Gpu,
        entry_point: &str,
    ) -> anyhow::Result<wgpu::ComputePipeline> {
        let pipeline_layout = gpu
            .device
            .create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                label: None,
                bind_group_layouts: &[Some(&self.shader.bind_group_layout)],
                immediate_size: 0,
            });

        if !self
            .shader
            .naga_module
            .entry_points
            .iter()
            .any(|ep| ep.name == entry_point && ep.stage == naga::ir::ShaderStage::Compute)
        {
            anyhow::bail!(
                "Shader entry point '{}' with compute type not found.",
                entry_point
            );
        }
        let compilation_options = wgpu::PipelineCompilationOptions {
            zero_initialize_workgroup_memory: self.zero_initialize,
            ..Default::default()
        };
        Ok(gpu
            .device
            .create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
                label: None,
                layout: Some(&pipeline_layout),
                module: &self.shader.shader_module,
                entry_point: Some(entry_point),
                compilation_options: compilation_options,
                cache: None,
            }))
    }

    // Build a RenderPipeline. Assumes the shader has exactly one vertex and one fragment shader.
    pub fn build_render(self, gpu: &Gpu) -> anyhow::Result<wgpu::RenderPipeline> {
        let pipeline_layout = gpu
            .device
            .create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                label: None,
                bind_group_layouts: &[Some(&self.shader.bind_group_layout)],
                immediate_size: 0,
            });
        let vertex_shader = self
            .shader
            .get_shader_by_stage(naga::ir::ShaderStage::Vertex)?;
        let fragment_shader = self
            .shader
            .get_shader_by_stage(naga::ir::ShaderStage::Fragment)?;
        let compilation_options = wgpu::PipelineCompilationOptions {
            zero_initialize_workgroup_memory: self.zero_initialize,
            ..Default::default()
        };
        let Some(color_target) = self.color_target else {
            anyhow::bail!("color_target is not set");
        };
        let Some(vertex_buffer_layout) = self.vertex_buffer_layout else {
            anyhow::bail!("vertex_buffer_layout is not set");
        };
        Ok(gpu
            .device
            .create_render_pipeline(&wgpu::RenderPipelineDescriptor {
                label: None,
                layout: Some(&pipeline_layout),
                vertex: wgpu::VertexState {
                    module: &self.shader.shader_module,
                    entry_point: Some(&vertex_shader.name),
                    compilation_options: compilation_options.clone(),
                    buffers: &vec![Some(vertex_buffer_layout)],
                },
                fragment: Some(wgpu::FragmentState {
                    module: &self.shader.shader_module,
                    entry_point: Some(&fragment_shader.name),
                    targets: &vec![Some(color_target)],
                    compilation_options: compilation_options,
                }),
                primitive: wgpu::PrimitiveState::default(),
                depth_stencil: None,
                multisample: wgpu::MultisampleState::default(),
                multiview_mask: None,
                cache: None,
            }))
    }
}

#[derive(Default)]
pub struct Transform {
    pub position: [f32; 3],
    pub scale: [f32; 3],
    // TODO abstract to support euler angles and quaternions
    pub orientation: [f32; 3],
}

pub struct Object {
    pub transform: Transform,
    pub vertex_buffer: wgpu::Buffer,
    // Could hoist this into Scene to have more than one mesh per pipeline, but
    // this is good enough for prototype.
    pub pipeline: wgpu::RenderPipeline,
    pub camera_bind_group: Option<wgpu::BindGroup>,
}

pub struct ObjectBuilder<'a> {
    transform: Transform,
    vertex_buffer: Option<wgpu::Buffer>,
    pipeline: Option<wgpu::RenderPipeline>,
    camera: Option<&'a Camera>,
}

impl<'a> ObjectBuilder<'a> {
    pub fn new() -> Self {
        Self {
            transform: Transform::default(),
            vertex_buffer: None,
            pipeline: None,
            camera: None,
        }
    }

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

    pub fn vertex_buffer(mut self, vertex_buffer: wgpu::Buffer) -> Self {
        self.vertex_buffer = Some(vertex_buffer);
        self
    }

    pub fn pipeline(mut self, pipeline: wgpu::RenderPipeline) -> Self {
        self.pipeline = Some(pipeline);
        self
    }

    pub fn camera(mut self, camera: &'a Camera) -> Self {
        self.camera = Some(camera);
        self
    }

    pub fn build(self, gpu: &Gpu, shader: &Shader) -> anyhow::Result<Object> {
        let Some(vertex_buffer) = self.vertex_buffer else {
            anyhow::bail!("Cannot build Object without vertex_buffer");
        };
        let Some(pipeline) = self.pipeline else {
            anyhow::bail!("Cannot build Object without pipeline");
        };
        let camera_bind_group = if shader.uses_camera() {
            Some(self.camera.context("Shader requires a camera")?.bind_group(gpu, shader))
        } else {
            None
        };
        Ok(Object {
            transform: self.transform,
            vertex_buffer,
            pipeline,
            camera_bind_group,
        })
    }
}

#[derive(Builder)]
#[builder(build_fn(skip))]
pub struct Camera {
    pub eye: glam::Vec3,
    pub target: glam::Vec3,
    pub up: glam::Vec3,
    pub vertical_fov: f32,
    pub near: f32,
    pub far: f32,
    #[builder(setter(custom))]
    pub buffer: wgpu::Buffer,
}

impl CameraBuilder {
    pub fn build(&self, gpu: &Gpu, aspect: f32) -> anyhow::Result<Camera> {
        let buffer = gpu.buffer_by_sz(64).uniform().copy_dst().build();
        let camera = Camera {
            eye: self.eye.context("eye is not set")?,
            target: self.target.context("target is not set")?,
            up: self.up.context("up is not set")?,
            vertical_fov: self.vertical_fov.context("vertical_fov is not set")?,
            near: self.near.context("near is not set")?,
            far: self.far.context("far is not set")?,
            buffer,
        };
        camera.update(gpu, aspect);
        Ok(camera)
    }
}

impl Camera {
    fn bind_group(&self, gpu: &Gpu, shader: &Shader) -> wgpu::BindGroup {
        gpu.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("Object camera bindings"),
            layout: &shader.bind_group_layout,
            entries: &[wgpu::BindGroupEntry {
                binding: 0,
                resource: self.buffer.as_entire_binding(),
            }],
        })
    }

    pub fn update(&self, gpu: &Gpu, aspect: f32) {
        let matrix = self.view_projection(aspect).to_cols_array_2d();
        gpu.queue.write_buffer(&self.buffer, 0, bytemuck::bytes_of(&matrix));
    }

    fn view_projection(&self, aspect: f32) -> glam::Mat4 {
        // right x, up y. right handed -> near z
        use glam::camera::rh::{proj, view};
        let view = view::look_at_mat4(self.eye, self.target, self.up);
        // 0-1 depth.
        let projection = proj::directx::perspective(self.vertical_fov, aspect, self.near, self.far);
        projection * view
    }
}

pub struct Scene {
    pub camera: Camera,
    pub objects: Vec<Object>,
}

pub struct SceneBuilder {
    camera: Option<Camera>,
    objects: Vec<Object>,
}

impl SceneBuilder {
    pub fn new() -> Self {
        Self {
            camera: None,
            objects: Vec::new(),
        }
    }

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

    pub fn object(mut self, object: Object) -> Self {
        self.objects.push(object);
        self
    }

    pub fn build(self) -> anyhow::Result<Scene> {
        let Some(camera) = self.camera else {
            anyhow::bail!("camera is not set");
        };
        Ok(Scene {
            camera,
            objects: self.objects,
        })
    }
}