From 157be4de23679ffca1375ccbeb4f15d5302e0467 Mon Sep 17 00:00:00 2001
From: aldrenean <65921954+aldrenean@users.noreply.github.com>
Date: Thu, 14 Nov 2024 12:18:19 -0800
Subject: [PATCH] feat (shields rework): Allow deactivation of shield units
 (#3841)

* Enable the on/off toggle for all units with shields
* Set shield charge to 0 when disabled
* The above only apply when the "Shields Rework" experimental mod is enabled
* Toggle visual shield effects along with building state
---
 gamedata/alldefs_post.lua                  | 9 +++++++++
 luamocks/SpringLuaApi.lua                  | 4 ++++
 luarules/gadgets/lups_shield.lua           | 8 +++++++-
 luarules/gadgets/unit_shield_behaviour.lua | 4 ++++
 4 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/gamedata/alldefs_post.lua b/gamedata/alldefs_post.lua
index 7283385d816..0f1585be82e 100644
--- a/gamedata/alldefs_post.lua
+++ b/gamedata/alldefs_post.lua
@@ -860,6 +860,15 @@ function UnitDef_Post(name, uDef)
 		end
 	end
 
+	-- Shield Rework
+	if modOptions.shieldsrework == true and uDef.weapondefs then
+		for _, weapon in pairs(uDef.weapondefs) do
+			if weapon.shield and weapon.shield.repulser then
+				uDef.onoffable = true
+			end
+		end
+	end
+
 	--- EMP rework
 	if modOptions.emprework == true then
 		if name == "armstil" then
diff --git a/luamocks/SpringLuaApi.lua b/luamocks/SpringLuaApi.lua
index 1402cf43d98..3101303d8ed 100644
--- a/luamocks/SpringLuaApi.lua
+++ b/luamocks/SpringLuaApi.lua
@@ -1910,6 +1910,10 @@ assert(type(unitID) == "number","Argument unitID is of invalid type - expected n
 return  booleanMock
  end
 
+function Spring.GetUnitIsActive (unitID)
+assert(type(unitID) == "number","Argument unitID is of invalid type - expected number");
+return  booleanMock
+ end
 
 function Spring.SetLastMessagePosition (  x, y, z)
 assert(type(x) == "number","Argument x is of invalid type - expected number");
diff --git a/luarules/gadgets/lups_shield.lua b/luarules/gadgets/lups_shield.lua
index 013ce0eead9..a80ba843577 100644
--- a/luarules/gadgets/lups_shield.lua
+++ b/luarules/gadgets/lups_shield.lua
@@ -228,6 +228,12 @@ local function UpdateVisibility(unitID, unitData, unitVisible, forceUpdate)
 		unitVisible = GetVisibleSearch(ux, uz, unitData.search)
 	end
 
+	local unitIsActive = Spring.GetUnitIsActive(unitID) 
+	if unitIsActive ~= unitData.isActive then
+		forceUpdate = true
+		unitData.isActive = unitIsActive
+	end
+
 	if unitVisible == unitData.unitVisible and not forceUpdate then
 		return
 	end
@@ -237,7 +243,7 @@ local function UpdateVisibility(unitID, unitData, unitVisible, forceUpdate)
 		local fxID = unitData.fxTable[i]
 		local fx = Lups.GetParticles(fxID)
 		if fx then
-			fx.visibleToMyAllyTeam = unitVisible
+			fx.visibleToMyAllyTeam = unitIsActive and unitVisible
 		end
 	end
 end
diff --git a/luarules/gadgets/unit_shield_behaviour.lua b/luarules/gadgets/unit_shield_behaviour.lua
index 36524f01a0d..64263ed635d 100644
--- a/luarules/gadgets/unit_shield_behaviour.lua
+++ b/luarules/gadgets/unit_shield_behaviour.lua
@@ -37,6 +37,7 @@ local spGetUnitsInSphere = Spring.GetUnitsInSphere
 local spGetProjectilesInRectangle = Spring.GetProjectilesInRectangle
 local spGetUnitIsStunned = Spring.GetUnitIsStunned
 local spAreTeamsAllied = Spring.AreTeamsAllied
+local spGetUnitIsActive = Spring.GetUnitIsActive
 
 local shieldUnitDefs = {}
 local shieldUnitsData = {}
@@ -216,6 +217,9 @@ function gadget:GameFrame(frame)
 		end
 
 		if frame % 10 == 0 then
+			if not spGetUnitIsActive(shieldUnitID) then
+				spSetUnitShieldState(shieldUnitID, shieldData.shieldWeaponNumber, 0)
+			end
 			if not shieldData.shieldEnabled and shieldData.downtimeReset ~= 0 and shieldData.downtimeReset <= gameSeconds then
 				shieldData.downtimeReset = 0
 				shieldData.shieldEnabled = true