wrgpgpu 0.1.5

use std::{marker::PhantomData, sync::Arc};

use image::RgbaImage;
use wgpu::util::DeviceExt;

use super::Bind;

/// Type of texture binding
pub trait TextureBindType {
	fn binding_type(format: wgpu::TextureFormat) -> wgpu::BindingType;
	fn texture_usages() -> wgpu::TextureUsages;
}

#[derive(Debug)]
/// Normal Texture Binding
pub struct TextureType;
#[derive(Debug)]
/// Read/Write Storage Texture Binding
///
/// Needs a special features in wgpu to use, and doesn't work on web
pub struct StorageReadWriteTextureType;
#[derive(Debug)]
/// Read Only Storage Texture Binding
///
/// Needs a special features in wgpu to use, and doesn't work on web
pub struct StorageReadOnlyTextureType;
#[derive(Debug)]
/// Write Storage Texture Binding
pub struct StorageWriteOnlyTextureType;

impl TextureBindType for TextureType {
	fn binding_type(_format: wgpu::TextureFormat) -> wgpu::BindingType {
		wgpu::BindingType::Texture {
			sample_type: wgpu::TextureSampleType::Float { filterable: true },
			view_dimension: wgpu::TextureViewDimension::D2,
			multisampled: false,
		}
	}
	fn texture_usages() -> wgpu::TextureUsages {
		wgpu::TextureUsages::TEXTURE_BINDING
	}
}
impl TextureBindType for StorageReadWriteTextureType {
	fn binding_type(format: wgpu::TextureFormat) -> wgpu::BindingType {
		wgpu::BindingType::StorageTexture {
			access: wgpu::StorageTextureAccess::ReadWrite,
			format,
			view_dimension: wgpu::TextureViewDimension::D2,
		}
	}
	fn texture_usages() -> wgpu::TextureUsages {
		wgpu::TextureUsages::STORAGE_BINDING
	}
}
impl TextureBindType for StorageReadOnlyTextureType {
	fn binding_type(format: wgpu::TextureFormat) -> wgpu::BindingType {
		wgpu::BindingType::StorageTexture {
			access: wgpu::StorageTextureAccess::ReadOnly,
			format,
			view_dimension: wgpu::TextureViewDimension::D2,
		}
	}
	fn texture_usages() -> wgpu::TextureUsages {
		wgpu::TextureUsages::STORAGE_BINDING
	}
}
impl TextureBindType for StorageWriteOnlyTextureType {
	fn binding_type(format: wgpu::TextureFormat) -> wgpu::BindingType {
		wgpu::BindingType::StorageTexture {
			access: wgpu::StorageTextureAccess::WriteOnly,
			format,
			view_dimension: wgpu::TextureViewDimension::D2,
		}
	}
	fn texture_usages() -> wgpu::TextureUsages {
		wgpu::TextureUsages::STORAGE_BINDING
	}
}

/// A regular texture bind, does not have a sampler
pub type PlainTextureBind<I> = TextureBind<I, TextureType>;
/// A storage texture bind, can only be written
pub type StorageTextureBind<I> = TextureBind<I, StorageWriteOnlyTextureType>;
/// A storage texture bind backed by an [`image::RgbaImage`]
pub type RgbaStorageTextureBind = TextureBind<RgbaImage, StorageWriteOnlyTextureType>;

/// Images that can be bound to the gpu
pub trait BindableImage {
	fn from_vec(width: u32, height: u32, vec: Vec<u8>) -> Self;
	fn size(&self) -> (u32, u32);
	fn as_slice(&self) -> &[u8];
	fn format() -> wgpu::TextureFormat;
}

impl BindableImage for RgbaImage {
	fn from_vec(width: u32, height: u32, vec: Vec<u8>) -> Self {
		RgbaImage::from_vec(width, height, vec).unwrap()
	}

	fn size(&self) -> (u32, u32) {
		(self.width(), self.height())
	}

	fn as_slice(&self) -> &[u8] {
		self.as_raw()
	}

	fn format() -> wgpu::TextureFormat {
		wgpu::TextureFormat::Rgba8Unorm
	}
}

/// Represents a texture binding on the gpu
pub struct TextureBind<I: BindableImage, Y: TextureBindType> {
	pub texture: Arc<wgpu::Texture>,
	view: wgpu::TextureView,
	_img: PhantomData<I>,
	_ty: PhantomData<Y>,
}

impl<I: BindableImage, Y: TextureBindType> Bind for TextureBind<I, Y> {
	type Data = I;
	type CreateInfo = (u32, u32);

	fn binding_type() -> wgpu::BindingType {
		Y::binding_type(I::format())
	}

