bevy_state_stack 0.2.1

use std::{collections::HashMap, fmt::Debug};

use bevy_ecs::{
	prelude::*,
	schedule::{IntoSystemDescriptor, ShouldRun, StateData},
	system::Resource,
};

enum Hook {
	Enter,
	Exit,
	Pause,
	Resume,
}

/// Function used to generate the function used for the run condition
pub fn in_state<T: PartialEq + Resource>(state: T) -> impl Fn(Res<TopState<T>>) -> ShouldRun {
	move |res| {
		if res.0 == state {
			ShouldRun::Yes
		} else {
			ShouldRun::No
		}
	}
}

mod r#trait;
pub use r#trait::AppStateStackExt;

pub struct TopState<T: PartialEq>(pub T);

impl<T: PartialEq + Debug> Debug for TopState<T> {
	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
		write!(f, "TopState({:?})", self.0)
	}
}

/// Resource to change the stack
pub enum Stack<T> {
	/// Replace the top level of the state stack.
	Set(T),
	/// Push another state on top of the stack.
	Push(T),
	/// Pop the top of the state at the top of the stack.
	Pop,
	/// Reset the stack to value,
	Reset(T),
}

struct StackStage<T: StateData> {
	enter_stages: HashMap<T, Box<dyn Stage>>,
	resume_stages: HashMap<T, Box<dyn Stage>>,
	exit_stages: HashMap<T, Box<dyn Stage>>,
	pause_stages: HashMap<T, Box<dyn Stage>>,
	stack: Vec<T>,
}

impl<T: StateData> StackStage<T> {
	fn add_system<Params>(&mut self, h: Hook, state: T, system: impl IntoSystemDescriptor<Params>) {
		let map = match h {
			Hook::Enter => &mut self.enter_stages,
			Hook::Exit => &mut self.exit_stages,
			Hook::Pause => &mut self.pause_stages,
			Hook::Resume => &mut self.resume_stages,
		};

		if let Some(stage) = map.get_mut(&state) {
			stage
				.downcast_mut::<SystemStage>()
				.expect("State is not a SystemStage! This should be impossible!")
				.add_system(system);
		} else {
			let mut stage = SystemStage::parallel();
			stage.add_system(system);
			map.insert(state.clone(), Box::new(stage));
		}
	}

	fn new(init: T) -> Self {
		Self {
			enter_stages: Default::default(),
			exit_stages: Default::default(),
			pause_stages: Default::default(),
			resume_stages: Default::default(),
			stack: vec![init],
		}
	}
}

impl<T: StateData> Stage for StackStage<T> {
	fn run(&mut self, world: &mut World) {
		if !world.contains_resource::<TopState<T>>() {
			let init = &self.stack[0];
			world.insert_resource(TopState(init.clone()));

			if let Some(stage) = self.enter_stages.get_mut(init) {
				stage.run(world);
			}
		}

		if let Some(change) = world.remove_resource::<Stack<T>>() {
			let cur = &world.get_resource::<TopState<T>>().unwrap().0;
			match change {
				Stack::Set(s) => {
					// exit the current stage
					if let Some(stage) = self.exit_stages.get_mut(cur) {
						stage.run(world);
					}

					world.insert_resource(TopState(s.clone()));

					// enter the new stage
					if let Some(stage) = self.enter_stages.get_mut(&s) {
						stage.run(world);
					}

					let last = self.stack.last_mut().unwrap();
					*last = s;
				}
				Stack::Push(s) => {
					// pause the current stage
					if let Some(stage) = self.pause_stages.get_mut(cur) {
						stage.run(world);
					}

					world.insert_resource(TopState(s.clone()));

					// enter the new stage
					if let Some(stage) = self.enter_stages.get_mut(&s) {
						stage.run(world);
					}

					self.stack.push(s)
				}
				Stack::Pop => {
					if self.stack.len() == 1 {
						eprintln!("Can't pop the only state in stack!")
					} else {
						if let Some(stage) = self.exit_stages.get_mut(cur) {
							stage.run(world);
						}

						self.stack.pop();

						let state = self.stack.last().unwrap();

						if let Some(stage) = self.resume_stages.get_mut(state) {
							stage.run(world);
						}

						world.insert_resource(TopState(state.clone()))
					}
				}
				Stack::Reset(s) => {
					for state in self.stack.iter().rev() {
						if let Some(stage) = self.exit_stages.get_mut(state) {
							stage.run(world)
						}
					}

					world.insert_resource(TopState(s.clone()));

					// enter the new stage
					if let Some(stage) = self.enter_stages.get_mut(&s) {
						stage.run(world);
					}

					self.stack = vec![s];
				}
			}
		}
	}
}