HEDZReloaded/crates/hedz_reloaded/src/heads/mod.rs

use crate::{
    GameState,
    animation::AnimationFlags,
    backpack::{Backpack, BackpackSwapEvent},
    control::{SelectLeftPressed, SelectRightPressed},
    global_observer,
    heads_database::HeadsDatabase,
    hitpoints::Hitpoints,
    player::Player,
    protocol::{ClientToController, PlaySound, is_server},
};
use bevy::prelude::*;
use bevy_replicon::prelude::FromClient;
use serde::{Deserialize, Serialize};

pub mod heads_ui;

pub static HEAD_COUNT: usize = 18;
pub static HEAD_SLOTS: usize = 5;

#[derive(Resource, Default)]
pub struct HeadsImages {
    pub heads: Vec<Handle<Image>>,
}

#[derive(Clone, Copy, Debug, PartialEq, Reflect, Serialize, Deserialize)]
pub struct HeadState {
    pub head: usize,
    pub health: u32,
    pub health_max: u32,
    pub ammo: u32,
    pub ammo_max: u32,
    pub reload_duration: f32,
    pub last_use: f32,
}

impl HeadState {
    pub fn new(head: usize, heads_db: &HeadsDatabase) -> Self {
        let ammo = heads_db.head_stats(head).ammo;
        Self {
            head,
            health: 100,
            health_max: 100,
            ammo,
            ammo_max: ammo,
            reload_duration: 5.,
            last_use: 0.,
        }
    }

    pub fn has_ammo(&self) -> bool {
        self.ammo > 0
    }
}

#[derive(Component, Default, Reflect, Debug, Serialize, Deserialize, PartialEq)]
#[reflect(Component)]
pub struct ActiveHeads {
    heads: [Option<HeadState>; 5],
    current_slot: usize,
    selected_slot: usize,
}

impl ActiveHeads {
    pub fn new(heads: [Option<HeadState>; 5]) -> Self {
        Self {
            heads,
            current_slot: 0,
            selected_slot: 0,
        }
    }

    pub fn current(&self) -> Option<HeadState> {
        self.heads[self.current_slot]
    }

    pub fn use_ammo(&mut self, time: f32) {
        let Some(head) = &mut self.heads[self.current_slot] else {
            error!("cannot use ammo of empty head");
            return;
        };

        head.last_use = time;
        head.ammo = head.ammo.saturating_sub(1);
    }

    pub fn medic_heal(&mut self, heal_amount: u32, time: f32) -> Option<u32> {
        let mut healed = false;
        for (index, head) in self.heads.iter_mut().enumerate() {
            if index == self.current_slot {
                continue;
            }

            if let Some(head) = head
                && head.health < head.health_max
            {
                head.health = head
                    .health
                    .saturating_add(heal_amount)
                    .clamp(0, head.health_max);
                healed = true;
            }
        }

        if healed {
            let Some(head) = &mut self.heads[self.current_slot] else {
                error!("cannot heal with empty head");
                return None;
            };

            head.last_use = time;
            head.health = head.health.saturating_sub(1);

            Some(head.health)
        } else {
            None
        }
    }

    pub fn head(&self, slot: usize) -> Option<HeadState> {
        self.heads[slot]
    }

    pub fn reloading(&self) -> bool {
        for head in self.heads {
            let Some(head) = head else {
                continue;
            };

            if head.ammo == 0 {
                return true;
            }
        }

        false
    }

    pub fn hp(&self) -> Hitpoints {
        if let Some(head) = &self.heads[self.current_slot] {
            Hitpoints::new(head.health_max).with_health(head.health)
        } else {
            Hitpoints::new(0)
        }
    }

    pub fn set_hitpoint(&mut self, hp: &Hitpoints) {
        let Some(head) = &mut self.heads[self.current_slot] else {
            error!("cannot use ammo of empty head");
            return;
        };

        (head.health, head.health_max) = hp.get()
    }

    // returns new current head id
    pub fn loose_current(&mut self) -> Option<usize> {
        self.heads[self.current_slot] = None;
        self.next_head()
    }

    fn next_head(&mut self) -> Option<usize> {
        let start_idx = self.current_slot;

        for offset in 1..5 {
            let new_idx = (start_idx + offset) % 5;
            if let Some(head) = self.heads[new_idx] {
                self.current_slot = new_idx;
                return Some(head.head);
            }
        }

        None
    }

    pub fn contains(&self, head: usize) -> bool {
        self.heads.iter().any(|h| h.is_some_and(|h| h.head == head))
    }
}

#[derive(Event)]
pub struct HeadChanged(pub usize);

