path: root/scripts/src/libs/system.rs

use std::borrow::Cow;
use std::mem::align_of;

use chrono::Offset;
use wasmtime::{Caller, Memory};

use crate::{ScriptManager, WasmScriptState};

use super::LIBRARY_NAME;

pub fn _get_memory<'a>(caller: &'a mut Caller<WasmScriptState>) -> (Memory, &'a mut [u8]) {
	let memory = caller.get_export("memory").unwrap().into_memory().unwrap();
	let mem_ptr = memory.data_mut(caller);
	(memory, mem_ptr)
}

fn _get_string<'a>(
	caller: &'a mut Caller<WasmScriptState>,
	offset: u32,
	length: u32,
) -> Option<Cow<'a, str>> {
	let (_, mem_ptr) = _get_memory(caller);
	if (offset as usize + length as usize) > mem_ptr.len() {
		return None;
	}

	Some(String::from_utf8_lossy(
		&mem_ptr[offset as usize..length as usize],
	))
}

fn _log(caller: &mut Caller<WasmScriptState>, level: log::Level, message: u32, length: u32) -> u32 {
	let Some(message) = _get_string(caller, message, length) else {
		return 4;
	};
	log::log!(level, "{}", message);
	0
}

pub fn _alloc<'a>(
	mut caller: &'a mut Caller<WasmScriptState>,
	size: u32,
	align: u32,
) -> Option<(u32, &'a mut [u8])> {
	// find bumper location in memory
	let bump_offset = caller.data().bump_pointer;
	let (memory, memory_ptr) = _get_memory(caller);
	let bump_ptr = &memory_ptr[bump_offset as usize..];

	// calculate pointer to the new allocation
	let offset = bump_ptr.as_ptr().align_offset(align as usize);
	let new_offset = bump_offset + offset as u32;

	// update bump pointer
	let bump_offset = &mut caller.data_mut().bump_pointer;
	*bump_offset = new_offset + size;

	// grow memory if necessary
	let bump_offset = *bump_offset;
	if memory.data_size(&caller) < bump_offset as usize {
		let delta = (size + offset as u32) / (64 * 1024) + 1;
		if memory.grow(&mut caller, delta as u64).is_err() {
			return None;
		}
	}

	let new_ptr = &mut _get_memory(caller).1[new_offset as usize..size as usize];
	Some((new_offset, new_ptr))
}

fn allocate(mut caller: Caller<WasmScriptState>, size: u32, align: u32) -> u32 {
	match _alloc(&mut caller, size, align) {
		Some((ptr, _)) => ptr,
		None => 0,
	}
}

fn free_sized(mut caller: Caller<'_, WasmScriptState>, ptr: u32, size: u32) -> u32 {
	let bump_offset = &mut caller.data_mut().bump_pointer;
	if ptr + size == *bump_offset {
		*bump_offset = ptr;
	}

	0
}

fn free_aligned_sized(
	caller: Caller<'_, WasmScriptState>,
	ptr: u32,
	_alignment: u32,
	size: u32,
) -> u32 {
	free_sized(caller, ptr, size)
}

fn env_var(mut caller: Caller<'_, WasmScriptState>, name: u32, length: u32) -> u32 {
	let Some(name) = _get_string(&mut caller, name, length) else {
		return 4;
	};
	let Ok(value) = std::env::var(name.to_string()) else {
		return 1;
	};

	let Some((offset, ptr)) = _alloc(&mut caller, value.len() as u32, align_of::<u8>() as u32)
	else {
		return 0;
	};
	ptr.clone_from_slice(value.as_bytes());

	offset
}

fn random_number() -> u64 {
	rand::random()
}

fn current_time_utc() -> (i64, u32) {
	let duration = chrono::Utc::now() - chrono::DateTime::<chrono::Utc>::MIN_UTC;
	(
		duration.num_seconds(),
		duration.to_std().unwrap().subsec_nanos(),
	)
}

fn local_offset() -> i32 {
	chrono::Local::now().offset().fix().local_minus_utc()
}

fn log(mut caller: Caller<'_, WasmScriptState>, message: u32, length: u32) {
	_log(&mut caller, log::Level::Info, message, length);
}

fn log_warn(mut caller: Caller<'_, WasmScriptState>, message: u32, length: u32) {
	_log(&mut caller, log::Level::Warn, message, length);
}

fn log_error(mut caller: Caller<'_, WasmScriptState>, message: u32, length: u32) {
	_log(&mut caller, log::Level::Error, message, length);
}

pub fn library(manager: &mut ScriptManager) -> Option<()> {
	manager.add_library_function(LIBRARY_NAME, "aligned_alloc", allocate)?;
	manager.add_library_function(LIBRARY_NAME, "free_sized", free_sized)?;
	manager.add_library_function(LIBRARY_NAME, "free_aligned_sized", free_aligned_sized)?;
	manager.add_library_function(LIBRARY_NAME, "env_var", env_var)?;
	manager.add_library_function(LIBRARY_NAME, "random_number", random_number)?;
	manager.add_library_function(LIBRARY_NAME, "current_time_utc", current_time_utc)?;
	manager.add_library_function(LIBRARY_NAME, "local_offset", local_offset)?;
	manager.add_library_function(LIBRARY_NAME, "log", log)?;
	manager.add_library_function(LIBRARY_NAME, "log_warn", log_warn)?;
	manager.add_library_function(LIBRARY_NAME, "log_error", log_error)?;

	Some(())
}