use glam::*;
use refraction::mesh_loader::load_mesh;
use refraction::mesh_tracer::{trace_to_mesh, Mesh};
use show_image::event::{ElementState, VirtualKeyCode, WindowEvent};
use show_image::{exit, ImageInfo, ImageView, WindowOptions};
use std::env;
use std::error::Error;
use std::f32::consts::PI;
use std::fs::File;
use std::io::BufReader;
use std::sync::atomic::{AtomicUsize, Ordering::Relaxed};

const W: i32 = 320;
const H: i32 = 240;

#[derive(Copy, Clone)]
struct Color(u8, u8, u8);

struct Image {
	w: i32,
	h: i32,
	data: Vec<u8>,
}

impl Image {
	fn data(&self) -> &[u8] {
		self.data.as_slice()
	}

	fn put_pixel(&mut self, x: i32, y: i32, color: Color) {
		if x < 0 || x >= self.w || y < 0 || y > self.h {
			return;
		}
		let index = 3 * (x + self.w * y) as usize;
		self.data[index] = color.0;
		self.data[index + 1] = color.1;
		self.data[index + 2] = color.2;
	}
}

fn ypr_to_mat(ypr: Vec3) -> Mat3 {
	let Vec3 {
		x: yaw,
		y: pitch,
		z: roll,
	} = ypr;
	let m_roll = mat3(
		vec3(roll.cos(), roll.sin(), 0.0),
		vec3(-roll.sin(), roll.cos(), 0.0),
		vec3(0.0, 0.0, 1.0),
	);
	let m_yaw = mat3(
		vec3(yaw.cos(), 0.0, yaw.sin()),
		vec3(0.0, 1.0, 0.0),
		vec3(-yaw.sin(), 0.0, yaw.cos()),
	);
	let m_pitch = mat3(
		vec3(1.0, 0.0, 0.0),
		vec3(0.0, pitch.cos(), -pitch.sin()),
		vec3(0.0, pitch.sin(), pitch.cos()),
	);
	m_roll * m_pitch * m_yaw
}

fn render(mesh: &Mesh, camera: impl Fn(Vec2) -> (Vec3, Vec3)) -> Image {
	let bkg = vec3(0.0, 0.0, 0.0);
	let mut img = Image {
		w: W,
		h: H,
		data: vec![0; (3 * W * H) as usize],
	};
	let img_size = vec2(W as f32, H as f32);
	for y in 0..H {
		for x in 0..W {
			let img_coords = vec2(x as f32, y as f32);
			let off = (img_coords - img_size * 0.5) / img_size.y;
			let (base, ray) = camera(off);
			let color = if let Some(r) = trace_to_mesh(mesh, base, ray.normalize()) {
				// to_vec3(0.45) * dot(r.normal, normalize(vec3(-1.0, 1.0, -1.0))) + 0.50
				r.normal * 0.45 + 0.50
			} else {
				bkg
			};
			let color = (color * 255.0).as_ivec3().clamp(IVec3::splat(0), IVec3::splat(255));
			img.put_pixel(x, y, Color(color.x as u8, color.y as u8, color.z as u8));
		}
	}
	img
}

fn persp(dist: f32, off: Vec2) -> (Vec3, Vec3) {
	(vec3(0., 0., -dist), vec3(off.x, off.y, dist))
}

fn ortho(dist: f32, off: Vec2) -> (Vec3, Vec3) {
	(vec3(off.x, off.y, -dist), vec3(0., 0., 1.))
}

#[test]
fn test_projs() {
	fn check(f: fn(dist: f32, off: Vec2) -> (Vec3, Vec3), x: f32, y: f32, z: f32) {
		let (base, ray) = f(z, vec2(x, y));
		let at_dist = base + ray * (z / ray.z);
		assert_eq!(at_dist, vec3(x, y, 0.));
	}
	check(persp, 1., 2., 3.);
	check(ortho, 1., 2., 3.);
	check(persp, 5., 3., 7.);
	check(ortho, 9., 1., 8.);
}

// add_event_handler wants 'static + Send. Let it be so.
static PROJ_INDEX: AtomicUsize = AtomicUsize::new(0);
static PROJS: [fn(dist: f32, off: Vec2) -> (Vec3, Vec3); 2] = [persp, ortho];

#[show_image::main]
fn main() -> Result<(), Box<dyn Error>> {
	let args: Vec<String> = env::args().collect();
	if args.len() != 2 {
		println!("Usage: {} path/to/model.obj", args[0]);
		exit(1);
	}
	let mesh = {
		let f = File::open(&args[1])?;
		let mut f = BufReader::new(f);
		load_mesh(&mut f)?
	};
	let window = show_image::create_window("Raytracing", WindowOptions::default())?;
	window.add_event_handler(|_wnd, ev, _ctl| {
		if let WindowEvent::KeyboardInput(ev) = ev {
			if ev.input.state != ElementState::Pressed {
				return;
			}
			if let Some(VirtualKeyCode::Tab) = ev.input.key_code {
				PROJ_INDEX.store((PROJ_INDEX.load(Relaxed) + 1) % PROJS.len(), Relaxed);
			}
		}
	})?;
	loop {
		for phi in 0..360 {
			let proj = PROJS[PROJ_INDEX.load(Relaxed)];
			let m_view = ypr_to_mat(vec3((135.0 + phi as f32) * PI / 180.0, -30.0 * PI / 180.0, 0.0f32));
			let m_camera = m_view.transpose();
			let img = render(mesh.as_slice(), |off| {
				let (base, ray) = proj(40., 20. * off);
				(m_camera * base, m_camera * ray)
			});

			let image = ImageView::new(ImageInfo::rgb8(W as u32, H as u32), img.data());
			window.set_image("image", image)?;
		}
	}
}