#![feature(gen_blocks)]

use std::{error::Error, f32::consts::PI};

use glam::{Mat4, UVec2, Vec2, Vec3, vec3};
use rand_distr::Distribution as _;
use wgpu::rwh::{RawDisplayHandle, RawWindowHandle};

use crate::{
	camera::OrbitalCamera,
	iter::loop_list,
	ray::Ray,
	render::{
		DEPTH_FORMAT, OUTPUT_FORMAT,
		lines::{LookParams, Mesh, Pipeline, Vertex},
	},
	trace::{Hit, Lambertian, Reflector, Scene, Source, Sphere},
};

pub use crate::ffi::{RedrawArgs, SphericalPosition, UseNormal};

mod camera;
mod iter;
mod ray;
mod render;
mod shape;
mod trace;

#[diplomat::bridge]
mod ffi {
	use std::{f32::consts::PI, num::NonZero, ptr::NonNull};

	use glam::uvec2;
	use wgpu::rwh::{XcbDisplayHandle, XcbWindowHandle};

	#[derive(Debug, Clone, Copy, PartialEq)]
	pub struct SphericalPosition {
		pub yaw: f32,
		pub pitch: f32,
		pub distance: f32,
	}

	#[derive(Debug)]
	pub enum UseNormal {
		Light,
		Camera,
	}

	#[derive(Debug, Clone, Copy)]
	pub struct RedrawArgs {
		pub camera_position: SphericalPosition,
		pub light_position: SphericalPosition,
		pub light_radius: f32,
		pub light_spread: f32,
		pub accum_sigma: f32,
		pub accum_scale: f32,
		pub reflections: u8,
		pub use_normal: UseNormal,
		pub show_axes: bool,
		pub show_shapes: bool,
		pub show_hit_emission: bool,
		pub show_miss_emission: bool,
		pub show_direct_hit: bool,
		pub show_indirect_hit: bool,
		pub show_light: bool,
	}

	impl Default for RedrawArgs {
		fn default() -> Self {
			Self {
				camera_position: SphericalPosition {
					yaw: PI / 4.,
					pitch: 0.5f32.sqrt().atan(),
					distance: 3.0,
				},
				light_position: SphericalPosition {
					yaw: 0.0,
					pitch: PI / 3.,
					distance: 1.0,
				},
				light_radius: 0.125,
				light_spread: 0.125,
				accum_sigma: 0.025,
				accum_scale: 0.01,
				reflections: 2,
				use_normal: UseNormal::Light,
				show_axes: true,
				show_shapes: true,
				show_hit_emission: false,
				show_miss_emission: false,
				show_direct_hit: false,
				show_indirect_hit: false,
				show_light: true,
			}
		}
	}

	#[diplomat::opaque_mut]
	pub struct Core(crate::Core);

	impl Core {
		pub fn from_xcb(connection: usize, window: u32, width: u32, height: u32) -> Box<Core> {
			let display =
				XcbDisplayHandle::new(Some(NonNull::new(connection as *mut _).unwrap()), 0);
			let window = XcbWindowHandle::new(NonZero::new(window).unwrap());
			let gpu = unsafe { crate::Gpu::new_raw(display, window) };
			let size = uvec2(width, height);
			Box::new(Self(crate::Core::new(gpu, size)))
		}

		pub fn configure(&mut self, width: u32, height: u32) {
			self.0.configure(uvec2(width, height));
		}

		pub fn redraw(&mut self, args: &RedrawArgs) {
			self.0.redraw(args);
		}
	}
}

pub struct Gpu {
	device: wgpu::Device,
	queue: wgpu::Queue,
	surface: wgpu::Surface<'static>,
}

impl Gpu {
	unsafe fn new_raw(
		display: impl Into<RawDisplayHandle>,
		window: impl Into<RawWindowHandle>,
	) -> Self {
		let target = wgpu::SurfaceTargetUnsafe::RawHandle {
			raw_display_handle: display.into(),
			raw_window_handle: window.into(),
		};
		let gpu = pollster::block_on(init_gpu_inner(|instance| unsafe {
			instance.create_surface_unsafe(target)
		}))
		.unwrap();
		gpu
	}
}

pub struct Core {
	device: wgpu::Device,
	queue: wgpu::Queue,
	surface: wgpu::Surface<'static>,

	depth: wgpu::Texture,

	pipeline: Pipeline,
	mesh_pipe: render::faces::Pipeline,
	tripod: Mesh,
}