pub fn plugin(app: &mut App) {
    app.add_plugins(heads_ui::plugin);

    app.register_type::<ActiveHeads>();

    app.add_systems(OnEnter(GameState::Playing), setup);
    app.add_systems(
        FixedUpdate,
        (
            (reload, sync_hp).run_if(in_state(GameState::Playing)),
            on_select_active_head,
        )
            .run_if(is_server),
    );

    global_observer!(app, on_swap_backpack);
}

fn setup(mut commands: Commands, asset_server: Res<AssetServer>, heads: Res<HeadsDatabase>) {
    // TODO: load via asset loader
    let heads = (0usize..HEAD_COUNT)
        .map(|i| asset_server.load(format!("ui/heads/{}.png", heads.head_key(i))))
        .collect();

    commands.insert_resource(HeadsImages { heads });
}

fn sync_hp(mut query: Query<(&mut ActiveHeads, &Hitpoints)>) {
    for (mut active_heads, hp) in query.iter_mut() {
        if active_heads.hp().get() != hp.get() {
            active_heads.set_hitpoint(hp);
        }
    }
}

fn reload(
    mut commands: Commands,
    mut active: Query<&mut ActiveHeads>,
    time: Res<Time>,
    mut flags: Single<&mut AnimationFlags, With<Player>>,
) {
    for mut active in active.iter_mut() {
        if !active.reloading() {
            continue;
        }

        for head in active.heads.iter_mut() {
            let Some(head) = head else {
                continue;
            };

            if !head.has_ammo() && (head.last_use + head.reload_duration <= time.elapsed_secs()) {
                // only for player?

                use bevy_replicon::prelude::{SendMode, ServerTriggerExt, ToClients};
                commands.server_trigger(ToClients {
                    mode: SendMode::Broadcast,
                    message: PlaySound::Reloaded,
                });
                flags.restart_shooting = true;
                head.ammo = head.ammo_max;
            }
        }
    }
}

fn on_select_active_head(
    mut commands: Commands,
    mut query: Query<(&mut ActiveHeads, &mut Hitpoints), With<Player>>,
    mut select_lefts: MessageReader<FromClient<SelectLeftPressed>>,
    mut select_rights: MessageReader<FromClient<SelectRightPressed>>,
    controllers: ClientToController,
) {
    for press in select_lefts.read() {
        use bevy_replicon::prelude::{SendMode, ServerTriggerExt, ToClients};

        let player = controllers.get_controller(press.client_id);
        let (mut active_heads, mut hp) = query.get_mut(player).unwrap();

        active_heads.selected_slot = (active_heads.selected_slot + (HEAD_SLOTS - 1)) % HEAD_SLOTS;

        commands.server_trigger(ToClients {
            mode: SendMode::Broadcast,
            message: PlaySound::Selection,
        });

        if active_heads.head(active_heads.selected_slot).is_some() {
            active_heads.current_slot = active_heads.selected_slot;
            hp.set_health(active_heads.current().unwrap().health);

            commands.trigger(HeadChanged(
                active_heads.heads[active_heads.current_slot].unwrap().head,
            ));
        }
    }

    for press in select_rights.read() {
        use bevy_replicon::prelude::{SendMode, ServerTriggerExt, ToClients};

        let player = controllers.get_controller(press.client_id);
        let (mut active_heads, mut hp) = query.get_mut(player).unwrap();

        active_heads.selected_slot = (active_heads.selected_slot + 1) % HEAD_SLOTS;

        commands.server_trigger(ToClients {
            mode: SendMode::Broadcast,
            message: PlaySound::Selection,
        });

        if active_heads.head(active_heads.selected_slot).is_some() {
            active_heads.current_slot = active_heads.selected_slot;
            hp.set_health(active_heads.current().unwrap().health);

            commands.trigger(HeadChanged(
                active_heads.heads[active_heads.current_slot].unwrap().head,
            ));
        }
    }
}

fn on_swap_backpack(
    trigger: On<FromClient<BackpackSwapEvent>>,
    mut commands: Commands,
    mut query: Query<(&mut ActiveHeads, &mut Hitpoints, &mut Backpack), With<Player>>,
) {
    let backpack_slot = trigger.event().0;

    let Ok((mut active_heads, mut hp, mut backpack)) = query.single_mut() else {
        return;
    };

    let head = backpack.heads.get(backpack_slot).unwrap();

    let selected_slot = active_heads.selected_slot;

    let selected_head = active_heads.heads[selected_slot];
    active_heads.heads[selected_slot] = Some(*head);

    if let Some(old_active) = selected_head {
        backpack.heads[backpack_slot] = old_active;
    } else {
        backpack.heads.remove(backpack_slot);
    }

    hp.set_health(active_heads.current().unwrap().health);

    commands.trigger(HeadChanged(
        active_heads.heads[active_heads.selected_slot].unwrap().head,
    ));
}