server/src/nodes/spatial.rs

use super::Node;
use crate::core::client::Client;
use anyhow::{anyhow, ensure, Result};
use glam::{Mat4, Quat};
use mint::Vector3;
use parking_lot::Mutex;

use serde::Deserialize;
use stardust_xr::schemas::flex::{deserialize, serialize};
use stardust_xr::values::Transform;
use std::ptr;
use std::sync::{Arc, Weak};

pub struct Spatial {
	pub(super) node: Weak<Node>,
	parent: Mutex<Option<Arc<Spatial>>>,
	pub(super) transform: Mutex<Mat4>,
}

impl Spatial {
	pub fn new(node: Weak<Node>, parent: Option<Arc<Spatial>>, transform: Mat4) -> Arc<Self> {
		Arc::new(Spatial {
			node,
			parent: Mutex::new(parent),
			transform: Mutex::new(transform),
		})
	}
	pub fn add_to(
		node: &Arc<Node>,
		parent: Option<Arc<Spatial>>,
		transform: Mat4,
	) -> Result<Arc<Spatial>> {
		ensure!(
			node.spatial.get().is_none(),
			"Internal: Node already has a Spatial aspect!"
		);
		let spatial = Spatial {
			node: Arc::downgrade(node),
			parent: Mutex::new(parent),
			transform: Mutex::new(transform),
		};
		node.add_local_method("getTransform", Spatial::get_transform_flex);
		node.add_local_signal("setTransform", Spatial::set_transform_flex);
		node.add_local_signal("setSpatialParent", Spatial::set_spatial_parent_flex);
		node.add_local_signal(
			"setSpatialParentInPlace",
			Spatial::set_spatial_parent_in_place_flex,
		);
		let spatial_arc = Arc::new(spatial);
		let _ = node.spatial.set(spatial_arc.clone());
		Ok(spatial_arc)
	}

	pub fn space_to_space_matrix(from: Option<&Spatial>, to: Option<&Spatial>) -> Mat4 {
		let space_to_world_matrix = from.map_or(Mat4::IDENTITY, |from| from.global_transform());
		let world_to_space_matrix = to.map_or(Mat4::IDENTITY, |to| to.global_transform().inverse());
		world_to_space_matrix * space_to_world_matrix
	}

	pub fn local_transform(&self) -> Mat4 {
		*self.transform.lock()
	}
	pub fn global_transform(&self) -> Mat4 {
		match self.parent.lock().clone() {
			Some(value) => value.global_transform() * *self.transform.lock(),
			None => *self.transform.lock(),
		}
	}
	pub fn set_local_transform(&self, transform: Mat4) {
		*self.transform.lock() = transform;
	}
	pub fn set_local_transform_components(
		&self,
		reference_space: Option<&Spatial>,
		transform: Transform,
	) {
		let reference_to_parent_transform = reference_space
			.map(|reference_space| {
				Spatial::space_to_space_matrix(Some(reference_space), self.parent.lock().as_deref())
			})
			.unwrap_or(Mat4::IDENTITY);
		let mut local_transform_in_reference_space =
			reference_to_parent_transform.inverse() * self.local_transform();
		let (mut reference_space_scl, mut reference_space_rot, mut reference_space_pos) =
			local_transform_in_reference_space.to_scale_rotation_translation();

		if let Some(pos) = transform.position {
			reference_space_pos = pos.into()
		}
		if let Some(rot) = transform.rotation {
			reference_space_rot = rot.into()
		} else if reference_space_rot.is_nan() {
			reference_space_rot = Quat::IDENTITY;
		}
		if let Some(scl) = transform.scale {
			reference_space_scl = scl.into()
		}

		local_transform_in_reference_space = Mat4::from_scale_rotation_translation(
			reference_space_scl,
			reference_space_rot,
			reference_space_pos,
		);
		self.set_local_transform(
			reference_to_parent_transform * local_transform_in_reference_space,
		);
	}

	pub fn is_ancestor_of(&self, spatial: Arc<Spatial>) -> bool {
		let mut current_ancestor = spatial;
		loop {
			if Arc::as_ptr(&current_ancestor) == ptr::addr_of!(*self) {
				return true;
			}

			let current_ancestor_parent = current_ancestor.parent.lock().clone();
			if let Some(parent) = current_ancestor_parent {
				current_ancestor = parent;
			} else {
				return false;
			}
		}
	}

	pub fn set_spatial_parent(&self, parent: Option<&Arc<Spatial>>) -> Result<()> {
		let is_ancestor = parent
			.map(|parent| self.is_ancestor_of(parent.clone()))
			.unwrap_or(false);
		if is_ancestor {
			return Err(anyhow!("Setting spatial parent would cause a loop"));
		}

		*self.parent.lock() = parent.cloned();

		Ok(())
	}

	pub fn set_spatial_parent_in_place(&self, parent: Option<&Arc<Spatial>>) -> Result<()> {
		let is_ancestor = parent
			.map(|parent| self.is_ancestor_of(parent.clone()))
			.unwrap_or(false);
		if is_ancestor {
			return Err(anyhow!("Setting spatial parent would cause a loop"));
		}

		self.set_local_transform(Spatial::space_to_space_matrix(
			Some(self),
			parent.cloned().as_deref(),
		));
		*self.parent.lock() = parent.cloned();

		Ok(())
	}

