use std::time::Instant;

use glam::{uvec2, vec3, Vec3};
use winit::{
	event::*,
	event_loop::EventLoop,
	keyboard::{KeyCode, PhysicalKey},
	window::{Window, WindowBuilder},
};

mod camera;
mod lines;
mod scene;
mod viewport;

// The coordinate system:
// * X: forward
// * Y: left
// * Z: up

fn prepare_scene(device: &wgpu::Device) -> Vec<lines::Line> {
	scene::build()
		.into_iter()
		.map(|line| lines::Line::new_strip(device, lines::Attrs { color: line.color }, line.pts))
		.collect()
}

mod camctl {
	use glam::{vec3, Mat4, Quat, Vec3};

	pub struct CameraLocation {
		pos: Vec3,
		rot: Quat,
	}

	impl CameraLocation {
		pub fn new() -> CameraLocation {
			let rot = Quat::from_euler(glam::EulerRot::ZYX, std::f32::consts::FRAC_PI_4, 0., 0.);
			let pos = rot * vec3(-200., 0., 50.);
			CameraLocation { pos, rot }
		}

		pub fn view_mtx(&self) -> Mat4 {
			Mat4::from_quat(self.rot.inverse()) * Mat4::from_translation(-self.pos)
		}

		pub fn move_rel(&mut self, offset: Vec3) {
			self.pos += self.rot * offset;
		}

		pub fn rotate_rel_ypr(&mut self, ypr: Vec3) {
			self.rotate_rel_quat(Quat::from_euler(glam::EulerRot::XYZ, ypr.x, ypr.y, ypr.z));
		}

		pub fn rotate_rel_quat(&mut self, rot: Quat) {
			self.rot *= rot;
		}
	}
}

mod keyctl {
	use std::{collections::HashSet, iter::Sum};
	use winit::{event::ElementState, keyboard::PhysicalKey};

	pub struct Keyboard {
		pressed: HashSet<PhysicalKey>,
	}

	impl Keyboard {
		pub fn new() -> Self {
			Keyboard {
				pressed: Default::default(),
			}
		}

		pub fn is_pressed(&self, key: PhysicalKey) -> bool {
			self.pressed.contains(&key)
		}

		pub fn set_key_state(&mut self, key: PhysicalKey, state: ElementState) {
			match state {
				ElementState::Pressed => self.pressed.insert(key),
				ElementState::Released => self.pressed.remove(&key),
			};
		}

		pub fn control<T: Copy + Sum>(&self, keymap: &[(PhysicalKey, T)]) -> T {
			keymap
				.iter()
				.copied()
				.filter_map(|(key, ctl)| self.is_pressed(key).then_some(ctl))
				.sum()
		}
	}
}

static KEYS_MOVE: &'static [(PhysicalKey, Vec3)] = &[
	(PhysicalKey::Code(KeyCode::KeyW), vec3(1., 0., 0.)),
	(PhysicalKey::Code(KeyCode::KeyS), vec3(-1., 0., 0.)),
	(PhysicalKey::Code(KeyCode::KeyA), vec3(0., 1., 0.)),
	(PhysicalKey::Code(KeyCode::KeyD), vec3(0., -1., 0.)),
	(PhysicalKey::Code(KeyCode::Space), vec3(0., 0., 1.)),
	(PhysicalKey::Code(KeyCode::ShiftLeft), vec3(0., 0., -1.)),
];

static KEYS_ROTATE: &'static [(PhysicalKey, Vec3)] = &[
	(PhysicalKey::Code(KeyCode::Numpad9), vec3(1., 0., 0.)),
	(PhysicalKey::Code(KeyCode::Numpad7), vec3(-1., 0., 0.)),
	(PhysicalKey::Code(KeyCode::Numpad5), vec3(0., 1., 0.)),
	(PhysicalKey::Code(KeyCode::Numpad8), vec3(0., -1., 0.)),
	(PhysicalKey::Code(KeyCode::Numpad4), vec3(0., 0., 1.)),
	(PhysicalKey::Code(KeyCode::Numpad6), vec3(0., 0., -1.)),
];

struct State<'a> {
	device: wgpu::Device,
	queue: wgpu::Queue,

	fps: fps::Counter,
	kbd: keyctl::Keyboard,
	t1: Instant,

	viewport: viewport::Viewport<'a>,
	cam_loc: camctl::CameraLocation,
	cam_obj: camera::Camera,
	line_rend: lines::LineRenderer,

	scene: Vec<lines::Line>,

	window: &'a Window,
}

