use core::{
	alloc::Layout,
	ops::{Deref, DerefMut},
};

use alloc::{
	alloc::{alloc, dealloc},
	borrow::Cow,
	string::String,
};
use serde::{Deserialize, Serialize};
use spin::lock_api::{MappedRwLockWriteGuard, RwLock, RwLockWriteGuard};

use crate::ScriptContext;

mod tortuise {
	#[link(wasm_import_module = "tortuise")]
	unsafe extern "C" {
		pub unsafe fn entity_size(context: u64, name: *const u8, name_len: usize) -> usize;
		pub unsafe fn load_entity(context: u64, name: *const u8, name_len: usize, offset: *mut u8);
		pub unsafe fn save_entity(
			context: u64,
			name: *const u8,
			name_len: usize,
			offset: *const u8,
			size: usize,
		);
	}
}

fn load_entity<'de, T: Deserialize<'de>>(context: u64, name: &str) -> Option<T> {
	unsafe {
		let size = tortuise::entity_size(context, name.as_ptr(), name.len());
		let Ok(layout) = Layout::array::<u8>(size) else {
			log::error!("Failed to create layout for size: {size}");
			panic!();
		};
		if layout.size() == 0 {
			return None;
		}
		let ptr = alloc(layout);
		tortuise::load_entity(context, name.as_ptr(), name.len(), ptr);
		let data = core::slice::from_raw_parts(ptr, size);
		let data = match messagepack_serde::from_slice(data) {
			Ok(data) => data,
			Err(e) => {
				log::error!("Failed to parse entity: {e}");
				return None;
			}
		};
		dealloc(ptr, layout);
		Some(data)
	}
}

fn save_entity<T: Serialize>(context: u64, name: &str, state: &T) {
	unsafe {
		let data = match messagepack_serde::to_vec(state) {
			Ok(data) => data,
			Err(e) => {
				log::error!("Failed to save entity: {e}");
				panic!();
			}
		};
		tortuise::save_entity(
			context,
			name.as_ptr(),
			name.len(),
			data.as_ptr(),
			data.len(),
		);
	}
}

pub struct Entity<T> {
	name: Cow<'static, str>,
	data: RwLock<Option<T>>,
	default: fn(&ScriptContext) -> T,
}

pub struct EntityRef<'entity, T> {
	data: MappedRwLockWriteGuard<'entity, T>,
}

pub struct EntityMut<'entity, T: Serialize> {
	context: u64,
	name: &'entity str,
	data: MappedRwLockWriteGuard<'entity, T>,
}

impl<'de, T: Serialize + Deserialize<'de>> Entity<T> {
	pub const fn new(name: &'static str, default: fn(&ScriptContext) -> T) -> Self {
		Self {
			name: Cow::Borrowed(name),
			data: RwLock::new(None),
			default,
		}
	}

	pub const fn from_string(name: String, default: fn(&ScriptContext) -> T) -> Self {
		Self {
			name: Cow::Owned(name),
			data: RwLock::new(None),
			default,
		}
	}

	pub fn get(&self, context: &ScriptContext) -> EntityRef<'_, T> {
		let data = RwLockWriteGuard::map(self.data.write(), |data| {
			data.get_or_insert_with(|| {
				load_entity(context.0, &self.name).unwrap_or_else(|| (self.default)(context))
			})
		});
		EntityRef { data }
	}

	pub fn get_mut(&self, context: &ScriptContext) -> EntityMut<'_, T> {
		let data = RwLockWriteGuard::map(self.data.write(), |data| {
			data.get_or_insert_with(|| {
				load_entity(context.0, &self.name).unwrap_or_else(|| (self.default)(context))
			})
		});
		EntityMut {
			context: context.0,
			name: &self.name,
			data,
		}
	}
}

impl<T> Deref for EntityRef<'_, T> {
	type Target = T;

	fn deref(&self) -> &Self::Target {
		&self.data
	}
}

impl<T: Serialize> Deref for EntityMut<'_, T> {
	type Target = T;

	fn deref(&self) -> &Self::Target {
		&self.data
	}
}

impl<T: Serialize> DerefMut for EntityMut<'_, T> {
	fn deref_mut(&mut self) -> &mut Self::Target {
		&mut self.data
	}
}

impl<T: Serialize> Drop for EntityMut<'_, T> {
	fn drop(&mut self) {
		save_entity(self.context, self.name, &*self.data);
	}
}