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

use spin::lock_api::{MappedMutexGuard, Mutex, MutexGuard};

use crate::ScriptContext;

pub struct ScriptState<T> {
	data: Mutex<Option<T>>,
	default: fn(&ScriptContext) -> T,
}

pub struct StateRef<'state, T> {
	data: MappedMutexGuard<'state, T>,
}

pub struct StateMut<'state, T> {
	data: MappedMutexGuard<'state, T>,
}

impl<T> ScriptState<T> {
	pub const fn new(default: fn(&ScriptContext) -> T) -> Self {
		Self {
			data: Mutex::new(None),
			default,
		}
	}

	pub fn get(&self, context: &ScriptContext) -> StateRef<'_, T> {
		puffin::profile_function!();
		let data = MutexGuard::map(self.data.lock(), |data| {
			data.get_or_insert_with(|| (self.default)(context))
		});
		StateRef { data }
	}

	pub fn get_mut(&self, context: &ScriptContext) -> StateMut<'_, T> {
		puffin::profile_function!();
		let data = MutexGuard::map(self.data.lock(), |data| {
			data.get_or_insert_with(|| (self.default)(context))
		});
		StateMut { data }
	}
}

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

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

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

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

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