HEDZReloaded/crates/shared/src/ai/mod.rs

use crate::{
    GameState,
    abilities::{HeadAbility, TriggerData, TriggerThrow},
    aim::AimTarget,
    heads::ActiveHeads,
    heads_database::HeadsDatabase,
    player::Player,
};
use bevy::prelude::*;

#[derive(Component, Reflect)]
#[reflect(Component)]
pub struct Ai;

#[derive(Component, Reflect, Clone)]
#[reflect(Component)]
struct WaitForAnyPlayer;

#[derive(Component, Reflect, Clone)]
#[reflect(Component)]
struct Engage(Entity);

#[derive(Component, Reflect, Clone)]
#[reflect(Component)]
struct Reload;

pub fn plugin(app: &mut App) {
    app.add_systems(
        Update,
        (
            engage_and_throw,
            wait_for_player,
            out_of_range,
            detect_reload,
            detect_reload_done,
            rotate,
        )
            .run_if(in_state(GameState::Playing)),
    );
    app.add_systems(Update, on_ai_added);
}

fn on_ai_added(mut commands: Commands, query: Query<Entity, Added<Ai>>) {
    for entity in query.iter() {
        commands.entity(entity).insert(WaitForAnyPlayer);
    }
}

fn wait_for_player(
    mut commands: Commands,
    agents: Query<Entity, With<WaitForAnyPlayer>>,
    transform: Query<&Transform>,
    players: Query<Entity, With<Player>>,
) {
    for agent in agents.iter() {
        if let Some(player) = in_range(50., agent, &players, &transform) {
            info!("[{agent}] Engage: {player}");
            if let Ok(mut agent) = commands.get_entity(agent) {
                agent.remove::<WaitForAnyPlayer>().insert(Engage(player));
            }
        }
    }
}

fn out_of_range(
    mut commands: Commands,
    agents: Query<Entity, With<Engage>>,
    transform: Query<&Transform>,
    players: Query<Entity, With<Player>>,
) {
    for agent in agents.iter() {
        if in_range(100., agent, &players, &transform).is_none() {
            info!("[{agent}] Player out of range");
            commands
                .entity(agent)
                .remove::<Engage>()
                .insert(WaitForAnyPlayer);
        }
    }
}

fn detect_reload(mut commands: Commands, agents: Query<(Entity, &ActiveHeads), With<Engage>>) {
    for (e, head) in agents.iter() {
        if head.reloading() {
            info!("[{e}] Reload started");
            commands.entity(e).remove::<Engage>().insert(Reload);
        }
    }
}

fn detect_reload_done(mut commands: Commands, agents: Query<(Entity, &ActiveHeads), With<Reload>>) {
    for (e, head) in agents.iter() {
        if !head.reloading() {
            info!("[{e}] Reload done");
            commands
                .entity(e)
                .remove::<Reload>()
                .insert(WaitForAnyPlayer);
        }
    }
}

fn in_range(
    range: f32,
    entity: Entity,
    players: &Query<'_, '_, Entity, With<Player>>,
    transform: &Query<'_, '_, &Transform>,
) -> Option<Entity> {
    let Ok(pos) = transform.get(entity).map(|t| t.translation) else {
        return None;
    };

    players
        .iter()
        .filter_map(|p| transform.get(p).ok().map(|t| (p, *t)))
        .find(|(_, t)| t.translation.distance(pos) < range)
        .map(|(e, _)| e)
}

fn rotate(agent: Query<(Entity, &Engage)>, mut transform: Query<&mut Transform>) {
    for (agent, Engage(target)) in agent.iter() {
        let Ok(target_pos) = transform.get(*target).map(|t| t.translation) else {
            continue;
        };
        let Ok(mut agent_transform) = transform.get_mut(agent) else {
            continue;
        };

        // Get the direction vector from the current position to the target
        let direction = (agent_transform.translation - target_pos).normalize();

        // Project the direction onto the XZ plane by zeroing out the Y component
        let xz_direction = Vec3::new(direction.x, 0.0, direction.z).normalize();

        agent_transform.rotation = Quat::from_rotation_arc(Vec3::Z, xz_direction);
    }
}

fn engage_and_throw(
    mut commands: Commands,
    mut query: Query<(&mut ActiveHeads, &AimTarget, &Transform), With<Engage>>,
    time: Res<Time>,
    heads_db: Res<HeadsDatabase>,
) {
    for (mut npc, target, t) in query.iter_mut() {
        if target.0.is_none() {
            continue;
        }

        let Some(npc_head) = npc.current() else {
            continue;
        };

        let ability = heads_db.head_stats(npc_head.head).ability;

        //TODO: support other abilities
        if ability != HeadAbility::Thrown {
            continue;
        }

        let can_shoot_again = npc_head.last_use + 1. < time.elapsed_secs();

        if can_shoot_again && npc_head.has_ammo() {
            npc.use_ammo(time.elapsed_secs());

            commands.trigger(TriggerThrow(TriggerData::new(
                target.0,
                t.forward(),
                t.translation,
                crate::physics_layers::GameLayer::Player,
                npc_head.head,
                // TODO: we probably need to make sure the ai's projectile does not get deduped, zero should not be used by anyone though
                0,
            )));
        }
    }
}