beet_core 0.0.8

use crate::prelude::*;
use beet_core_macros::BundleEffect;
use bevy::ecs::error::ErrorContext;
use bevy::ecs::relationship::RelatedSpawner;
use bevy::ecs::relationship::Relationship;
use bevy::ecs::spawn::SpawnRelatedBundle;
use bevy::ecs::spawn::SpawnWith;
use bevy::ecs::system::IntoObserverSystem;

/// Type helper for [`SpawnWith`], useful for spawning any number of related entities
/// like children.
pub fn spawn_with<T: RelationshipTarget, F>(
	func: F,
) -> SpawnRelatedBundle<T::Relationship, SpawnWith<F>>
where
	F: 'static + Send + Sync + FnOnce(&mut RelatedSpawner<T::Relationship>),
{
	T::spawn(SpawnWith(func))
}


/// A type erased [`BundleEffect`] that runs a function when the entity is spawned.
#[derive(BundleEffect)]
pub struct OnSpawn(
	pub Box<dyn 'static + Send + Sync + FnOnce(&mut EntityWorldMut)>,
);

/// Allow for Bundles and Results to be returned from methods
pub trait ApplyToEntity<M>: 'static + Send + Sync {
	fn apply(self, entity: &mut EntityWorldMut);
}

pub struct BundleApplyToEntityMarker;
pub struct ResultApplyToEntityMarker;

impl<T: Bundle> ApplyToEntity<BundleApplyToEntityMarker> for T {
	fn apply(self, entity: &mut EntityWorldMut) { entity.insert(self); }
}

impl<T: Bundle> ApplyToEntity<ResultApplyToEntityMarker> for Result<T> {
	fn apply(self, entity: &mut EntityWorldMut) {
		match self {
			Ok(bundle) => {
				entity.insert(bundle);
			}
			Err(err) => entity.world_scope(|world| {
				world.default_error_handler()(
					err.into(),
					ErrorContext::Command {
						name: "ApplyToEntity".into(),
					},
				);
			}),
		}
	}
}
impl OnSpawn {
	/// Create a new [`OnSpawn`] effect.
	pub fn new(
		func: impl 'static + Send + Sync + FnOnce(&mut EntityWorldMut),
	) -> Self {
		Self(Box::new(func))
	}
	/// Insert this bundle into the entity on spawn.
	pub fn insert(bundle: impl Bundle) -> Self {
		Self::new(move |entity| {
			entity.insert(bundle);
		})
	}
	/// Insert the bundle if it is `Some`
	pub fn insert_option(bundle: Option<impl Bundle>) -> Self {
		Self::new(move |entity| {
			if let Some(bundle) = bundle {
				entity.insert(bundle);
			}
		})
	}

	/// Run the system and insert the resulting bundle into the entity on spawn.
	pub fn run_insert<
		System: 'static + Send + Sync + IntoSystem<(), Out, M1>,
		M1,
		Out: ApplyToEntity<M2>,
		M2,
	>(
		system: System,
	) -> Self {
		Self::new(move |entity| {
			entity
				.world_scope(move |world| world.run_system_once(system))
				.unwrap()
				.apply(entity);
		})
	}

	/// Insert the resource into the world when the entity is spawned.
	pub fn insert_resource(resource: impl Resource) -> Self {
		Self::new(move |entity| {
			entity.world_scope(move |world| world.insert_resource(resource));
		})
	}

	pub fn trigger<M>(event: impl IntoEntityTargetEvent<M>) -> Self {
		Self::new(move |entity| {
			entity.trigger_target(event);
		})
	}
	pub fn trigger_option<M>(
		event: Option<impl IntoEntityTargetEvent<M>>,
	) -> Self {
		Self::new(move |entity| {
			if let Some(event) = event {
				entity.trigger_target(event);
			}
		})
	}

