use wasmtime::Caller;

use crate::{ScriptManager, WasmScriptState};

use super::{
	system::{_alloc, _get_memory},
	LIBRARY_NAME,
};

fn allocate_zeroed(mut caller: Caller<'_, WasmScriptState>, size: u32, align: u32) -> u32 {
	let Some((ptr, data)) = _alloc(&mut caller, size, align) else {
		return 0;
	};

	data.fill(0);
	ptr
}

fn reallocate(
	mut caller: Caller<'_, WasmScriptState>,
	ptr: u32,
	size: u32,
	new_size: u32,
	align: u32,
) -> u32 {
	if new_size < size {
		return ptr;
	}

	let bump_pointer = &mut caller.data_mut().bump_pointer;
	if *bump_pointer == ptr + size {
		*bump_pointer = ptr;
		let Some((new_ptr, _)) = _alloc(&mut caller, new_size, align) else {
			return 4;
		};
		return new_ptr;
	}

	let (_, original) = _get_memory(&mut caller);
	let Some(original) = original.get(ptr as usize..(ptr + size) as usize) else {
		return 4;
	};
	let original = original.to_vec();

	let Some((new_ptr, data)) = _alloc(&mut caller, new_size, align) else {
		return 0;
	};

	data.get_mut(..size as usize)
		.unwrap()
		.clone_from_slice(&original);
	new_ptr
}

fn duplicate(mut caller: Caller<'_, WasmScriptState>, ptr: u32, size: u32, align: u32) -> u32 {
	let (_, original) = _get_memory(&mut caller);
	let Some(original) = original.get(ptr as usize..(ptr + size) as usize) else {
		return 4;
	};
	let original = original.to_vec();

	let Some((new_ptr, data)) = _alloc(&mut caller, size, align) else {
		return 0;
	};

	data.clone_from_slice(&original);
	new_ptr
}

pub fn library(manager: &mut ScriptManager) -> Option<()> {
	manager.add_library_function(LIBRARY_NAME, "calloc", allocate_zeroed)?;
	manager.add_library_function(LIBRARY_NAME, "realloc", reallocate)?;
	manager.add_library_function(LIBRARY_NAME, "memdup", duplicate)?;

	Some(())
}