snow_trail/src/state.rs

use std::any::TypeId;
use std::collections::HashMap;

use crate::entity::EntityHandle;
use crate::world::{Storage, World};

pub trait PlayerState
{
    fn tick_time(&mut self, _delta: f32) {}
    fn on_enter(&mut self, _world: &mut World, _entity: EntityHandle) {}
    fn on_exit(&mut self, _world: &mut World, _entity: EntityHandle) {}
    fn on_update(&mut self, _world: &mut World, _entity: EntityHandle, _delta: f32) {}
    fn on_physics_update(&mut self, _world: &mut World, _entity: EntityHandle, _delta: f32) {}
}

struct StateOps
{
    name: &'static str,
    remove: Box<dyn Fn(&mut World, EntityHandle)>,
    insert_default: Box<dyn Fn(&mut World, EntityHandle)>,
    on_enter: Box<dyn Fn(&mut World, EntityHandle)>,
    on_exit: Box<dyn Fn(&mut World, EntityHandle)>,
    tick: Box<dyn Fn(&mut World, EntityHandle, f32)>,
    physics_tick: Box<dyn Fn(&mut World, EntityHandle, f32)>,
}

struct StateTransition
{
    to_state_id: TypeId,
    condition: Box<dyn Fn(&World) -> bool>,
}

pub struct StateMachine
{
    state_ops: HashMap<TypeId, StateOps>,
    state_transitions: HashMap<TypeId, Vec<StateTransition>>,
    current_state_id: TypeId,
}

fn short_type_name<T: 'static>() -> &'static str
{
    let full = std::any::type_name::<T>();
    full.rsplit("::").next().unwrap_or(full)
}

impl StateMachine
{
    pub fn new<S: 'static>() -> Self
    {
        Self {
            state_ops: HashMap::new(),
            state_transitions: HashMap::new(),
            current_state_id: TypeId::of::<S>(),
        }
    }

    pub fn register_state<S: PlayerState + Default + 'static>(
        &mut self,
        storage: fn(&mut World) -> &mut Storage<S>,
    )
    {
        let ops = StateOps {
            name: short_type_name::<S>(),
            remove: Box::new(move |world, entity| {
                storage(world).components.remove(&entity);
            }),
            insert_default: Box::new(move |world, entity| {
                storage(world).insert(entity, S::default());
            }),
            on_enter: Box::new(move |world, entity| {
                if let Some(mut state) = storage(world).components.remove(&entity)
                {
                    state.on_enter(world, entity);
                    storage(world).insert(entity, state);
                }
            }),
            on_exit: Box::new(move |world, entity| {
                if let Some(mut state) = storage(world).components.remove(&entity)
                {
                    state.on_exit(world, entity);
                    storage(world).insert(entity, state);
                }
            }),
            tick: Box::new(move |world, entity, delta| {
                if let Some(mut state) = storage(world).components.remove(&entity)
                {
                    state.tick_time(delta);
                    state.on_update(world, entity, delta);
                    storage(world).insert(entity, state);
                }
            }),
            physics_tick: Box::new(move |world, entity, delta| {
                if let Some(mut state) = storage(world).components.remove(&entity)
                {
                    state.on_physics_update(world, entity, delta);
                    storage(world).insert(entity, state);
                }
            }),
        };
        self.state_ops.insert(TypeId::of::<S>(), ops);
    }

    pub fn update(&mut self, world: &mut World, entity: EntityHandle, delta: f32)
    {
        if let Some(next_state_id) = self.get_transition_state_id(world)
        {
            self.apply_transition(world, entity, next_state_id);
        }

        if let Some(ops) = self.state_ops.get(&self.current_state_id)
        {
            (ops.tick)(world, entity, delta);
        }
    }

    pub fn physics_update(&mut self, world: &mut World, entity: EntityHandle, delta: f32)
    {
        if let Some(ops) = self.state_ops.get(&self.current_state_id)
        {
            (ops.physics_tick)(world, entity, delta);
        }
    }

    fn get_transition_state_id(&self, world: &World) -> Option<TypeId>
    {
        if let Some(transitions) = self.state_transitions.get(&self.current_state_id)
        {
            for transition in transitions
            {
                if (transition.condition)(world)
                {
                    return Some(transition.to_state_id);
                }
            }
        }
        None
    }

    fn apply_transition(&mut self, world: &mut World, entity: EntityHandle, next_id: TypeId)
    {
        if let Some(ops) = self.state_ops.get(&self.current_state_id)
        {
            (ops.on_exit)(world, entity);
            (ops.remove)(world, entity);
        }

        self.current_state_id = next_id;

        if let Some(ops) = self.state_ops.get(&next_id)
        {
            (ops.insert_default)(world, entity);
            (ops.on_enter)(world, entity);
        }
    }

    pub fn add_transition<TFrom: 'static, TTo: 'static>(
        &mut self,
        condition: impl Fn(&World) -> bool + 'static,
    )
    {
        let from_id = TypeId::of::<TFrom>();
        let to_id = TypeId::of::<TTo>();
        let transitions = self.state_transitions.entry(from_id).or_default();
        transitions.push(StateTransition {
            to_state_id: to_id,
            condition: Box::new(condition),
        });
    }

    pub fn get_current_state_name(&self) -> &'static str
    {
        self.state_ops
            .get(&self.current_state_id)
            .map(|ops| ops.name)
            .unwrap_or("Unknown")
    }
}