pub fn new_tripod(device: &wgpu::Device) -> Mesh {
	Mesh::new(
		device,
		&[
			Vertex::new(vec3(0., 0., 0.), vec3(1., 0., 0.)),
			Vertex::new(vec3(1., 0., 0.), vec3(1., 0., 0.)),
			Vertex::new(vec3(0., 0., 0.), vec3(0., 1., 0.)),
			Vertex::new(vec3(0., 1., 0.), vec3(0., 1., 0.)),
			Vertex::new(vec3(0., 0., 0.), vec3(0., 0., 1.)),
			Vertex::new(vec3(0., 0., 1.), vec3(0., 0., 1.)),
		],
	)
}

impl OrbitalCamera {
	fn make_ray(&self, rng: &mut impl rand::Rng) -> Ray {
		let off: f32 = rand_distr::StandardUniform.sample(rng);
		let side: Vec2 = rand_distr::UnitCircle.sample(rng).into();

		let m = self.transform().inverse();
		let fwd = 1. - 0.1 * off;
		let side_scale = (1. - fwd.powi(2)).sqrt();
		let dir = Vec3::from((side_scale * side, fwd));
		Ray {
			base: self.position(),
			dir: m.transform_vector3(dir),
		}
	}
}

impl Core {
	pub fn new(gpu: Gpu, pixel_size: UVec2) -> Self {
		let Gpu {
			device,
			queue,
			surface,
		} = gpu;

		let pipeline = Pipeline::new(&device);
		let mesh_pipe = render::faces::Pipeline::new(&device);
		let tripod = new_tripod(&device);
		let depth = Self::create_depth_buffer(&device, pixel_size);
		queue.submit([]); // flush buffer updates

		Self::configure_surface(&surface, &device, pixel_size);

		Self {
			device,
			queue,
			surface,
			depth,
			pipeline,
			tripod,
			mesh_pipe,
		}
	}

