use std::mem::size_of;

use bytemuck::{Pod, Zeroable};

/// The ID for an instance
#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub struct InstanceId(usize);

/// A sprite, that can be used by the alligator shader
#[repr(C)]
#[derive(Copy, Clone, Debug, PartialEq, Pod, Zeroable)]
pub struct Instance {
	/// Position on the screen
	pub position: [f32; 2],
	/// Relative size
	pub size: [f32; 2],
	/// The location of the texture in the texture atlas
	pub texture_coordinates: [f32; 2],
	/// The size of the sprite's texture
	pub texture_size: [f32; 2],
	/// The index of the texture atlas to use
	pub texture_atlas_index: u32,
	/// Rotation, in radians
	pub rotation: f32,
	/// z-index
	pub z_index: f32,
}

impl Default for Instance {
	fn default() -> Self {
		Self {
			position: [0.0; 2],
			size: [1.0; 2],
			rotation: 0.0,
			z_index: 0.0,
			texture_coordinates: [0.0; 2],
			texture_size: [1.0; 2],
			texture_atlas_index: 0,
		}
	}
}

impl Instance {
	// whenever this is updated, please also update `sprite.wgsl`
	const ATTRIBUTES: [wgpu::VertexAttribute; 7] = wgpu::vertex_attr_array![
		1 => Float32x2, 2 => Float32x2, 3 => Float32x2, 4 => Float32x2,
		5 => Uint32, 6 => Float32, 7 => Float32
	];

	pub(crate) fn desc<'a>() -> wgpu::VertexBufferLayout<'a> {
		// make sure these two don't conflict
		debug_assert_eq!(
			Self::ATTRIBUTES[0].shader_location as usize,
			crate::Vertex::ATTRIBUTES.len()
		);
		wgpu::VertexBufferLayout {
			array_stride: size_of::<Self>() as wgpu::BufferAddress,
			step_mode: wgpu::VertexStepMode::Instance,
			attributes: &Self::ATTRIBUTES,
		}
	}
}

/// A buffer of sprites, for both CPU and GPU memory
pub struct InstanceBuffer {
	instances: Vec<Instance>,
	instance_buffer: wgpu::Buffer,
	instance_buffer_size: usize,
}

fn create_buffer(device: &wgpu::Device, instances: &Vec<Instance>) -> wgpu::Buffer {
	device.create_buffer(&wgpu::BufferDescriptor {
		label: Some("Sprite Instance Buffer"),
		size: (instances.capacity() * size_of::<Instance>()) as wgpu::BufferAddress,
		usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
		mapped_at_creation: false,
	})
}

impl InstanceBuffer {
	/// Create a new buffer with the given capacity
	pub(crate) fn new(device: &wgpu::Device, capacity: usize) -> Self {
		let instances = Vec::with_capacity(capacity);
		let instance_buffer_size = instances.capacity();
		let instance_buffer = create_buffer(device, &instances);

		Self {
			instances,
			instance_buffer,
			instance_buffer_size,
		}
	}

	/// The number of sprites
	pub fn len(&self) -> u32 {
		self.instances
			.len()
			.try_into()
			.expect("expected less than 3 billion instances")
	}

	/// Returns `true` if there are no sprites
	pub fn is_empty(&self) -> bool {
		self.instances.is_empty()
	}

	/// The capacity of the buffer
	pub const fn buffer_size(&self) -> usize {
		self.instance_buffer_size
	}

	/// Get a slice containing the entire buffer
	pub(crate) fn buffer_slice(&self) -> wgpu::BufferSlice {
		self.instance_buffer.slice(..)
	}

	/// Add a new sprite. The new sprite's `InstanceId` is returned. This ID
	/// becomes invalid if the buffer is cleared.
	pub fn push_instance(&mut self, instance: Instance) -> InstanceId {
		let index = self.instances.len();
		self.instances.push(instance);
		InstanceId(index)
	}

	/// Get a specific instance
	pub fn get_instance(&self, id: InstanceId) -> Option<&Instance> {
		self.instances.get(id.0)
	}

	/// Get a mutable reference to a specific sprite
	pub fn get_instance_mut(&mut self, id: InstanceId) -> Option<&mut Instance> {
		self.instances.get_mut(id.0)
	}

	/// Clear the instance buffer. This invalidates all `InstanceId`'s
	pub fn clear(&mut self) {
		self.instances.clear();
	}

	/// Increase the capacity of the buffer
	fn expand_buffer(&mut self, device: &wgpu::Device) {
		self.instance_buffer_size = self.instances.capacity();
		self.instance_buffer = create_buffer(device, &self.instances);
	}

	/// Fill the GPU buffer with the sprites in the CPU buffer.
	#[profiling::function]
	pub(crate) fn fill_buffer(&mut self, device: &wgpu::Device, queue: &wgpu::Queue) {
		if self.instances.len() > self.instance_buffer_size {
			self.expand_buffer(device);
		}

		// the instances must be sorted by z-index before being handed to the GPU
		let mut sorted = self.instances.clone();
		sorted.sort_by(|a, b| a.z_index.total_cmp(&b.z_index));

		queue.write_buffer(
			&self.instance_buffer,
			0 as wgpu::BufferAddress,
			bytemuck::cast_slice(&sorted),
		);
	}
}