low-luajit 0.2.0

use std::ffi::CStr;
use low_luajit::{AmountReturned, LuaState, LuaType, reg};

const NUM_VECTOR_TYPE_NAME: &CStr =  c"numerical vector";

fn check_vec(lua: &LuaState) -> &Vec<f64> {
	// SAFETY: "numerical vector" userdata should always be valid and of this type, and there should be no other mutable reference to it.
	unsafe { lua.check_full_userdata::<Vec<f64>>(1, NUM_VECTOR_TYPE_NAME) }
}

fn check_vec_mut(lua: &mut LuaState) -> &mut Vec<f64> {
	// SAFETY: "numerical vector" userdata should always be valid and of this type, and there should be no other reference to it.
	unsafe { lua.check_full_userdata::<Vec<f64>>(1, NUM_VECTOR_TYPE_NAME) }
}

extern "C-unwind" fn to_string(lua: LuaState) -> AmountReturned {
	let vec = check_vec(&lua);

	let string = format!("{vec:?}");
	lua.return_string(string)
}

extern "C-unwind" fn push(mut lua: LuaState) -> AmountReturned {
	let n = lua.check_number(2);
	let vec = check_vec_mut(&mut lua);

	vec.push(n);

	lua.return_none()
}

extern "C-unwind" fn pop(mut lua: LuaState) -> AmountReturned {
	let vec = check_vec_mut(&mut lua);

	match vec.pop() {
		Some(n) => lua.return_number(n),
		None => lua.return_nil(),
	}
}

extern "C-unwind" fn index(mut lua: LuaState) -> AmountReturned {
	let vec = check_vec(&lua);

	match lua.type_of(2) {
		// The index is a string, so it's probably trying to access the vector's methods.
		// Let's quickly generate the methods' table and get the requested method.
		// (It would be better to cache it somewhere, but for this example this is fine enough.)
		LuaType::String => {
			// We can create the table as a "module" for convenience. This way we can quickly add more methods if needed.
			lua.new_module(reg![push, pop]);

			// The key needs to be the top element, so we move the table to the 2nd element instead.
			lua.insert(2);

			// Index the method table.
			lua.get_table(2);
		},

		// If the index is just a number, attempt to index the vector itself
		LuaType::Number if let Some(n) = vec.get(lua.to_integer(2).cast_unsigned()) => lua.push_number(*n),

		// Invalid index, return nil.
		_ => lua.push_nil(),
	}

	lua.return_any(1)
}

#[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
extern "C-unwind" fn new_index(mut lua: LuaState) -> AmountReturned {
	let i = lua.check_number(2);
	let new = lua.check_number(3);
	let vec = check_vec_mut(&mut lua);

	let Some(value) = vec.get_mut(i as usize) else { lua.error(c"invalid index"); };
	*value = new;

	lua.return_none()
}

extern "C-unwind" fn constructor(mut lua: LuaState) -> AmountReturned {
	let mut vec = Vec::new();

	// Initialize the vector if the called passed some numbers.
	for i in 1..=lua.get_top() {
		vec.push(lua.to_number(i));
	}

	// Push the vector to the stack as userdata.
	lua.new_full_userdata::<Vec<f64>>(vec);

	// Set the vector's metatable, initializing the metatable itself if's been first created.
	if lua.new_userdata_metatable(NUM_VECTOR_TYPE_NAME) {
		// Vectors allocate their own memory, so we need to have them be dropped when garbage collected.
		// SAFETY: The metatable is being used by a valid full userdata of the same type.
		unsafe { lua.set_drop_gc_metamethod::<Vec<f64>>(); }

		// __tostring metamethod.
		lua.push_cfunction(to_string);
		lua.set_field(-2, c"__tostring");

		// __index metamethod.
		lua.push_cfunction(index);
		lua.set_field(-2, c"__index");

		// __newindex metamethod.
		lua.push_cfunction(new_index);
		lua.set_field(-2, c"__newindex");
	}
	lua.set_metatable(-2);

	lua.return_any(1)
}

#[unsafe(no_mangle)]
extern "C-unwind" fn luaopen_num_vector(lua: LuaState) -> AmountReturned {
	lua.return_cfunction(constructor)
}