use std::path::Path;

use rustc_hash::FxHashMap;

use crate::{buffer::Buffer, file::FileManager, input::InputManager, time::TimeContext};

mod wasm;

pub struct ScriptManager<'a> {
	scripts: &'a [ScriptDeclaration],
	pub active_scripts: Vec<usize>,
	script_states: FxHashMap<usize, Box<[u8]>>,
}

pub struct WasmRunner {
	pub engine: wasmi::Engine,
	pub linker: wasmi::Linker<&'static mut ScriptContext<'static>>,
}

pub struct ScriptContext<'a> {
	pub buffer: &'a mut Buffer,
	pub input: &'a InputManager,
	pub time: &'a TimeContext,
}

type StaticScriptFunction = Option<fn(&mut ScriptContext<'_>, &mut Box<[u8]>)>;

#[derive(Debug, Clone)]
pub enum ScriptDeclaration {
	Static {
		start: StaticScriptFunction,
		update: StaticScriptFunction,
		end: StaticScriptFunction,
	},
	Wasm {
		path: Box<Path>,
	},
}

impl<'a> ScriptContext<'a> {
	pub fn exit(&mut self) -> ! {
		let _ = self.buffer.finish();
		std::process::exit(0);
	}
}

fn get_script_state(states: &mut FxHashMap<usize, Box<[u8]>>, script_id: usize) -> &mut Box<[u8]> {
	states.entry(script_id).or_insert(Box::from([]))
}

impl Default for WasmRunner {
	fn default() -> Self {
		Self::new()
	}
}

impl WasmRunner {
	pub fn new() -> Self {
		let engine = wasmi::Engine::default();
		let linker = unsafe {
			std::mem::transmute::<
				wasmi::Linker<&mut ScriptContext<'_>>,
				wasmi::Linker<&mut ScriptContext<'_>>,
			>(wasm::linker(&engine).unwrap())
		};
		Self { engine, linker }
	}

	pub fn compile(&self, data: &[u8]) -> Result<wasmi::Module, wasmi::Error> {
		wasmi::Module::new(&self.engine, data)
	}

	pub fn run_function(
		&self,
		module: &wasmi::Module,
		context: &mut ScriptContext<'_>,
		function_name: &str,
	) -> Result<(), wasmi::Error> {
		// safety: the store won't be kept around after this function ends. Only the
		//         lifetimes are being transmuted
		let mut store = {
			puffin::profile_scope!("create store");
			wasmi::Store::new(&self.engine, unsafe {
				std::mem::transmute::<&mut ScriptContext<'_>, &mut ScriptContext<'_>>(context)
			})
		};
		let instance = {
			puffin::profile_scope!("instantiate");
			self.linker.instantiate_and_start(&mut store, module)?
		};
		let function = {
			puffin::profile_scope!("get function");
			instance.get_func(&store, function_name).ok_or_else(|| {
				wasmi::Error::new(format!("No function named {function_name} was exported"))
			})?
		};
		{
			puffin::profile_scope!("call function");
			function.call(&mut store, &[], &mut [])
		}
	}
}

impl<'a> ScriptManager<'a> {
	pub fn new(scripts: &'a [ScriptDeclaration]) -> Self {
		Self {
			scripts,
			active_scripts: Vec::new(),
			script_states: FxHashMap::default(),
		}
	}

	pub fn preload_script(
		&self,
		script_id: usize,
		files: &mut FileManager,
		wasm_runner: &WasmRunner,
	) {
		let script = &self.scripts[script_id];
		match script {
			ScriptDeclaration::Static { .. } => (),
			ScriptDeclaration::Wasm { path } => files.start_loading_file(path, &wasm_runner.engine),
		}
	}

	pub fn is_script_loaded(&self, script_id: usize, files: &mut FileManager) -> bool {
		let script = &self.scripts[script_id];
		match script {
			ScriptDeclaration::Static { .. } => true,
			ScriptDeclaration::Wasm { path } => files.is_file_loaded(path),
		}
	}

	pub fn unload_script(&self, script_id: usize, files: &mut FileManager) {
		let script = &self.scripts[script_id];
		match script {
			ScriptDeclaration::Static { .. } => (),
			ScriptDeclaration::Wasm { path } => files.unload_file(path),
		}
	}

	pub fn activate_script(
		&mut self,
		script_id: usize,
		context: &mut ScriptContext<'_>,
		wasm_runner: &WasmRunner,
		files: &mut FileManager,
	) {
		let Some(script_declaration) = self.scripts.get(script_id) else {
			unreachable!("The active script does not exist");
		};
		match script_declaration {
			ScriptDeclaration::Static { start, .. } => {
				start.map(|s| {
					s(
						context,
						get_script_state(&mut self.script_states, script_id),
					)
				});
			}
			ScriptDeclaration::Wasm { path } => {
				let module = files
					.load_file(path.clone(), &wasm_runner.engine)
					.unwrap()
					.unwrap_wasm_module();
				wasm_runner.run_function(module, context, "start");
			}
		}
		self.active_scripts.push(script_id);
	}

	pub fn deactivate_script(
		&mut self,
		script_id: usize,
		context: &mut ScriptContext<'_>,
		wasm_runner: &WasmRunner,
		files: &mut FileManager,
	) {
		self.active_scripts.retain(|element| {
			if *element != script_id {
				let Some(script) = self.scripts.get(script_id) else {
					unreachable!("The active script does not exist");
				};
				match script {
					ScriptDeclaration::Static { end, .. } => {
						end.map(|s| {
							s(
								context,
								get_script_state(&mut self.script_states, script_id),
							);
						});
					}
					ScriptDeclaration::Wasm { path } => {
						let module = files
							.load_file(path.clone(), &wasm_runner.engine)
							.unwrap()
							.unwrap_wasm_module();
						wasm_runner.run_function(module, context, "end");
					}
				};

				true
			} else {
				false
			}
		});
	}

	pub fn run_scripts(
		&mut self,
		context: &mut ScriptContext<'_>,
		wasm_runner: &WasmRunner,
		files: &mut FileManager,
	) {
		puffin::profile_function!();
		for script_id in &self.active_scripts {
			let Some(script) = self.scripts.get(*script_id) else {
				unreachable!("The active script does not exist");
			};
			match script {
				ScriptDeclaration::Static { update, .. } => {
					if let Some(s) = update {
						s(
							context,
							get_script_state(&mut self.script_states, *script_id),
						)
					}
				}
				ScriptDeclaration::Wasm { path } => {
					let module = {
						puffin::profile_scope!("load module");
						files
							.load_file(path.clone(), &wasm_runner.engine)
							.unwrap()
							.unwrap_wasm_module()
					};
					wasm_runner.run_function(module, context, "update");
				}
			}
		}
	}
}