wrgpgpu 0.1.5

use std::marker::PhantomData;

use crate::Device;

/// Buffer Bindings
pub mod buffer;
/// Texture Bindings
pub mod texture;

/// Defines some data that can be put on the gpu and used in a shader
pub trait Bind {
	type Data;
	type CreateInfo;

	/// Get the binding type
	fn binding_type() -> wgpu::BindingType;
	/// Bind to the wgpu as a binding resource, to be put into a bindgroup
	fn bind(&self) -> wgpu::BindingResource;
	/// Download data from the gpu from this binding
	fn download(&self, device: &crate::Device) -> Self::Data;

	// Create a new, empty binding
	fn new_empty(device: &crate::Device, create_info: Self::CreateInfo) -> Self;
	// Create a new binding with data
	//
	// Usually, if you would like to change the data from the CPU side, a new buffer should be
	// created.
	fn new_init(device: &crate::Device, inner: Self::Data) -> Self;
}

/// Buffers & Textures associated with a bind group
pub trait BindGroupData {
	const COUNT: usize;

	fn binding_types() -> Vec<wgpu::BindingType>;
	fn binds(&self) -> Vec<wgpu::BindingResource>;
}

impl BindGroupData for () {
	const COUNT: usize = 0;

	fn binding_types() -> Vec<wgpu::BindingType> {
		vec![]
	}

	fn binds(&self) -> Vec<wgpu::BindingResource> {
		vec![]
	}
}

impl<B: Bind> BindGroupData for B {
	const COUNT: usize = 1;

	fn binding_types() -> Vec<wgpu::BindingType> {
		vec![<B as Bind>::binding_type()]
	}

	fn binds(&self) -> Vec<wgpu::BindingResource> {
		vec![<B as Bind>::bind(self)]
	}
}

impl<B1: BindGroupData, B2: BindGroupData> BindGroupData for (B1, B2) {
	const COUNT: usize = B1::COUNT + B2::COUNT;

	fn binding_types() -> Vec<wgpu::BindingType> {
		let mut types = B1::binding_types();
		types.append(&mut B2::binding_types());
		types
	}

	fn binds(&self) -> Vec<wgpu::BindingResource> {
		let mut binds = self.0.binds();
		binds.append(&mut self.1.binds());
		binds
	}
}

/// Represents a Bind Group that can be attached to the shader
pub struct BindGroup<B: BindGroupData> {
	pub bind_group: wgpu::BindGroup,
	_data: PhantomData<B>,
}

impl<B: BindGroupData> BindGroup<B> {
	pub(crate) fn get_layout(device: &Device) -> wgpu::BindGroupLayout {
		device
			.device
			.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
				label: Some(std::any::type_name::<B>()),
				entries: &B::binding_types()
					.into_iter()
					.enumerate()
					.map(|(i, ty)| wgpu::BindGroupLayoutEntry {
						binding: i as u32,
						visibility: wgpu::ShaderStages::COMPUTE,
						ty,
						count: None,
					})
					.collect::<Vec<_>>(),
			})
	}
	pub(crate) fn new(device: &Device, layout: &wgpu::BindGroupLayout, data: &B) -> Self {
		Self {
			bind_group: device.device.create_bind_group(&wgpu::BindGroupDescriptor {
				label: Some(std::any::type_name::<B>()),
				layout,
				entries: &data
					.binds()
					.into_iter()
					.enumerate()
					.map(|(i, resource)| wgpu::BindGroupEntry {
						binding: i as u32,
						resource,
					})
					.collect::<Vec<_>>(),
			}),
			_data: PhantomData,
		}
	}
}

/// Represents the set of bind groups associated with a shader
pub trait BindGroups {
	const LEN: usize;

	fn get_layouts(device: &Device) -> Vec<wgpu::BindGroupLayout>;
	fn get_bind_group_n(&self, n: usize) -> &wgpu::BindGroup;
}

impl BindGroups for () {
	const LEN: usize = 0;

	fn get_layouts(_device: &Device) -> Vec<wgpu::BindGroupLayout> {
		vec![]
	}
	fn get_bind_group_n(&self, _n: usize) -> &wgpu::BindGroup {
		panic!("No bind groups!")
	}
}

impl<B: BindGroupData> BindGroups for BindGroup<B> {
	const LEN: usize = 1;

	fn get_layouts(device: &Device) -> Vec<wgpu::BindGroupLayout> {
		vec![Self::get_layout(device)]
	}

	fn get_bind_group_n(&self, n: usize) -> &wgpu::BindGroup {
		if n == 0 {
			&self.bind_group
		} else {
			panic!("No bind group at {n} > 0!")
		}
	}
}

impl<B1: BindGroups, B2: BindGroups> BindGroups for (B1, B2) {
	const LEN: usize = B1::LEN + B2::LEN;

	fn get_layouts(device: &Device) -> Vec<wgpu::BindGroupLayout> {
		let mut layouts = B1::get_layouts(device);
		layouts.append(&mut B2::get_layouts(device));
		layouts
	}

	fn get_bind_group_n(&self, n: usize) -> &wgpu::BindGroup {
		if n < B1::LEN {
			self.0.get_bind_group_n(n)
		} else {
			self.1.get_bind_group_n(n - B1::LEN)
		}
	}
}