proxisim/src/effect.rs

use std::{any::TypeId, collections::HashMap, sync::Mutex};

use bevy_ecs::{prelude::*, system::SystemParam};

use crate::{
    hierarchy::{HierarchyBuilder, Parent},
    player::{Current, Defender, Player},
    Stages,
};

const SECONDS_PER_ACTION: f32 = 1.1;

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Component)]
pub struct EffectId(usize);

pub struct EffectInfo {
    pub apply: Box<dyn System<In = Vec<Entity>, Out = ()>>,
    pub teardown: Box<dyn System<In = Vec<Entity>, Out = ()>>,
}

// TODO: remove need for unsafe code by splitting the system info from this struct
#[derive(Resource, Default)]
struct EffectRegistry {
    system_info: HashMap<EffectId, EffectInfo>,
    type_map: HashMap<TypeId, EffectId>,
    id_counter: usize,
}

#[derive(Resource, Default)]
struct ScheduledEffects {
    create: Mutex<HashMap<EffectId, Vec<Entity>>>,
    teardown: Mutex<HashMap<EffectId, Vec<Entity>>>,
}

#[derive(Component)]
pub struct TurnLimitedEffect {
    pub player: Entity,
    pub turns: u16,
}

impl TurnLimitedEffect {
    pub fn new(player: Entity, turns: u16) -> Self {
        Self { player, turns }
    }
}

#[derive(Component)]
pub struct TimeLimitedEffect(pub f32);

#[derive(Component)]
struct Timestamp(f32);

#[derive(Component)]
struct Permanent;

#[derive(Resource, Default)]
pub struct Clock(pub f32);

/// Marker for effects that last until the beginning of the next round
#[derive(Component)]
pub struct FullRoundEffect;

#[derive(SystemParam)]
pub struct Effects<'w, 's> {
    registry: Res<'w, EffectRegistry>,
    scheduled: Res<'w, ScheduledEffects>,
    effect_q: Query<'w, 's, &'static EffectId>,
    commands: Commands<'w, 's>,
}

// TODO: Schedule effects using commands in order to avoid the need for having a separate
// `Schedule` dedicated to them.
impl<'w, 's> Effects<'w, 's> {
    pub fn spawn_and_insert<T: Component + 'static>(
        &mut self,
        effect: T,
        to: Entity,
        addins: impl Bundle,
    ) -> Entity {
        let id = self.registry.type_map.get(&TypeId::of::<T>()).unwrap();

        let spawned_effect = self
            .commands
            .spawn(effect)
            .insert(*id)
            .insert(addins)
            .set_parent(to)
            .id();

        self.scheduled
            .create
            .lock()
            .unwrap()
            .entry(*id)
            .or_default()
            .push(spawned_effect);

        spawned_effect
    }

    pub fn spawn<T: Component + 'static>(&mut self, effect: T, to: Entity) -> Entity {
        self.spawn_and_insert(effect, to, ())
    }

    pub fn remove(&mut self, entity: Entity) {
        let id = self.effect_q.get(entity).unwrap();
        self.scheduled
            .teardown
            .lock()
            .unwrap()
            .entry(*id)
            .or_default()
            .push(entity);
    }
}

pub struct EffectBuilder<'w> {
    world: &'w mut World,
    apply: Option<Box<dyn System<In = Vec<Entity>, Out = ()>>>,
    teardown: Option<Box<dyn System<In = Vec<Entity>, Out = ()>>>,
    id: EffectId,
}

impl<'r> EffectBuilder<'r> {
    #[must_use]
    pub fn apply<A, S>(&mut self, system: S) -> &mut Self
    where
        S: IntoSystem<Vec<Entity>, (), A> + 'static,
    {
        let mut system = IntoSystem::into_system(system);
        system.initialize(self.world);
        self.apply = Some(Box::new(system));
        self
    }

    #[must_use]
    pub fn teardown<T, S>(&mut self, system: S) -> &mut Self
    where
        S: IntoSystem<Vec<Entity>, (), T> + 'static,
    {
        let mut system = IntoSystem::into_system(system);
        system.initialize(self.world);
        self.teardown = Some(Box::new(system));
        self
    }

    pub fn build(&mut self) {
        let mut registry = self.world.resource_mut::<EffectRegistry>();
        registry.system_info.insert(
            self.id,
            EffectInfo {
                apply: self.apply.take().unwrap(),
                teardown: self.teardown.take().unwrap(),
            },
        );
    }
}