	pub fn observe<E: Event, B: Bundle, M>(
		observer: impl 'static + Send + Sync + IntoObserverSystem<E, B, M>,
	) -> Self {
		Self::new(move |entity| {
			entity.observe_any(observer);
		})
	}


	fn effect(self, entity: &mut EntityWorldMut) { (self.0)(entity); }

	/// Create a new [`OnSpawn`] effect.
	pub fn new_async<Fut, Out>(
		func: impl 'static + Send + Sync + FnOnce(AsyncEntity) -> Fut,
	) -> Self
	where
		Fut: 'static + Send + Sync + Future<Output = Out>,
		Out: 'static + AsyncTaskOut,
	{
		Self(Box::new(move |entity| {
			let id = entity.id();
			entity.world_scope(move |world| {
				world
					.run_async(async move |world| func(world.entity(id)).await);
			});
		}))
	}
	/// Create a new [`OnSpawn`] effect.
	pub fn new_async_local<Fut, Out>(
		func: impl 'static + Send + Sync + FnOnce(AsyncEntity) -> Fut,
	) -> Self
	where
		Fut: 'static + Future<Output = Out>,
		Out: 'static + AsyncTaskOut,
	{
		Self(Box::new(move |entity| {
			let id = entity.id();
			entity.world_scope(move |world| {
				world.run_async_local(async move |world| {
					func(world.entity(id)).await
				});
			});
		}))
	}
}


/// A [`BundleEffect`] that runs a function when the entity is spawned.
#[derive(Clone, BundleEffect)]
pub struct OnSpawnTyped<F: 'static + Send + Sync + FnOnce(&mut EntityWorldMut)>(
	pub F,
);

impl<F: Send + Sync + FnOnce(&mut EntityWorldMut)> OnSpawnTyped<F> {
	/// Create a new [`OnSpawn`] effect.
	pub fn new(func: F) -> Self { Self(func) }

	fn effect(self, entity: &mut EntityWorldMut) { self.0(entity); }
}

#[derive(Component)]
pub struct OnSpawnDeferred(
	pub Box<dyn 'static + Send + Sync + FnOnce(&mut EntityWorldMut) -> Result>,
);

impl OnSpawnDeferred {
	/// Create a new [`OnSpawnDeferred`] effect.
	pub fn new(
		func: impl 'static + Send + Sync + FnOnce(&mut EntityWorldMut) -> Result,
	) -> Self {
		Self(Box::new(func))
	}

	/// Run the function for the parent of this entity
	pub fn parent<R: Relationship>(
		func: impl 'static + Send + Sync + FnOnce(&mut EntityWorldMut) -> Result,
	) -> Self {
		Self::new(move |entity| {
			let Some(parent) = entity.get::<R>() else {
				bevybail!(
					"OnSpawnDeferred::insert_parent: Entity does not have a parent"
				);
			};
			let parent = parent.get();
			entity.world_scope(move |world| func(&mut world.entity_mut(parent)))
		})
	}

	/// Insert this bundle into the entity on spawn.
	pub fn insert(bundle: impl Bundle) -> Self {
		Self::new(move |entity| {
			entity.insert(bundle);
			Ok(())
		})
	}

	/// When flushed, insert this bundle into the parent of the entity.
	pub fn insert_parent<R: Relationship>(bundle: impl Bundle) -> Self {
		Self::parent::<R>(move |entity| {
			entity.insert(bundle);
			Ok(())
		})
	}

	pub fn trigger_target<M>(ev: impl IntoEntityTargetEvent<M>) -> Self {
		Self::new(move |entity| {
			entity.trigger_target(ev);
			Ok(())
		})
	}

	/// Run all [`OnSpawnDeferred`]
	pub fn flush(
		mut commands: Commands,
		mut query: Query<(Entity, &mut Self)>,
	) {
		for (entity, mut on_spawn) in query.iter_mut() {
			commands.entity(entity).remove::<Self>();
			let func = on_spawn.take();
			commands.queue(move |world: &mut World| {
				let mut entity = world.entity_mut(entity);
				func.call(&mut entity)
			});
		}
	}

