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

use alloc::alloc::{alloc, dealloc};
use serde::{Deserialize, Serialize};

use crate::ScriptContext;

mod tortuise {
	#[link(wasm_import_module = "tortuise")]
	unsafe extern "C" {
		pub safe fn storage_size(context: u64) -> usize;
		pub unsafe fn load_storage(context: u64, ptr: *mut u8);
		pub unsafe fn save_storage(context: u64, ptr: *const u8, size: usize);
	}
}

fn load_storage<'de, T: Deserialize<'de>>(context: u64) -> Option<T> {
	unsafe {
		let size = tortuise::storage_size(context);
		let layout = Layout::array::<u8>(size).unwrap();
		let ptr = alloc(layout);
		tortuise::load_storage(context, ptr);
		let data = core::slice::from_raw_parts(ptr, size);
		let data = messagepack_serde::from_slice(data).ok()?;
		dealloc(ptr, layout);
		Some(data)
	}
}

fn save_storage<T: Serialize>(context: u64, state: &T) {
	unsafe {
		let data = messagepack_serde::to_vec(state).unwrap();
		tortuise::save_storage(context, data.as_ptr(), data.len());
	}
}

pub struct Storage<T> {
	initial_state: fn() -> T,
}

pub struct StorageMut<T: Serialize> {
	context: u64,
	data: T,
}

impl<'de, T: Serialize + Deserialize<'de>> Storage<T> {
	pub const fn new(initial_state: fn() -> T) -> Self {
		Self { initial_state }
	}

	pub fn load(&self, context: &ScriptContext) -> T {
		load_storage(context.0).unwrap_or_else(self.initial_state)
	}

	pub fn load_mut(&self, context: &ScriptContext) -> StorageMut<T> {
		StorageMut {
			context: context.0,
			data: load_storage(context.0).unwrap_or_else(self.initial_state),
		}
	}
}

impl<T: Serialize> Deref for StorageMut<T> {
	type Target = T;

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

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

impl<T: Serialize> Drop for StorageMut<T> {
	fn drop(&mut self) {
		save_storage(self.context, &self.data);
	}
}