1 files changed, 37 insertions, 5 deletions

diff --git a/src/shared/ability.rs b/src/shared/ability.rs
index b841661..80c74b2 100644
--- a/src/shared/ability.rs
+++ b/src/shared/ability.rs
@@ -59,6 +59,7 @@ impl TargetedAbility {
 pub enum ActivatedAbility {
     Focus(Focus),
     Haste(Haste),
+    Shield(Shield),
     Speed(Speed),
 }
 
@@ -73,6 +74,12 @@ pub struct Haste {
 }
 
 #[derive(Copy, Clone, PartialEq, Debug, Deserialize, Serialize)]
+pub struct Shield {
+    pub duration: f32,
+    pub blocked_damage: f32,
+}
+
+#[derive(Copy, Clone, PartialEq, Debug, Deserialize, Serialize)]
 pub struct Speed {
     pub duration: f32,
 }
@@ -84,12 +91,13 @@ impl ActivatedAbility {
         match self {
             ActivatedAbility::Focus(focus) => focus_activation(focus),
             ActivatedAbility::Haste(haste) => haste_activation(haste),
+            ActivatedAbility::Shield(shield) => shield_activation(shield),
             ActivatedAbility::Speed(speed) => speed_activation(speed),
         }
     }
 }
 
-fn speed_activation(speed: Speed) -> ActivatedAbilityActivation {
+fn focus_activation(focus: Focus) -> ActivatedAbilityActivation {
     Box::new(move |commands: &mut Commands, source_id: PlayerId| {
         commands.add(move |world: &mut World| {
             world.run_system_once(
@@ -100,7 +108,8 @@ fn speed_activation(speed: Speed) -> ActivatedAbilityActivation {
                     let Ok(mut buffs) = buffses.get_mut(*entity_id) else {
                         return;
                     };
-                    buffs.speed = Some(Timer::from_seconds(speed.duration, TimerMode::Once));
+                    buffs.haste = Some(Timer::from_seconds(focus.duration, TimerMode::Once));
+                    buffs.speed = Some(Timer::from_seconds(focus.duration, TimerMode::Once));
                 },
             )
         });
@@ -125,7 +134,31 @@ fn haste_activation(haste: Haste) -> ActivatedAbilityActivation {
     })
 }
 
-fn focus_activation(focus: Focus) -> ActivatedAbilityActivation {
+fn shield_activation(shield: Shield) -> ActivatedAbilityActivation {
+    Box::new(move |commands: &mut Commands, source_id: PlayerId| {
+        commands.add(move |world: &mut World| {
+            world.run_system_once(
+                move |entity_map: Res<EntityMap>,
+                      mut healths: Query<&mut Health>,
+                      mut buffses: Query<&mut Buffs>| {
+                    let Some(entity_id) = entity_map.0.get(&source_id.0) else {
+                        return;
+                    };
+                    let Ok(mut health) = healths.get_mut(*entity_id) else {
+                        return;
+                    };
+                    let Ok(mut buffs) = buffses.get_mut(*entity_id) else {
+                        return;
+                    };
+                    health.shield = health.shield + shield.blocked_damage;
+                    buffs.shield = Some(Timer::from_seconds(shield.duration, TimerMode::Once));
+                },
+            )
+        });
+    })
+}
+
+fn speed_activation(speed: Speed) -> ActivatedAbilityActivation {
     Box::new(move |commands: &mut Commands, source_id: PlayerId| {
         commands.add(move |world: &mut World| {
             world.run_system_once(
@@ -136,8 +169,7 @@ fn focus_activation(focus: Focus) -> ActivatedAbilityActivation {
                     let Ok(mut buffs) = buffses.get_mut(*entity_id) else {
                         return;
                     };
-                    buffs.haste = Some(Timer::from_seconds(focus.duration, TimerMode::Once));
-                    buffs.speed = Some(Timer::from_seconds(focus.duration, TimerMode::Once));
+                    buffs.speed = Some(Timer::from_seconds(speed.duration, TimerMode::Once));
                 },
             )
         });