From 7d06145a3485293cc5ba50a9c0caac3c3feb9b8b Mon Sep 17 00:00:00 2001
From: b3ni15 <ballobeni08@gmail.com>
Date: Mon, 15 Dec 2025 09:19:46 +0100
Subject: [PATCH] feat: Implement BuildingRemoveHook to manage building removal
 and prevent infinite loops fix: Enhance SaveTransferPacket handling for
 out-of-order delivery and reset transfer state

---
 Main.cs                             | 33 +++++++++++++++++
 Packets/Lobby/SaveTransferPacket.cs | 56 +++++++++++++----------------
 2 files changed, 58 insertions(+), 31 deletions(-)

diff --git a/Main.cs b/Main.cs
index 998e6be..a21f99c 100644
--- a/Main.cs
+++ b/Main.cs
@@ -1015,6 +1015,39 @@ namespace KCM
             }
         }
 
+        [HarmonyPatch(typeof(Building), "Remove")]
+        public class BuildingRemoveHook
+        {
+            public static void Prefix(Building __instance)
+            {
+                try
+                {
+                    // Skip if we're processing a remove packet (prevents infinite loop)
+                    if (Packets.Game.GameBuilding.BuildingRemovePacket.isProcessingPacket)
+                        return;
+
+                    if (KCClient.client.IsConnected)
+                    {
+                        // Only send packet if this building belongs to a player
+                        if (__instance.TeamID() >= 0)
+                        {
+                            Main.helper.Log($"Building {__instance.UniqueName} (guid: {__instance.guid}) being removed, sending packet");
+
+                            new Packets.Game.GameBuilding.BuildingRemovePacket()
+                            {
+                                guid = __instance.guid
+                            }.Send();
+                        }
+                    }
+                }
+                catch (Exception e)
+                {
+                    helper.Log($"Error in BuildingRemoveHook: {e.Message}");
+                    helper.Log(e.StackTrace);
+                }
+            }
+        }
+
         #endregion
 
         #region "Time Hooks"
diff --git a/Packets/Lobby/SaveTransferPacket.cs b/Packets/Lobby/SaveTransferPacket.cs
index 19dd74a..6e84f82 100644
--- a/Packets/Lobby/SaveTransferPacket.cs
+++ b/Packets/Lobby/SaveTransferPacket.cs
@@ -30,9 +30,11 @@ namespace KCM.Packets.Lobby
 
         public override void HandlePacketClient()
         {
-            if (chunkId == 0)
+            // Initialize on first chunk OR if arrays aren't properly sized yet
+            // This handles out-of-order packet delivery
+            if (!loadingSave || saveData.Length != saveSize || chunksReceived.Length != totalChunks)
             {
-                Main.helper.Log("Save Transfer started! Resetting static transfer state.");
+                Main.helper.Log($"Save Transfer initializing. saveSize={saveSize}, totalChunks={totalChunks}");
                 loadingSave = true;
 
                 saveData = new byte[saveSize];
@@ -42,6 +44,12 @@ namespace KCM.Packets.Lobby
                 ServerLobbyScript.LoadingSave.SetActive(true);
             }
 
+            // Skip if we already received this chunk (duplicate packet)
+            if (chunksReceived[chunkId])
+            {
+                Main.helper.Log($"[SaveTransfer] Duplicate chunk {chunkId} received, skipping.");
+                return;
+            }
 
             Array.Copy(saveDataChunk, 0, saveData, saveDataIndex, saveDataChunk.Length);
 
@@ -51,6 +59,7 @@ namespace KCM.Packets.Lobby
 
             Main.helper.Log($"[SaveTransfer] Processed chunk {chunkId}/{totalChunks}. Received: {received} bytes of {saveSize}.");
 
+            // Update progress bar
             if (saveSize > 0)
             {
                 float savePercent = (float)received / (float)saveSize;
@@ -68,26 +77,13 @@ namespace KCM.Packets.Lobby
                 ServerLobbyScript.ProgressText.text = "0.00 KB / 0.00 KB";
             }
 
-            if (!IsTransferComplete())
-            {
-                Main.helper.Log($"[SaveTransfer] Transfer not yet complete after chunk {chunkId}. Missing: {WhichIsNotComplete()}");
-            }
-
-            if (chunkId + 1 == totalChunks)
-            {
-                Main.helper.Log($"Received last save transfer packet. Final check: IsComplete={IsTransferComplete()}");
-
-                if (!IsTransferComplete())
-                {
-                    Main.helper.Log($"[SaveTransfer] WARNING: Transfer is NOT complete even after last chunk. Missing: {WhichIsNotComplete()}");
-                }
-            }
-
-
             if (IsTransferComplete())
             {
                 Main.helper.Log("Save Transfer complete!");
 
+                // Reset the loading state before processing
+                loadingSave = false;
+
                 LoadSaveLoadHook.saveBytes = saveData;
                 LoadSaveLoadHook.memoryStreamHook = true;
 
@@ -99,27 +95,25 @@ namespace KCM.Packets.Lobby
                 Broadcast.OnLoadedEvent.Broadcast(new OnLoadedEvent());
 
                 ServerLobbyScript.LoadingSave.SetActive(false);
+
+                // Reset static state for next transfer
+                ResetTransferState();
             }
         }
 
+        public static void ResetTransferState()
+        {
+            saveData = new byte[1];
+            chunksReceived = new bool[1];
+            loadingSave = false;
+            received = 0;
+        }
+
         public static bool IsTransferComplete()
         {
             return chunksReceived.All(x => x == true);
         }
 
-        public static string WhichIsNotComplete()
-        {
-            string notComplete = "";
-            for (int i = 0; i < chunksReceived.Length; i++)
-            {
-                if (!chunksReceived[i])
-                {
-                    notComplete += i + ", ";
-                }
-            }
-            return notComplete;
-        }
-
         public override void HandlePacketServer()
         {
         }