From 2b10ee6b8fba497815930d4d04c8b1bd61ad9b33 Mon Sep 17 00:00:00 2001
From: Vera Aguilera Puerto <zddm@outlook.es>
Date: Fri, 12 Feb 2021 11:27:11 +0100
Subject: [PATCH] Add cooldown to melee weapons on pickup/selected Fixes #3149
 Fixes kitchen knives not having ItemCooldown component.

---
 .../Weapon/Melee/MeleeWeaponComponent.cs      | 47 ++++++++++++++-----
 .../Entities/Objects/Weapons/Melee/knife.yml  |  1 +
 2 files changed, 37 insertions(+), 11 deletions(-)

diff --git a/Content.Server/GameObjects/Components/Weapon/Melee/MeleeWeaponComponent.cs b/Content.Server/GameObjects/Components/Weapon/Melee/MeleeWeaponComponent.cs
index c1ccbf089c..364837a117 100644
--- a/Content.Server/GameObjects/Components/Weapon/Melee/MeleeWeaponComponent.cs
+++ b/Content.Server/GameObjects/Components/Weapon/Melee/MeleeWeaponComponent.cs
@@ -19,7 +19,7 @@ using Robust.Shared.ViewVariables;
 namespace Content.Server.GameObjects.Components.Weapon.Melee
 {
     [RegisterComponent]
-    public class MeleeWeaponComponent : Component, IAttack
+    public class MeleeWeaponComponent : Component, IAttack, IHandSelected
     {
         [Dependency] private readonly IPhysicsManager _physicsManager = default!;
         [Dependency] private readonly IGameTiming _gameTiming = default!;
@@ -126,11 +126,7 @@ namespace Content.Server.GameObjects.Components.Weapon.Melee
             _lastAttackTime = curTime;
             _cooldownEnd = _lastAttackTime + TimeSpan.FromSeconds(ArcCooldownTime);
 
-            if (Owner.TryGetComponent(out ItemCooldownComponent cooldown))
-            {
-                cooldown.CooldownStart = _lastAttackTime;
-                cooldown.CooldownEnd = _cooldownEnd;
-            }
+            RefreshItemCooldown();
 
             return true;
         }
@@ -180,11 +176,7 @@ namespace Content.Server.GameObjects.Components.Weapon.Melee
             _lastAttackTime = curTime;
             _cooldownEnd = _lastAttackTime + TimeSpan.FromSeconds(CooldownTime);
 
-            if (Owner.TryGetComponent(out ItemCooldownComponent cooldown))
-            {
-                cooldown.CooldownStart = _lastAttackTime;
-                cooldown.CooldownEnd = _cooldownEnd;
-            }
+            RefreshItemCooldown();
 
             return true;
         }
@@ -211,6 +203,39 @@ namespace Content.Server.GameObjects.Components.Weapon.Melee
 
             return resSet;
         }
+
+        void IHandSelected.HandSelected(HandSelectedEventArgs eventArgs)
+        {
+            var curTime = _gameTiming.CurTime;
+            var cool = TimeSpan.FromSeconds(CooldownTime * 0.5f);
+
+            if (curTime < _cooldownEnd)
+            {
+                if (_cooldownEnd - curTime < cool)
+                {
+                    _lastAttackTime = curTime;
+                    _cooldownEnd += cool;
+                }
+                else
+                    return;
+            }
+            else
+            {
+                _lastAttackTime = curTime;
+                _cooldownEnd = curTime + cool;
+            }
+
+            RefreshItemCooldown();
+        }
+
+        private void RefreshItemCooldown()
+        {
+            if (Owner.TryGetComponent(out ItemCooldownComponent cooldown))
+            {
+                cooldown.CooldownStart = _lastAttackTime;
+                cooldown.CooldownEnd = _cooldownEnd;
+            }
+        }
     }
 
     public class MeleeHitMessage : ComponentMessage
diff --git a/Resources/Prototypes/Entities/Objects/Weapons/Melee/knife.yml b/Resources/Prototypes/Entities/Objects/Weapons/Melee/knife.yml
index 193afab293..a90199dd7a 100644
--- a/Resources/Prototypes/Entities/Objects/Weapons/Melee/knife.yml
+++ b/Resources/Prototypes/Entities/Objects/Weapons/Melee/knife.yml
@@ -6,6 +6,7 @@
   - type: Utensil
     types:
       - Knife
+  - type: ItemCooldown
   - type: MeleeWeapon
     hitSound: /Audio/Weapons/bladeslice.ogg
     damage: 12