use std::mem::size_of;

use bytemuck::{Pod, Zeroable};

#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub struct InstanceId(usize);

#[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: u32,
}

impl Default for Instance {
	fn default() -> Self {
		Self {
			position: [0.0; 2],
			size: [1.0; 2],
			rotation: 0.0,
			z_index: 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 => Uint32
	];

	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,
		}
	}
}

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 {
	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 = device.create_buffer(&wgpu::BufferDescriptor {
			label: Some("Sprite Instance Buffer"),
			size: (instance_buffer_size * size_of::<Instance>()) as wgpu::BufferAddress,
			usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
			mapped_at_creation: false,
		});

		Self {
			instances,
			instance_buffer,
			instance_buffer_size,
		}
	}

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

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

	pub const fn buffer_size(&self) -> usize {
		self.instance_buffer_size
	}

	pub(crate) fn buffer_slice(&self) -> wgpu::BufferSlice {
		self.instance_buffer.slice(..)
	}

	pub fn push_instance(&mut self, instance: Instance) -> InstanceId {
		let index = self.instances.len();
		self.instances.push(instance);
		InstanceId(index)
	}

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

	pub fn get_instance_mut(&mut self, id: InstanceId) -> Option<&mut Instance> {
		self.instances.get_mut(id.0)
	}

	pub fn clear(&mut self) {
		self.instances.clear();
	}

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

	#[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);
		}

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