	fn bind(&self) -> wgpu::BindingResource {
		wgpu::BindingResource::TextureView(&self.view)
	}

	fn download(&self, device: &crate::Device) -> Self::Data {
		let extent = self.texture.size();
		let download_buffer = Arc::new(device.device.create_buffer(&wgpu::BufferDescriptor {
			size: (extent.width
				* extent.height
				* extent.depth_or_array_layers
				* self.texture.format().block_copy_size(None).unwrap_or(4)) as u64,
			usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
			mapped_at_creation: false,
			label: None,
		}));

		let mut encoder = device
			.device
			.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
		encoder.copy_texture_to_buffer(
			wgpu::TexelCopyTextureInfo {
				texture: &self.texture,
				mip_level: 0,
				origin: wgpu::Origin3d { x: 0, y: 0, z: 0 },
				aspect: wgpu::TextureAspect::All,
			},
			wgpu::TexelCopyBufferInfo {
				buffer: &download_buffer,
				layout: wgpu::TexelCopyBufferLayout {
					offset: 0,
					bytes_per_row: Some(((extent.width * 4 + 255) / 256) * 256),
					rows_per_image: None,
				},
			},
			extent,
		);
		let command_buffer: wgpu::CommandBuffer = encoder.finish();
		device.queue.submit(Some(command_buffer));

		let buffer_content = Arc::new(std::sync::Mutex::new(Some(Vec::new())));
		let buffer_content_clone = buffer_content.clone();
		download_buffer
			.clone()
			.slice(..)
			.map_async(wgpu::MapMode::Read, move |result| {
				result.unwrap();
				buffer_content_clone
					.lock()
					.unwrap()
					.as_mut()
					.unwrap()
					.extend_from_slice(&download_buffer.slice(..).get_mapped_range());
			});
		device.device.poll(wgpu::Maintain::Wait);
		let data = buffer_content.lock().unwrap().take().unwrap();
		I::from_vec(extent.width, extent.height, data)
	}

	fn new_empty(device: &crate::Device, size: (u32, u32)) -> Self {
		let texture = device.device.create_texture(&wgpu::TextureDescriptor {
			label: Some(std::any::type_name::<I>()),
			size: wgpu::Extent3d {
				width: size.0,
				height: size.1,
				depth_or_array_layers: 1,
			},
			mip_level_count: 1,
			sample_count: 1,
			dimension: wgpu::TextureDimension::D2,
			format: I::format(),
			usage: wgpu::TextureUsages::COPY_SRC
				| wgpu::TextureUsages::COPY_DST
				| wgpu::TextureUsages::TEXTURE_BINDING
				| Y::texture_usages(),
			view_formats: &[],
		});
		let view = texture.create_view(&wgpu::TextureViewDescriptor {
			label: Some(std::any::type_name::<I>()),
			usage: None,
			format: Some(I::format()),
			dimension: Some(wgpu::TextureViewDimension::D2),
			aspect: wgpu::TextureAspect::All,
			base_mip_level: 0,
			mip_level_count: None,
			base_array_layer: 0,
			array_layer_count: None,
		});
		Self {
			texture: Arc::new(texture),
			view,
			_img: PhantomData,
			_ty: PhantomData,
		}
	}

	fn new_init(device: &crate::Device, inner: Self::Data) -> Self {
		let size = inner.size();
		let texture = device.device.create_texture_with_data(
			&device.queue,
			&wgpu::TextureDescriptor {
				label: Some(std::any::type_name::<I>()),
				size: wgpu::Extent3d {
					width: size.0,
					height: size.1,
					depth_or_array_layers: 1,
				},
				mip_level_count: 1,
				sample_count: 1,
				dimension: wgpu::TextureDimension::D2,
				format: I::format(),
				usage: wgpu::TextureUsages::COPY_SRC
					| wgpu::TextureUsages::COPY_DST
					| wgpu::TextureUsages::TEXTURE_BINDING
					| Y::texture_usages(),
				view_formats: &[],
			},
			wgpu::util::TextureDataOrder::LayerMajor,
			inner.as_slice(),
		);
		let view = texture.create_view(&wgpu::TextureViewDescriptor {
			label: Some(std::any::type_name::<I>()),
			usage: None,
			format: Some(I::format()),
			dimension: Some(wgpu::TextureViewDimension::D2),
			aspect: wgpu::TextureAspect::All,
			base_mip_level: 0,
			mip_level_count: None,
			base_array_layer: 0,
			array_layer_count: None,
		});
		Self {
			texture: Arc::new(texture),
			view,
			_img: PhantomData,
			_ty: PhantomData,
		}
	}
}