	fn render(&self, output: &wgpu::Texture, args: &RedrawArgs) {
		let camera = OrbitalCamera {
			position_yaw: args.camera_position.yaw,
			position_pitch: args.camera_position.pitch,
			distance: args.camera_position.distance,
		};
		let aspect = {
			let size = output.size();
			let w = size.width as f32;
			let h = size.height as f32;
			w / h
		};
		let perspective = Mat4::perspective_lh(PI / 3., aspect, 1e-2, 1e2);
		self.pipeline.set_look(
			&self.queue,
			LookParams {
				m: perspective * camera.transform(),
			},
		);
		self.mesh_pipe.set_look(
			&self.queue,
			render::faces::LookParams {
				m: perspective * camera.transform(),
			},
		);
		self.queue.submit([]); // flush buffer updates

		let view = output.create_view(&wgpu::TextureViewDescriptor::default());
		let depth_view = self
			.depth
			.create_view(&wgpu::TextureViewDescriptor::default());
		let mut encoder = self
			.device
			.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
		let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
			color_attachments: &[Some(wgpu::RenderPassColorAttachment {
				view: &view,
				depth_slice: None,
				resolve_target: None,
				ops: wgpu::Operations {
					load: wgpu::LoadOp::Clear(wgpu::Color::BLACK),
					store: wgpu::StoreOp::Store,
				},
			})],
			depth_stencil_attachment: Some(wgpu::RenderPassDepthStencilAttachment {
				view: &depth_view,
				depth_ops: Some(wgpu::Operations {
					load: wgpu::LoadOp::Clear(1.),
					store: wgpu::StoreOp::Discard,
				}),
				stencil_ops: None,
			}),
			..Default::default()
		});
		if args.show_axes {
			self.pipeline.render(&mut pass, [&self.tripod]);
		}

		let source = Source {
			position_yaw: args.light_position.yaw,
			position_pitch: args.light_position.pitch,
			distance: args.light_position.distance,
			radius: args.light_radius,
			spread: args.light_spread,
		};

		if args.show_shapes {
			let contour: Vec<Vertex> = loop_list(source.contour(17))
				.map(|pos| Vertex {
					pos,
					color: vec3(1., 1., 1.),
				})
				.collect();
			self.pipeline
				.render(&mut pass, [&Mesh::new(&self.device, &contour)]);
		}

		const BASE_R: f32 = 2.;
		const BASE_POS: Vec3 = vec3(0., 0., -BASE_R);
		const BASE: Sphere = Sphere {
			position: vec3(0., 0., -BASE_R),
			radius: BASE_R,
			color: vec3(0.01, 0.03, 0.30),
		};
		fn sphere(pos: Vec3, color: Vec3) -> Sphere {
			Sphere {
				position: pos,
				radius: BASE_POS.distance(pos) - BASE_R,
				color,
			}
		}
		let scene = Scene {
			objects: vec![
				BASE,
				sphere(vec3(0., 0., 0.1), vec3(1., 1., 1.)),
				sphere(vec3(0.3, 0., 0.1), vec3(1., 1., 1.)),
				sphere(vec3(0.1, 0.3, 0.1), vec3(0.5, 0.03, 0.01)),
			],
		};

		#[derive(Debug, Clone, Copy)]
		struct LightHit<'a> {
			incident: Ray,
			normal: Vec3,
			light: Vec3,
			object: &'a Sphere,
		}
		let mut prng = rand_pcg::Pcg64::new(42, 0);
		let source_rays: Vec<Ray> = (0..10240).map(|_| source.make_ray(&mut prng)).collect();
		let camera_rays: Vec<Ray> = (0..10240).map(|_| camera.make_ray(&mut prng)).collect();
		let mut source_ray_display: Vec<Vertex> = Vec::with_capacity(source_rays.len());
		let mut hits: Vec<LightHit> = Vec::with_capacity(source_rays.len());
		for ray in source_rays {
			let light = Vec3::splat(1.);
			if let Some(hit) = scene.trace_ray(ray) {
				hits.push(LightHit {
					incident: hit.incident,
					normal: hit.normal,
					light,
					object: hit.object,
				});
				if args.show_hit_emission {
					source_ray_display.extend([
						Vertex {
							pos: ray.base,
							color: vec3(1., 1., 1.),
						},
						Vertex {
							pos: ray.base + 0.1 * ray.dir,
							color: vec3(0., 1., 0.),
						},
					]);
				}
				if args.show_direct_hit {
					source_ray_display.extend([
						Vertex {
							pos: hit.incident.base - 0.02 * hit.incident.dir,
							color: light,
						},
						Vertex {
							pos: hit.incident.base,
							color: vec3(1., 1., 1.),
						},
					]);
				}
			} else {
				if args.show_miss_emission {
					source_ray_display.extend([
						Vertex {
							pos: ray.base,
							color: vec3(1., 1., 1.),
						},
						Vertex {
							pos: ray.base + 0.1 * ray.dir,
							color: vec3(1., 0., 0.),
						},
					]);
				}
			}
		}
		if args.reflections > 0 {
			let mut hits1 = hits.clone();
			for _ in 0..args.reflections {
				let mut hits2: Vec<LightHit> = Vec::with_capacity(hits1.len());
				for hit in &hits1 {
					let reflector = Lambertian;
					let reflected = reflector.reflect(&mut prng, hit.normal, hit.incident.dir);
					let light = hit.light * hit.object.color;
					let ray = Ray::new(hit.incident.base, reflected);
					let Some(hit2) = scene.trace_ray(ray) else {
						continue;
					};
					hits2.push(LightHit {
						incident: hit2.incident,
						normal: hit2.normal,
						light: light,
						object: hit2.object,
					});
					if args.show_indirect_hit {
						source_ray_display.extend([
							Vertex {
								pos: hit2.incident.base - 0.02 * hit2.incident.dir,
								color: light,
							},
							Vertex {
								pos: hit2.incident.base,
								color: vec3(1., 1., 1.),
							},
						]);
					}
				}
				hits.extend(&hits2);
				hits1 = hits2;
			}
		}
		let light_at = {
			let sigma2 = args.accum_sigma.powi(2);
			let accum_normalizator = (2. * PI * sigma2).recip();
			let hits = &hits;
			move |hit: Hit| -> Vec3 {
				let mut total_cd = Vec3::splat(0.);
				for light_hit in hits {
					let d2 = hit.incident.base.distance_squared(light_hit.incident.base);
					if d2 > 9. * sigma2 {
						continue;
					}
					let normal = match args.use_normal {
						UseNormal::Light => light_hit.normal,
						UseNormal::Camera => hit.normal,
					};
					assert!(normal.is_normalized());
					assert!(hit.incident.dir.is_normalized());
					let reflector = Lambertian;
					let in_lm = light_hit.light;
					let out_cd = in_lm
						* reflector.brdf(normal, light_hit.incident.dir, -hit.incident.dir)
						* hit.object.color;
					let weight = accum_normalizator * (-0.5 * d2 / sigma2).exp();
					total_cd += weight * out_cd;
				}
				total_cd * args.accum_scale
			}
		};
		let colormap = |light: Vec3| {
			let light = light.dot(Vec3::splat(1. / 3.));
			let brightness = 3. * (1. - (1. + light).recip());
			vec3(brightness, brightness - 1., brightness - 2.)
				.clamp(Vec3::splat(0.), Vec3::splat(1.))
		};
		let mut camera_ray_display: Vec<Vertex> = Vec::with_capacity(camera_rays.len());
		if args.show_light {
			let r = args.accum_sigma;
			for ray in camera_rays {
				let Some(hit) = scene.trace_ray(ray) else {
					continue;
				};
				let light = light_at(hit);
				let color = colormap(light);
				let vertex = |off: Vec3| Vertex {
					pos: hit.incident.base + r * off,
					color,
				};
				camera_ray_display.extend([
					vertex(-Vec3::X),
					vertex(Vec3::X),
					vertex(-Vec3::Y),
					vertex(Vec3::Y),
					vertex(-Vec3::Z),
					vertex(Vec3::Z),
				]);
			}
		}
		if args.show_shapes {
			let mut meshes = Vec::new();
			for obj in &scene.objects {
				let mesh = shape::octo_sphere((obj.radius * 45.) as usize + 7);
				let obj_mesh = render::faces::Mesh::new(
					&self.device,
					&mesh
						.vertices
						.iter()
						.map(|&v| {
							let pos = obj.position + obj.radius * v;
							let normal = v;
							let mut dir = (pos - camera.position()).normalize();
							if !dir.is_finite() {
								dir = camera.direction();
							}
							let color = light_at(Hit {
								incident: Ray::new(pos, dir),
								normal,
								object: obj,
							});
							render::faces::Vertex { pos, color }
						})
						.collect::<Vec<_>>(),
					bytemuck::cast_slice(&mesh.indices),
				);
				meshes.push(obj_mesh);
			}
			self.mesh_pipe.render(&mut pass, &meshes);
		}
		if !source_ray_display.is_empty() {
			self.pipeline
				.render(&mut pass, [&Mesh::new(&self.device, &source_ray_display)]);
		}
		if !camera_ray_display.is_empty() {
			self.pipeline
				.render(&mut pass, [&Mesh::new(&self.device, &camera_ray_display)]);
		}

		drop(pass);
		self.queue.submit(std::iter::once(encoder.finish()));
	}

	fn create_depth_buffer(device: &wgpu::Device, pixel_size: UVec2) -> wgpu::Texture {
		device.create_texture(&wgpu::TextureDescriptor {
			label: Some("depth buffer"),
			size: wgpu::Extent3d {
				width: pixel_size.x,
				height: pixel_size.y,
				depth_or_array_layers: 1,
			},
			mip_level_count: 1,
			sample_count: 1,
			dimension: wgpu::TextureDimension::D2,
			format: DEPTH_FORMAT,
			usage: wgpu::TextureUsages::RENDER_ATTACHMENT,
			view_formats: &[],
		})
	}

	fn configure_surface(surface: &wgpu::Surface, device: &wgpu::Device, pixel_size: UVec2) {
		surface.configure(
			device,
			&wgpu::SurfaceConfiguration {
				usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_DST,
				format: OUTPUT_FORMAT,
				width: pixel_size.x,
				height: pixel_size.y,
				present_mode: wgpu::PresentMode::Fifo,
				alpha_mode: wgpu::CompositeAlphaMode::Auto,
				view_formats: vec![],
				desired_maximum_frame_latency: 2,
			},
		);
	}

	/// Configures the renderer for a given target size.
	pub fn configure(&mut self, pixel_size: UVec2) {
		Self::configure_surface(&self.surface, &self.device, pixel_size);
		self.depth = Self::create_depth_buffer(&self.device, pixel_size);
	}

	/// Redraws the entire surface.
	///
	/// [`Self::configure`] must be called at least once before this.
	pub fn redraw(&mut self, args: &RedrawArgs) {
		let output = self.surface.get_current_texture().unwrap();
		self.render(&output.texture, args);
		output.present();
	}
}

pub async fn init_gpu_inner<E: Error + 'static>(
	make_surface: impl FnOnce(&wgpu::Instance) -> Result<wgpu::Surface<'static>, E>,
) -> Result<Gpu, Box<dyn Error>> {
	let instance = wgpu::Instance::new(&wgpu::InstanceDescriptor {
		backends: wgpu::Backends::PRIMARY,
		..Default::default()
	});
	let surface = make_surface(&instance)?;
	let adapter = instance
		.request_adapter(&wgpu::RequestAdapterOptions {
			power_preference: wgpu::PowerPreference::default(),
			compatible_surface: Some(&surface),
			force_fallback_adapter: false,
		})
		.await
		.unwrap();
	let (device, queue) = adapter
		.request_device(&wgpu::DeviceDescriptor::default())
		.await
		.unwrap();
	Ok(Gpu {
		device,
		queue,
		surface,
	})
}