use std::{path::Path, sync::Arc};

use rustc_hash::FxHashMap;

use crate::{
	buffer::{Buffer, Cell},
	file::{FileManager, GameManifest},
	input::InputManager,
	scene::{preload_scene, switch_scenes},
	script::wasm::WasmModule,
	storage::StorageManager,
	time::TimeContext,
};

pub mod lua;
pub mod wasm;

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

#[derive(Clone)]
pub struct WasmRunner {
	pub engine: tinywasm::Engine,
	pub store: Arc<spin::Mutex<tinywasm::Store>>,
	pub linker: tinywasm::Imports,
}

pub struct ScriptContext<'a> {
	pub manifest: &'a GameManifest,
	pub buffer: &'a mut Buffer,
	pub input: &'a InputManager,
	pub time: &'a TimeContext,
	pub storage: &'a mut StorageManager,
	pub operation_queue: Vec<Operation>,
}

pub struct FullScriptContext<'a> {
	pub script_id: usize,
	pub manifest: &'a GameManifest,
	pub buffer: &'a Buffer,
	pub input: &'a InputManager,
	pub time: &'a TimeContext,
	pub scripts: &'a ScriptManager<'a>,
	pub files: Box<[Box<Path>]>,
	pub storage: &'a StorageManager,
	pub wasm_runner: &'a WasmRunner,
	pub operation_queue: &'a mut Vec<Operation>,
}

pub enum Operation {
	AddCell(Cell),
	PreloadScene(usize),
	SwitchScene(usize),
	PreloadScript(usize),
	UnloadScript(usize),
	ActivateScript(usize),
	DeactivateScript(usize),
	UpdateState {
		script_id: usize,
		new_state: Box<[u8]>,
	},
	UpdateEntity {
		name: Box<str>,
		new_state: Box<[u8]>,
	},
	UpdateStorage {
		new_save: Box<[u8]>,
	},
	Exit,
}

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

#[derive(Debug, Clone)]
pub enum ScriptDeclaration {
	Static {
		name: Box<str>,
		start: StaticScriptFunction,
		update: StaticScriptFunction,
		end: StaticScriptFunction,
	},
	Lua {
		name: Box<str>,
		path: Box<Path>,
	},
	Wasm {
		name: Box<str>,
		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 = tinywasm::Engine::default();
		let mut store = tinywasm::Store::new(engine.clone());
		let linker = wasm::linker(&mut store).unwrap();
		let store = Arc::new(spin::Mutex::new(store));
		Self {
			engine,
			store,
			linker,
		}
	}

	pub fn compile(&self, data: &[u8]) -> Result<WasmModule, tinywasm::Error> {
		let module = tinywasm::parse_bytes(data)?;
		let instance = tinywasm::ModuleInstance::instantiate(
			&mut self.store.lock(),
			&module,
			Some(self.linker.clone()),
		)?;
		Ok(WasmModule { instance })
	}

