use std::mem::offset_of;

use bytemuck::{Pod, Zeroable, bytes_of, cast_slice};
use glam::{Mat4, Vec3};
use wgpu::util::DeviceExt as _;

use crate::render::{DEPTH_FORMAT, OUTPUT_FORMAT};

#[derive(Debug, Clone, Copy, Pod, Zeroable)]
#[repr(C)]
pub struct LookParams {
	pub m: Mat4,
}

#[derive(Debug, Clone, Copy, Pod, Zeroable)]
#[repr(C)]
pub struct Vertex {
	pub pos: Vec3,
	pub color: Vec3,
}

impl Vertex {
	pub fn new(pos: Vec3, color: Vec3) -> Self {
		Self { pos, color }
	}
}

pub struct Mesh {
	vertex_buffer: wgpu::Buffer,
	index_buffer: wgpu::Buffer,
	index_count: u32,
}

impl Mesh {
	pub fn new(device: &wgpu::Device, vertices: &[Vertex], indices: &[u16]) -> Self {
		if vertices.len() >= 1 << 16 {
			panic!("too many vertices");
		}
		for index in indices {
			if *index as usize >= vertices.len() {
				panic!(
					"vertex index out of bounds: {index} out of {}",
					vertices.len()
				);
			}
		}
		let vertex_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
			label: None,
			usage: wgpu::BufferUsages::VERTEX,
			contents: cast_slice(vertices),
		});
		let index_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
			label: None,
			usage: wgpu::BufferUsages::INDEX,
			contents: cast_slice(indices),
		});
		Self {
			vertex_buffer,
			index_buffer,
			index_count: indices.len().try_into().expect("too many indices"),
		}
	}
}

pub struct Pipeline {
	look_buf: wgpu::Buffer,
	bindings: wgpu::BindGroup,
	pipeline: wgpu::RenderPipeline,
}

impl Pipeline {
	pub fn new(device: &wgpu::Device) -> Self {
		let look_buf = device.create_buffer(&wgpu::BufferDescriptor {
			label: None,
			size: size_of::<LookParams>() as u64,
			usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
			mapped_at_creation: false,
		});

		let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
			label: None,
			source: wgpu::ShaderSource::Wgsl(super::COLORMAP_SHADER.into()),
		});
		let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
			label: None,
			layout: None,
			vertex: wgpu::VertexState {
				module: &shader,
				entry_point: None,
				compilation_options: wgpu::PipelineCompilationOptions::default(),
				buffers: &[wgpu::VertexBufferLayout {
					array_stride: size_of::<Vertex>() as u64,
					step_mode: wgpu::VertexStepMode::Vertex,
					attributes: &[
						wgpu::VertexAttribute {
							shader_location: 0,
							offset: offset_of!(Vertex, pos) as u64,
							format: wgpu::VertexFormat::Float32x3,
						},
						wgpu::VertexAttribute {
							shader_location: 1,
							offset: offset_of!(Vertex, color) as u64,
							format: wgpu::VertexFormat::Float32x3,
						},
					],
				}],
			},
			primitive: wgpu::PrimitiveState {
				topology: wgpu::PrimitiveTopology::TriangleList,
				..Default::default()
			},
			depth_stencil: Some(wgpu::DepthStencilState {
				format: DEPTH_FORMAT,
				depth_write_enabled: true,
				depth_compare: wgpu::CompareFunction::LessEqual,
				stencil: wgpu::StencilState::default(),
				bias: wgpu::DepthBiasState::default(),
			}),
			multisample: wgpu::MultisampleState {
				count: 1,
				mask: !0,
				alpha_to_coverage_enabled: false,
			},
			fragment: Some(wgpu::FragmentState {
				module: &shader,
				entry_point: None,
				compilation_options: wgpu::PipelineCompilationOptions::default(),
				targets: &[Some(wgpu::ColorTargetState {
					format: OUTPUT_FORMAT,
					blend: Some(wgpu::BlendState::REPLACE),
					write_mask: wgpu::ColorWrites::ALL,
				})],
			}),
			multiview: None,
			cache: None,
		});
		let bindings = device.create_bind_group(&wgpu::BindGroupDescriptor {
			label: None,
			layout: &pipeline.get_bind_group_layout(0),
			entries: &[wgpu::BindGroupEntry {
				binding: 0,
				resource: look_buf.as_entire_binding(),
			}],
		});
		Self {
			look_buf,
			bindings,
			pipeline,
		}
	}

	pub fn set_look(&self, queue: &wgpu::Queue, look: LookParams) {
		queue.write_buffer(&self.look_buf, 0, bytes_of(&look));
	}

	pub fn render<'a>(
		&self,
		pass: &mut wgpu::RenderPass,
		meshes: impl IntoIterator<Item = &'a Mesh>,
	) {
		pass.set_pipeline(&self.pipeline);
		pass.set_bind_group(0, &self.bindings, &[]);
		for mesh in meshes {
			pass.set_vertex_buffer(0, mesh.vertex_buffer.slice(..));
			pass.set_index_buffer(mesh.index_buffer.slice(..), wgpu::IndexFormat::Uint16);
			pass.draw_indexed(0..mesh.index_count, 0, 0..1);
		}
	}
}