	pub fn get_transform_flex(
		node: &Node,
		calling_client: Arc<Client>,
		data: &[u8],
	) -> Result<Vec<u8>> {
		let this_spatial = node
			.spatial
			.get()
			.ok_or_else(|| anyhow!("Node doesn't have a spatial?"))?;
		let relative_spatial = calling_client
			.scenegraph
			.get_node(deserialize(data)?)
			.ok_or_else(|| anyhow!("Space not found"))?
			.spatial
			.get()
			.ok_or_else(|| anyhow!("Reference space node is not a spatial"))?
			.clone();

		let (scale, rotation, position) = Spatial::space_to_space_matrix(
			Some(this_spatial.as_ref()),
			Some(relative_spatial.as_ref()),
		)
		.to_scale_rotation_translation();

		serialize((
			mint::Vector3::from(position),
			mint::Quaternion::from(rotation),
			mint::Vector3::from(scale),
		))
		.map_err(|e| e.into())
	}
	pub fn set_transform_flex(node: &Node, calling_client: Arc<Client>, data: &[u8]) -> Result<()> {
		#[derive(Deserialize)]
		struct TransformArgs<'a> {
			reference_space_path: Option<&'a str>,
			transform: Transform,
		}
		let transform_args: TransformArgs = deserialize(data)?;
		let reference_space_transform = transform_args
			.reference_space_path
			.map(|path| -> Result<Arc<Spatial>> {
				Ok(calling_client
					.scenegraph
					.get_node(path)
					.ok_or_else(|| anyhow!("Other spatial node not found"))?
					.spatial
					.get()
					.ok_or_else(|| anyhow!("Node is not a Spatial!"))?
					.clone())
			})
			.transpose()?;

		node.spatial.get().unwrap().set_local_transform_components(
			reference_space_transform.as_deref(),
			transform_args.transform,
		);
		Ok(())
	}
	pub fn set_spatial_parent_flex(
		node: &Node,
		calling_client: Arc<Client>,
		data: &[u8],
	) -> Result<()> {
		let parent = calling_client
			.scenegraph
			.get_node(deserialize(data)?)
			.ok_or_else(|| anyhow!("Parent node not found"))?
			.spatial
			.get()
			.ok_or_else(|| anyhow!("Parent node is not a Spatial!"))?
			.clone();
		node.spatial
			.get()
			.unwrap()
			.set_spatial_parent(Some(&parent))?;
		Ok(())
	}
	pub fn set_spatial_parent_in_place_flex(
		node: &Node,
		calling_client: Arc<Client>,
		data: &[u8],
	) -> Result<()> {
		let parent_path = flexbuffers::Reader::get_root(data)?.get_str()?;
		let parent = calling_client
			.scenegraph
			.get_node(parent_path)
			.ok_or_else(|| anyhow!("Parent node not found"))?
			.spatial
			.get()
			.ok_or_else(|| anyhow!("Parent node is not a Spatial!"))?
			.clone();
		node.spatial
			.get()
			.unwrap()
			.set_spatial_parent_in_place(Some(&parent))?;
		Ok(())
	}
}

pub fn parse_transform(
	transform: Transform,
	translation: bool,
	rotation: bool,
	scale: bool,
) -> Result<Mat4> {
	let translation = translation
		.then_some(transform.position)
		.flatten()
		.unwrap_or_else(|| Vector3::from([0.0; 3]));
	let rotation = rotation
		.then_some(transform.rotation)
		.flatten()
		.unwrap_or_else(|| Quat::IDENTITY.into());
	let scale = scale
		.then_some(transform.scale)
		.flatten()
		.unwrap_or_else(|| Vector3::from([1.0; 3]));

	Ok(Mat4::from_scale_rotation_translation(
		scale.into(),
		rotation.into(),
		translation.into(),
	))
}

pub fn get_spatial_parent_flex(
	calling_client: &Arc<Client>,
	node_path: &str,
) -> Result<Arc<Spatial>> {
	Ok(calling_client
		.scenegraph
		.get_node(node_path)
		.ok_or_else(|| anyhow!("Spatial parent node not found"))?
		.spatial
		.get()
		.ok_or_else(|| anyhow!("Spatial parent node is not a spatial"))?
		.clone())
}

pub fn create_interface(client: &Arc<Client>) {
	let node = Node::create(client, "", "spatial", false);
	node.add_local_signal("createSpatial", create_spatial_flex);
	node.add_to_scenegraph();
}

pub fn create_spatial_flex(_node: &Node, calling_client: Arc<Client>, data: &[u8]) -> Result<()> {
	#[derive(Deserialize)]
	struct CreateSpatialInfo<'a> {
		name: &'a str,
		parent_path: &'a str,
		transform: Transform,
	}
	let info: CreateSpatialInfo = deserialize(data)?;
	let node = Node::create(&calling_client, "/spatial/spatial", info.name, true);
	let parent = get_spatial_parent_flex(&calling_client, info.parent_path)?;
	let transform = parse_transform(info.transform, true, true, true)?;
	let node = node.add_to_scenegraph();
	Spatial::add_to(&node, Some(parent), transform)?;
	Ok(())
}