From 26ba6810ade547816b0787d9ef91f16bb863f61f Mon Sep 17 00:00:00 2001
From: Tyfon <29051038+tyfon7@users.noreply.github.com>
Date: Fri, 21 Jun 2024 11:49:40 -0700
Subject: [PATCH] Stoppable multi-unload in raid

---
 MultiSelect.cs                | 17 +++++++++++++++++
 Patches/MultiSelectPatches.cs | 17 ++++++++++++++++-
 TaskSerializer.cs             | 24 ++++++++++++++++++++----
 3 files changed, 53 insertions(+), 5 deletions(-)

diff --git a/MultiSelect.cs b/MultiSelect.cs
index 9beebda..6fc04bd 100644
--- a/MultiSelect.cs
+++ b/MultiSelect.cs
@@ -16,6 +16,8 @@ namespace UIFixes
         private static readonly Dictionary<ItemContextClass, GridItemView> SelectedItems = [];
         private static readonly Dictionary<ItemContextClass, GridItemView> SecondaryItems = [];
 
+        private static ItemContextTaskSerializer UnloadSerializer = null;
+
         public static void Initialize()
         {
             // Grab the selection objects from ragfair as templates
@@ -257,14 +259,29 @@ namespace UIFixes
 
         public static void UnloadAmmoAll(ItemUiContext itemUiContext, bool allOrNothing)
         {
+            StopUnloading();
             if (!allOrNothing || InteractionCount(EItemInfoButton.UnloadAmmo, itemUiContext) == Count)
             {
+                // Call Initialize() before setting UnloadSerializer so that the initial synchronous call to StopProcesses()->StopUnloading() doesn't immediately cancel this
                 var taskSerializer = itemUiContext.GetOrAddComponent<ItemContextTaskSerializer>();
                 taskSerializer.Initialize(SortedItemContexts(), itemContext => itemUiContext.UnloadAmmo(itemContext.Item));
+
+                UnloadSerializer = taskSerializer;
                 itemUiContext.Tooltip?.Close();
             }
         }
 
+        public static void StopUnloading()
+        {
+            if (UnloadSerializer == null)
+            {
+                return;
+            }
+
+            UnloadSerializer.Cancel();
+            UnloadSerializer = null;
+        }
+
         private static void ShowSelection(GridItemView itemView)
         {
             GameObject selectedMark = itemView.transform.Find("SelectedMark")?.gameObject;
diff --git a/Patches/MultiSelectPatches.cs b/Patches/MultiSelectPatches.cs
index 4b505f2..26cbd97 100644
--- a/Patches/MultiSelectPatches.cs
+++ b/Patches/MultiSelectPatches.cs
@@ -1,6 +1,6 @@
 using Aki.Reflection.Patching;
-using Aki.Reflection.Utils;
 using Comfort.Common;
+using EFT;
 using EFT.Communications;
 using EFT.InventoryLogic;
 using EFT.UI;
@@ -55,6 +55,7 @@ namespace UIFixes
             // Actions
             new ItemViewClickPatch().Enable();
             new ContextActionsPatch().Enable();
+            new StopProcessesPatch().Enable();
 
             // GridView
             new GridViewCanAcceptPatch().Enable();
@@ -296,6 +297,20 @@ namespace UIFixes
             }
         }
 
+        public class StopProcessesPatch : ModulePatch
+        {
+            protected override MethodBase GetTargetMethod()
+            {
+                return AccessTools.DeclaredMethod(typeof(Player.PlayerInventoryController), nameof(Player.PlayerInventoryController.StopProcesses));
+            }
+
+            [PatchPostfix]
+            public static void Postfix()
+            {
+                MultiSelect.StopUnloading();
+            }
+        }
+
         // TradingItemView overrides GridItemView.OnClick and doesn't call base
         public class DeselectOnTradingItemViewClickPatch : ModulePatch
         {
diff --git a/TaskSerializer.cs b/TaskSerializer.cs
index 5fd203a..0e7b07e 100644
--- a/TaskSerializer.cs
+++ b/TaskSerializer.cs
@@ -26,9 +26,20 @@ namespace UIFixes
             return totalTask.Task;
         }
 
+        public void Cancel()
+        {
+            totalTask.TrySetCanceled();
+        }
+
         public void Update()
         {
-            if (!currentTask.IsCompleted)
+            if (currentTask.IsCanceled)
+            {
+                Complete();
+                return;
+            }
+
+            if (totalTask.Task.IsCompleted || !currentTask.IsCompleted)
             {
                 return;
             }
@@ -39,10 +50,15 @@ namespace UIFixes
             }
             else
             {
-                totalTask.Complete();
-                func = null;
-                Destroy(this);
+                Complete();
             }
         }
+
+        private void Complete()
+        {
+            totalTask.Complete();
+            func = null;
+            Destroy(this);
+        }
     }
 }