From e3e78c5ef2aa5d639e4ae515f3b592417b6f8b95 Mon Sep 17 00:00:00 2001
From: FynnTW <fynn.papadopoulos@gmail.com>
Date: Thu, 21 Nov 2024 08:48:59 +0200
Subject: [PATCH] hero abilities fix

---
 M2TWEOP Code/M2TWEOP library/types/unit.cpp | 15 ++++++++++++++-
 M2TWEOP Code/M2TWEOP library/types/unit.h   | 13 ++++++++++++-
 2 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/M2TWEOP Code/M2TWEOP library/types/unit.cpp b/M2TWEOP Code/M2TWEOP library/types/unit.cpp
index c5fcdd25..0505b20d 100644
--- a/M2TWEOP Code/M2TWEOP library/types/unit.cpp	
+++ b/M2TWEOP Code/M2TWEOP library/types/unit.cpp	
@@ -451,10 +451,22 @@ namespace unitActions
         unitUseSpecialAbility(const unit* un) : thisUnit(un) {}
     };
     
-    void useSpecialAbility(const unit* un)
+    void useSpecialAbility(const unit* un, bool heroOnly)
     {
         if (un == nullptr)
             return;
+    	const heroAbility* ability = nullptr;
+    	if (un->generalInfo)
+    		ability = un->generalInfo->ability;
+    	if (!ability && heroOnly)
+			return;
+    	if (ability)
+    	{
+		    if (ability->usesLeft > 0 && ability->abilityTimer > 0)
+				return;
+    		if (heroOnly && ability->usesLeft == 0)
+				return;
+    	}
         const auto order = std::make_shared<unitUseSpecialAbility>(un);
         gameHelpers::fireGameScriptFunc(order.get(), codes::offsets.useSpecialAbility);
     }
@@ -1778,6 +1790,7 @@ void luaPlugin::initUnits()
 	/***
 	Makes the unit perform their special ability.
 	@function unit:useSpecialAbility
+	@tparam bool heroOnly
 	@usage
 		unit:useSpecialAbility();
 	*/
diff --git a/M2TWEOP Code/M2TWEOP library/types/unit.h b/M2TWEOP Code/M2TWEOP library/types/unit.h
index b8006c53..909380d1 100644
--- a/M2TWEOP Code/M2TWEOP library/types/unit.h	
+++ b/M2TWEOP Code/M2TWEOP library/types/unit.h	
@@ -613,6 +613,17 @@ struct generalInfo
 	char pad_0074[144]; //0x0074
 }; //Size: 0x0104
 
+struct heroAbility
+{
+	char *name; //0x0000
+	char pad_0004[4]; //0x0004
+	void* generalStuff; //0x0008
+	int32_t usesLeft; //0x000C
+	int32_t cooldown; //0x0010
+	int32_t abilityTimer; //0x0014
+	int32_t duration; //0x0018
+};
+
 struct unitMoveStruct
 {
 	int xCoord;
@@ -2098,7 +2109,7 @@ namespace unitActions
     void moveToMissileRange(unit* un, const unit* targetUnit, bool run);
     void unitTurn(unit* un, int16_t angle, bool isRelative);
     void taunt(const unit* un);
-    void useSpecialAbility(const unit* un);
+    void useSpecialAbility(const unit* un, bool heroOnly);
     int getSiegeEngineType(const unit* un);
     void attackBuilding(unit* un, buildingBattle* building);
     void collectEngine(unit* un, siegeEngine* engine);