pub(crate) fn run_effects(world: &mut World) {
    loop {
        let (next_effect, entities) = {
            let scheduled = world.resource_mut::<ScheduledEffects>();
            let mut create = scheduled.create.lock().unwrap();
            let Some(next_effect) = create.keys().next().copied() else {
                break;
            };
            let entities = create.remove(&next_effect).unwrap();

            (next_effect, entities)
        };

        // safety: since the registry can't be mutably accessed by effect handlers this should be
        // fine
        // ...probably
        unsafe {
            let unsafe_world = world.as_unsafe_world_cell();
            let mut registry = unsafe_world.get_resource_mut::<EffectRegistry>().unwrap();
            let info = registry.system_info.get_mut(&next_effect).unwrap();
            info.apply.run(entities, unsafe_world.world_mut());
            info.apply.apply_deferred(unsafe_world.world_mut());
        };
    }

    loop {
        let (next_effect, entities) = {
            let scheduled = world.resource_mut::<ScheduledEffects>();
            let mut teardown = scheduled.teardown.lock().unwrap();
            let Some(next_effect) = teardown.keys().next().copied() else {
                break;
            };
            let entities = teardown.remove(&next_effect).unwrap();

            (next_effect, entities)
        };

        unsafe {
            let unsafe_world = world.as_unsafe_world_cell();
            let mut registry = unsafe_world.get_resource_mut::<EffectRegistry>().unwrap();
            let info = registry.system_info.get_mut(&next_effect).unwrap();
            info.teardown
                .run(entities.clone(), unsafe_world.world_mut());
            info.teardown.apply_deferred(unsafe_world.world_mut());
        };

        for entity in entities {
            world.despawn(entity);
        }
    }
}

#[must_use]
pub(crate) fn register_effect<Effect: 'static>(stages: &mut Stages) -> EffectBuilder {
    let mut registry = stages.world.resource_mut::<EffectRegistry>();
    let id = EffectId(registry.id_counter);
    registry.id_counter += 1;
    registry.type_map.insert(TypeId::of::<Effect>(), id);

    EffectBuilder {
        world: &mut stages.world,
        apply: None,
        teardown: None,
        id,
    }
}

fn update_round_limited_effects(
    mut effect_q: Query<(Entity, &mut TurnLimitedEffect)>,
    current_q: Query<Has<Current>>,
    mut effects: Effects,
) {
    for (entity, mut effect) in effect_q.iter_mut() {
        if !current_q.get(effect.player).unwrap() {
            continue;
        }

        if effect.turns == 0 {
            effects.remove(entity);
            continue;
        }
        effect.turns -= 1;
    }
}

fn advance_clock(
    mut clock: ResMut<Clock>,
    current_q: Query<Has<Defender>, (With<Current>, With<Player>)>,
) {
    let is_defender = current_q.single();
    if !is_defender {
        clock.0 += SECONDS_PER_ACTION;
    }
}

fn update_time_limited_effects(
    new_effects_q: Query<Entity, Added<TimeLimitedEffect>>,
    active_effects_q: Query<(Entity, &Timestamp, &TimeLimitedEffect)>,
    mut commands: Commands,
    mut effects: Effects,
    clock: Res<Clock>,
) {
    for entity in new_effects_q.iter() {
        commands.entity(entity).insert(Timestamp(clock.0));
    }

    for (entity, timestamp, effect) in active_effects_q.iter() {
        if (timestamp.0 + effect.0) < clock.0 {
            effects.remove(entity);
        }
    }
}

fn mark_permanent_effects(effect_q: Query<Entity, With<EffectId>>, mut commands: Commands) {
    for effect in effect_q.iter() {
        commands.entity(effect).insert(Permanent);
    }
}

fn remove_transient_effects(
    effect_q: Query<(Entity, &Parent), (With<EffectId>, Without<Permanent>)>,
    mut commands: Commands,
) {
    for (effect, target) in effect_q.iter() {
        commands.entity(effect).despawn();
        commands.entity(target.get()).remove_child(effect);
    }
}

pub(crate) fn configure(stages: &mut Stages) {
    stages.world.init_resource::<EffectRegistry>();
    stages.world.init_resource::<ScheduledEffects>();
    stages.world.init_resource::<Clock>();

    stages.snapshot.add_systems(mark_permanent_effects);
    stages
        .pre_turn
        .add_systems((advance_clock, update_round_limited_effects));
    stages.post_turn.add_systems(update_time_limited_effects);
    stages.restore.add_systems(remove_transient_effects);
}