	#[expect(clippy::too_many_arguments)]
	pub fn run_function(
		&self,
		script_id: usize,
		module: &mut WasmModule,
		scripts: &ScriptManager,
		files: Box<[Box<Path>]>,
		wasm_runner: &WasmRunner,
		context: &mut ScriptContext<'_>,
		function_name: &str,
	) -> Result<(), tinywasm::Error> {
		puffin::profile_function!();

		let context = std::ptr::from_mut(Box::leak(Box::new(FullScriptContext {
			script_id,
			manifest: context.manifest,
			buffer: context.buffer,
			input: context.input,
			time: context.time,
			storage: context.storage,
			scripts,
			files,
			wasm_runner,
			operation_queue: &mut context.operation_queue,
		})));

		let function = module
			.instance
			.func::<u64, ()>(&wasm_runner.store.lock(), function_name)?;
		let r = function.call(
			&mut wasm_runner.store.lock(),
			context.expose_provenance() as u64,
		);
		drop(unsafe { Box::from_raw(context) });
		r
	}
}

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

	pub fn script_id_by_name(&self, needle: &str) -> Option<usize> {
		for (i, script) in self.scripts.iter().enumerate() {
			match script {
				ScriptDeclaration::Static { name, .. } => {
					if &**name == needle {
						return Some(i);
					}
				}
				ScriptDeclaration::Lua { name, .. } => {
					if &**name == needle {
						return Some(i);
					}
				}
				ScriptDeclaration::Wasm { name, .. } => {
					if &**name == needle {
						return Some(i);
					}
				}
			}
		}

		None
	}

	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::Lua { path, .. } => files.start_loading_file(path, wasm_runner),
			ScriptDeclaration::Wasm { path, .. } => files.start_loading_file(path, wasm_runner),
		}
	}

	pub fn is_script_loaded(&self, script_id: usize, loaded_files: &[Box<Path>]) -> bool {
		let script = &self.scripts[script_id];
		match script {
			ScriptDeclaration::Static { .. } => true,
			ScriptDeclaration::Lua { path, .. } => loaded_files.contains(path),
			ScriptDeclaration::Wasm { path, .. } => loaded_files.contains(path),
		}
	}

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

	pub fn handle_queue(
		&mut self,
		context: &mut ScriptContext<'_>,
		wasm_runner: &WasmRunner,
		files: &mut FileManager,
	) {
		for operation in &context.operation_queue {
			match operation {
				Operation::AddCell(cell) => context.buffer.cells.push(*cell),
				Operation::PreloadScene(id) => {
					preload_scene(&context.manifest.scenes, self, files, wasm_runner, *id)
				}
				Operation::SwitchScene(id) => switch_scenes(
					self,
					&mut ScriptContext {
						manifest: context.manifest,
						buffer: context.buffer,
						input: context.input,
						time: context.time,
						storage: context.storage,
						operation_queue: Vec::new(),
					},
					files,
					wasm_runner,
					*id,
				),
				Operation::PreloadScript(id) => self.preload_script(*id, files, wasm_runner),
				Operation::UnloadScript(id) => self.unload_script(*id, files),
				Operation::ActivateScript(id) => self.activate_script(
					*id,
					&mut ScriptContext {
						manifest: context.manifest,
						buffer: context.buffer,
						input: context.input,
						time: context.time,
						storage: context.storage,
						operation_queue: Vec::new(),
					},
					wasm_runner,
					files,
				),
				Operation::DeactivateScript(id) => self.deactivate_script(
					*id,
					&mut ScriptContext {
						manifest: context.manifest,
						buffer: context.buffer,
						input: context.input,
						time: context.time,
						storage: context.storage,
						operation_queue: Vec::new(),
					},
					wasm_runner,
					files,
				),
				Operation::UpdateState {
					script_id,
					new_state,
				} => {
					self.script_states.insert(*script_id, new_state.clone());
				}
				Operation::UpdateEntity { name, new_state } => {
					self.entities.insert(name.clone(), new_state.clone());
				}
				Operation::UpdateStorage { new_save } => {
					context.storage.background_save(new_save.clone());
				}
				Operation::Exit => context.exit(),
			}
		}
		context.operation_queue.clear();
	}

	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::Lua { path, .. } => {
				files.load_file(path.clone(), wasm_runner).unwrap();
				let loaded_files = files.loaded_files();
				let module = files.get_cached_file_mut(path).unwrap().unwrap_lua_module();
				if let Err(e) = module.run_function(
					"start",
					FullScriptContext {
						script_id,
						manifest: context.manifest,
						buffer: context.buffer,
						input: context.input,
						time: context.time,
						scripts: self,
						files: loaded_files,
						storage: context.storage,
						wasm_runner,
						operation_queue: &mut context.operation_queue,
					},
				) {
					log::error!("{e}");
				}
			}
			ScriptDeclaration::Wasm { path, .. } => {
				files.load_file(path.clone(), wasm_runner).unwrap();
				let loaded_files = files.loaded_files();
				let module = files
					.get_cached_file_mut(path)
					.unwrap()
					.unwrap_wasm_module_mut();
				if let Err(e) = wasm_runner.run_function(
					script_id,
					module,
					self,
					loaded_files,
					wasm_runner,
					context,
					"start",
				) {
					log::error!("{e}");
				}
			}
		}
		self.active_scripts.push(script_id);
		self.handle_queue(context, wasm_runner, files);
	}

	pub fn deactivate_script(
		&mut self,
		script_id: usize,
		context: &mut ScriptContext<'_>,
		wasm_runner: &WasmRunner,
		files: &mut FileManager,
	) {
		for script_id in self
			.active_scripts
			.iter()
			.cloned()
			.filter(|element| *element == script_id)
		{
			let Some(script) = self.scripts.get(script_id) else {
				unreachable!("The active script does not exist");
			};
			self.script_states.remove(&script_id);
			match script {
				ScriptDeclaration::Static { end, .. } => {
					end.map(|s| {
						s(
							context,
							get_script_state(&mut self.script_states, script_id),
						);
					});
				}
				ScriptDeclaration::Lua { path, .. } => {
					files.load_file(path.clone(), wasm_runner).unwrap();
					let loaded_files = files.loaded_files();
					let module = files.get_cached_file_mut(path).unwrap().unwrap_lua_module();
					if let Err(e) = module.run_function(
						"end",
						FullScriptContext {
							script_id,
							manifest: context.manifest,
							buffer: context.buffer,
							input: context.input,
							time: context.time,
							scripts: self,
							files: loaded_files,
							storage: context.storage,
							wasm_runner,
							operation_queue: &mut context.operation_queue,
						},
					) {
						log::error!("{e}");
					}
				}
				ScriptDeclaration::Wasm { path, .. } => {
					files.load_file(path.clone(), wasm_runner).unwrap();
					let loaded_files = files.loaded_files();
					let module = files
						.get_cached_file_mut(path)
						.unwrap()
						.unwrap_wasm_module_mut();
					if let Err(e) = wasm_runner.run_function(
						script_id,
						module,
						self,
						loaded_files,
						wasm_runner,
						context,
						"end",
					) {
						log::error!("{e}");
					}
				}
			};
		}
		self.active_scripts.retain(|element| *element != script_id);
		self.handle_queue(context, wasm_runner, files);
	}

	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::Lua { path, .. } => {
					files.load_file(path.clone(), wasm_runner).unwrap();
					let loaded_files = files.loaded_files();
					let module = files.get_cached_file_mut(path).unwrap().unwrap_lua_module();
					if let Err(e) = module.run_function(
						"update",
						FullScriptContext {
							script_id: *script_id,
							manifest: context.manifest,
							buffer: context.buffer,
							input: context.input,
							time: context.time,
							scripts: self,
							files: loaded_files,
							storage: context.storage,
							wasm_runner,
							operation_queue: &mut context.operation_queue,
						},
					) {
						log::error!("{e}");
					}
				}
				ScriptDeclaration::Wasm { path, .. } => {
					files.load_file(path.clone(), wasm_runner).unwrap();
					let loaded_files = files.loaded_files();
					let module = files
						.get_cached_file_mut(path)
						.unwrap()
						.unwrap_wasm_module_mut();
					if let Err(e) = wasm_runner.run_function(
						*script_id,
						module,
						self,
						loaded_files,
						wasm_runner,
						context,
						"update",
					) {
						log::error!("Error while running scripts: {e}");
					}
				}
			}
		}
		self.handle_queue(context, wasm_runner, files);
	}
}