From b8431dea6ca2fee9be4b3a17bc3c95a64a1b074d Mon Sep 17 00:00:00 2001
From: lL1l1 <82986251+lL1l1@users.noreply.github.com>
Date: Sun, 17 Mar 2024 03:11:00 -0700
Subject: [PATCH] Track ACUs last taking damage for disconnect rules

---
 loc/US/strings_db.lua         |  6 +++---
 lua/aibrain.lua               | 12 +++---------
 lua/sim/units/ACUUnit.lua     | 14 ++++++++++++++
 lua/ui/lobby/lobbyOptions.lua |  6 +++---
 4 files changed, 23 insertions(+), 15 deletions(-)

diff --git a/loc/US/strings_db.lua b/loc/US/strings_db.lua
index a284c08cb04..11a060c0370 100644
--- a/loc/US/strings_db.lua
+++ b/loc/US/strings_db.lua
@@ -7727,7 +7727,7 @@ lobui_0796="Partial Share"
 lobui_0797="Your buildings and engineers will be transferred to your highest rated ally when you die.  Your other units will be destroyed when you die, except those captured by the enemy."
 
 lobui_0798="DC Share Conditions"
-lobui_0799="Set what happens to a player's units when they disconnect. In Assassination, only applies if an ACU is at full health/shield."
+lobui_0799="Set what happens to a player's units when they disconnect. In Assassination, only applies if an ACU has not been damaged in the last 2 minutes."
 lobui_0800="Same as Share Condition"
 lobui_0801="Treat disconnecting players the same as defeated players."
 
@@ -7740,9 +7740,9 @@ lobui_0806="Yes"
 lobui_0807="This game will not be rated."
 
 lobui_0808="Recall Disconnected ACUs"
-lobui_0809="Should disconnecting players' full health/shield ACUs be recalled, preventing their explosion?"
+lobui_0809="Should disconnecting players' ACUs be recalled, preventing their explosion if they were not damaged in the last 2 minutes?"
 lobui_0810="ACUs explode when their player disconnects."
-lobui_0811="ACUs that are at full health and shield are recalled when their player disconnects."
+lobui_0811="ACUs not damaged in the last 2 minutes are recalled when their player disconnects."
 
 aisettings_0001="AIx Cheat Multiplier"
 aisettings_0002="Set the cheat multiplier for the cheating AIs."
diff --git a/lua/aibrain.lua b/lua/aibrain.lua
index f76e9ba76c8..b53f60ac411 100644
--- a/lua/aibrain.lua
+++ b/lua/aibrain.lua
@@ -966,15 +966,9 @@ AIBrain = Class(AIBrainHQComponent, AIBrainStatisticsComponent, AIBrainJammerCom
 
                 local commanders = self:GetListOfUnits(categories.COMMAND, false)
                 for _, com in commanders do
-                    if com:GetHealth() == com:GetMaxHealth() then
-                        local comShield = com.MyShield
-                        if comShield then
-                            if comShield:GetHealth() == comShield:GetMaxHealth() then
-                                table.insert(safeCommanders, com)
-                            end
-                        else
-                            table.insert(safeCommanders, com)
-                        end
+                    -- 2 minutes since last damaged
+                    if com.LastTickDamaged == nil or com.LastTickDamaged + 1200 <= GetGameTick() then
+                        table.insert(safeCommanders, com)
                     end
                 end
 
diff --git a/lua/sim/units/ACUUnit.lua b/lua/sim/units/ACUUnit.lua
index 90680f0b4e6..67c3a1eb16b 100644
--- a/lua/sim/units/ACUUnit.lua
+++ b/lua/sim/units/ACUUnit.lua
@@ -1,6 +1,7 @@
 local CommandUnit = import("/lua/sim/units/commandunit.lua").CommandUnit
 
 ---@class ACUUnit : CommandUnit
+---@field LastTickDamaged number
 ACUUnit = ClassUnit(CommandUnit) {
     -- The "commander under attack" warnings.
     ---@param self ACUUnit
@@ -59,6 +60,19 @@ ACUUnit = ClassUnit(CommandUnit) {
         self.WeaponEnabled = {}
     end,
 
+    ---@param self ACUUnit
+    ---@param instigator Unit
+    ---@param amount number
+    ---@param vector Vector
+    ---@param damageType DamageType
+    OnDamage = function(self, instigator, amount, vector, damageType)
+        if self.CanTakeDamage and damageType ~= "TreeForce" and damageType ~= "TreeFire" then
+            self.LastTickDamaged = GetGameTick()
+        end
+
+        CommandUnit.OnDamage(self, instigator, amount, vector, damageType)
+    end,
+
     ---@param self ACUUnit
     ---@param instigator Unit
     ---@param amount number
diff --git a/lua/ui/lobby/lobbyOptions.lua b/lua/ui/lobby/lobbyOptions.lua
index b92e65a4df7..7709a88779a 100644
--- a/lua/ui/lobby/lobbyOptions.lua
+++ b/lua/ui/lobby/lobbyOptions.lua
@@ -236,7 +236,7 @@ globalOpts = {
     {
         default = 1,
         label = "<LOC lobui_0798>DC Share Conditions",
-        help = "<LOC lobui_0799>Set what happens to a player's units when they disconnect. In Assassination, only applies if an ACU is at full health/shield.",
+        help = "<LOC lobui_0799>Set what happens to a player's units when they disconnect. In Assassination, only applies if an ACU has not been damaged in the last 2 minutes.",
         key = 'AbandonmentShare',
         values = {
             {
@@ -279,7 +279,7 @@ globalOpts = {
     {
         default = 1,
         label = "<LOC lobui_0808>Recall Disconnected ACUs",
-        help = "<LOC lobui_0809>Should disconnecting players' full health/shield ACUs be recalled, preventing their explosion?",
+        help = "<LOC lobui_0809>Should disconnecting players' ACUs be recalled, preventing their explosion if they were not damaged in the last 2 minutes?",
         key = 'AbandonmentRecall',
         values = {
             {
@@ -289,7 +289,7 @@ globalOpts = {
             },
             {
                 text = "<LOC _Yes>Yes",
-                help = "<LOC lobui_0811>ACUs that are at full health and shield are recalled when their player disconnects.",
+                help = "<LOC lobui_0811>ACUs not damaged in the last 2 minutes are recalled when their player disconnects.",
                 key = true,
             },
         },