impl<'a> State<'a> {
	async fn new(window: &'a Window) -> State<'a> {
		let size = window.inner_size();

		let instance = wgpu::Instance::new(wgpu::InstanceDescriptor {
			backends: wgpu::Backends::PRIMARY,
			..Default::default()
		});
		let surface = instance.create_surface(window).unwrap();
		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 {
					label: None,
					required_features: wgpu::Features::PUSH_CONSTANTS
						| wgpu::Features::TEXTURE_ADAPTER_SPECIFIC_FORMAT_FEATURES,
					required_limits: wgpu::Limits {
						max_push_constant_size: 16,
						..wgpu::Limits::default()
					},
					memory_hints: Default::default(),
				},
				None, // Trace path
			)
			.await
			.unwrap();

		let viewport =
			viewport::Viewport::new(&adapter, &device, surface, uvec2(size.width, size.height));

		let kbd = keyctl::Keyboard::new();
		let cam_loc = camctl::CameraLocation::new();
		let t1 = Instant::now();

		let depth = None;
		let msaa = wgpu::MultisampleState {
			count: viewport.sample_count(),
			mask: !0,
			alpha_to_coverage_enabled: false,
		};

		let cam_obj = camera::Camera::new(&device);
		let line_rend = lines::LineRenderer::new(
			&device,
			cam_obj.bind_group_layout(),
			viewport.format(),
			depth,
			msaa,
		);

		let scene = prepare_scene(&device);

		let fps = fps::Counter::new();

		Self {
			device,
			queue,
			viewport,
			line_rend,
			kbd,
			fps,
			cam_loc,
			cam_obj,
			t1,
			scene,
			window,
		}
	}

	pub fn window(&self) -> &Window {
		&self.window
	}

	fn resize(&mut self, new_size: winit::dpi::PhysicalSize<u32>) {
		if new_size.width > 0 && new_size.height > 0 {
			self.viewport
				.resize(&self.device, uvec2(new_size.width, new_size.height));
		}
	}

	fn update(&mut self) {
		let dt = {
			let t2 = Instant::now();
			let dt = t2 - self.t1;
			self.t1 = t2;
			dt.as_secs_f32()
		};
		let size = self.viewport.size().as_vec2();

		self.cam_loc
			.move_rel(100. * dt * self.kbd.control(&KEYS_MOVE));
		self.cam_loc
			.rotate_rel_ypr(2. * dt * self.kbd.control(&KEYS_ROTATE));
		self.cam_obj.set(&self.queue, self.cam_loc.view_mtx(), size);
	}

	fn render(&mut self) -> Result<(), wgpu::SurfaceError> {
		self.fps.on_frame();
		self.window
			.set_title(&format!("Space Refraction ({:.1} FPS)", self.fps.get()));
		self.viewport
			.render_single_pass(&self.device, &self.queue, |mut render_pass| {
				self.line_rend.render(
					&mut render_pass,
					self.cam_obj.bind_group(),
					self.scene.iter(),
				);
			})
	}
}

pub async fn run() {
	let event_loop = EventLoop::new().unwrap();
	let window = WindowBuilder::new()
		.with_title("Refraction: Wireframe")
		.build(&event_loop)
		.unwrap();

	// State::new uses async code, so we're going to wait for it to finish
	let mut state = State::new(&window).await;
	let mut surface_configured = false;

	event_loop
		.run(move |event, control_flow| {
			match event {
				Event::WindowEvent {
					ref event,
					window_id,
				} if window_id == state.window().id() => {
					match event {
						WindowEvent::KeyboardInput {
							device_id: _,
							event,
							is_synthetic: _,
						} => {
							state.kbd.set_key_state(event.physical_key, event.state);
						}
						WindowEvent::CloseRequested => control_flow.exit(),
						WindowEvent::Resized(physical_size) => {
							surface_configured = true;
							state.resize(*physical_size);
						}
						WindowEvent::RedrawRequested => {
							// This tells winit that we want another frame after this one
							state.window().request_redraw();

							if !surface_configured {
								return;
							}

							state.update();
							match state.render() {
								Ok(_) => {}
								// Reconfigure the surface if it's lost or outdated
								Err(wgpu::SurfaceError::Lost | wgpu::SurfaceError::Outdated) => {
									state.viewport.configure(&state.device);
								}
								// The system is out of memory, we should probably quit
								Err(wgpu::SurfaceError::OutOfMemory) => {
									eprintln!("OutOfMemory");
									control_flow.exit();
								}

								// This happens when the a frame takes too long to present
								Err(wgpu::SurfaceError::Timeout) => {
									eprintln!("Surface timeout")
								}
							}
						}
						_ => {}
					}
				}
				_ => {}
			}
		})
		.unwrap();
}

fn main() {
	pollster::block_on(run());
}

mod fps {
	use std::time::{Duration, Instant};

	pub struct Counter {
		fps: f32,
		t1: Instant,
		frames: u32,
	}

	impl Counter {
		pub fn new() -> Self {
			Self {
				fps: 0.,
				t1: Instant::now(),
				frames: 0,
			}
		}

		pub fn get(&self) -> f32 {
			self.fps
		}

		pub fn on_frame(&mut self) {
			self.frames += 1;
			let t2 = Instant::now();
			let dt = t2 - self.t1;
			if dt >= Duration::from_secs(1) {
				*self = Self {
					fps: self.frames as f32 / dt.as_secs_f32(),
					t1: t2,
					frames: 0,
				}
			}
		}
	}
}