	pub fn into_command(
		self,
		entity: Entity,
	) -> impl FnOnce(&mut World) -> Result {
		move |world: &mut World| {
			let mut entity = world.entity_mut(entity);
			self.call(&mut entity)
		}
	}


	/// Call the deferred function.
	pub fn call(self, entity: &mut EntityWorldMut) -> Result {
		(self.0)(entity)
	}
	/// Convenience for getting the method from inside a system,
	/// this component should be removed when this is called
	///
	/// # Panics
	/// If the method has already been taken
	pub fn take(&mut self) -> Self {
		Self::new(std::mem::replace(
			&mut self.0,
			Box::new(|_| {
				panic!("OnSpawwnDeferred: This method has already been taken")
			}),
		))
	}
}

/// A [`Clone`] version of [`OnSpawn`]
#[derive(BundleEffect)]
pub struct OnSpawnClone(pub Box<dyn CloneEntityFunc>);

impl OnSpawnClone {
	/// Create a new [`OnSpawnCloneable`] effect.
	pub fn new(func: impl CloneEntityFunc) -> Self { Self(Box::new(func)) }
	/// Immediately inserts the bundle returned from this method
	pub fn insert<F, O>(func: F) -> Self
	where
		F: 'static + Send + Sync + Clone + FnOnce() -> O,
		O: Bundle,
	{
		Self::new(move |entity| {
			entity.insert(func.clone()());
		})
	}
	fn effect(self, entity: &mut EntityWorldMut) { (self.0)(entity); }
}


impl Clone for OnSpawnClone {
	fn clone(&self) -> Self { Self(self.0.box_clone()) }
}

pub trait CloneEntityFunc:
	'static + Send + Sync + Fn(&mut EntityWorldMut)
{
	fn box_clone(&self) -> Box<dyn CloneEntityFunc>;
}
impl<T> CloneEntityFunc for T
where
	T: 'static + Send + Sync + Clone + Fn(&mut EntityWorldMut),
{
	fn box_clone(&self) -> Box<dyn CloneEntityFunc> { Box::new(self.clone()) }
}


#[cfg(test)]
mod test {
	use crate::prelude::*;

	#[test]
	fn dfs() {
		let mut world = World::new();

		let numbers = Store::default();

		world.spawn((
			OnSpawnTyped::new(move |entity_world_mut| {
				numbers.push(1);
				entity_world_mut.insert(OnSpawnTyped::new(move |_| {
					numbers.push(2);
				}));
			}),
			OnSpawnTyped::new(move |_| {
				numbers.push(3);
			}),
		));

		numbers.get().xpect_eq(&[1, 2, 3]);
	}
	#[test]
	fn on_spawn_deferred() {
		let mut world = World::new();

		let numbers = Store::default();
		world.spawn((
			OnSpawnDeferred::new(move |entity_world_mut| {
				numbers.push(1);
				entity_world_mut.insert(OnSpawnTyped::new(move |_| {
					numbers.push(2);
				}));
				Ok(())
			}),
			children![OnSpawnDeferred::new(move |_| {
				numbers.push(3);
				Ok(())
			}),],
		));

		numbers.get().xpect_eq(&[] as &[u32]);
		world.run_system_cached(OnSpawnDeferred::flush).unwrap();

		// why is this?
		#[cfg(target_arch = "wasm32")]
		numbers.get().xpect_eq(&[1, 2, 3]);
		#[cfg(not(target_arch = "wasm32"))]
		numbers.get().xpect_eq(&[3, 1, 2]);
	}

	#[test]
	fn observe() {
		#[derive(EntityEvent)]
		struct Foo(Entity);

		let store = Store::default();
		let mut world = World::new();
		world
			.spawn(OnSpawn::observe(move |_: On<Foo>| store.set(3)))
			.trigger(Foo);

		store.get().xpect_eq(3);
	}
}