Fix grid reorder reshow

Patches/ReorderGridsPatch.cs

@@ -9,6 +9,15 @@ using UnityEngine;
 
 namespace UIFixes
 {
+    /* There are 3 cases to handle in TemplatedGridsView.Show
+     * 1. An item is shown for the first time
+     *      - It renders on its own, and the UI is correct
+     *      - Use the UI to sort Grids, and update GridViews to match
+     * 2. An item is shown for the 2nd+ time in a new context
+     *      - The GridViews will be recreated, so in the prefix we need to reorder them to match the Grids order
+     * 3. An existing TemplatedGridsView is reshown
+     *      - Everything is already in place, no action needed
+    */
     public class ReorderGridsPatch : ModulePatch
     {
         private static readonly HashSet<string> ReorderedItems = [];
@@ -20,13 +29,15 @@ namespace UIFixes
         }
 
         [PatchPrefix]
-        public static void Prefix(LootItemClass compoundItem, ref GridView[] ____presetGridViews)
+        public static void Prefix(TemplatedGridsView __instance, LootItemClass compoundItem, ref GridView[] ____presetGridViews)
         {
             if (!Settings.ReorderGrids.Value)
             {
+                // To properly support disabling this feature:
+                // 1. Items that sorted their Grids need to return them to original order
+                // 2. If this TemplatedGridsView was sorted, it needs to be unsorted to match
                 if (ReorderedItems.Contains(compoundItem.Id) && GridMaps.TryGetValue(compoundItem.TemplateId, out int[] unwantedMap))
                 {
-                    // Put it back
                     StashGridClass[] orderedGrids = new StashGridClass[compoundItem.Grids.Length];
                     for (int i = 0; i < compoundItem.Grids.Length; i++)
                     {
@@ -35,31 +46,40 @@ namespace UIFixes
 
                     compoundItem.Grids = orderedGrids;
                     ReorderedItems.Remove(compoundItem.Id);
+
+                    if (IsSorted(__instance))
+                    {
+                        GridView[] orderedGridView = new GridView[____presetGridViews.Length];
+                        for (int i = 0; i < ____presetGridViews.Length; i++)
+                        {
+                            orderedGridView[i] = ____presetGridViews[unwantedMap[i]];
+                        }
+
+                        ____presetGridViews = orderedGridView;
+                        MarkSorted(__instance, false);
+                    }
                 }
 
                 return;
             }
 
-            if (GridMaps.TryGetValue(compoundItem.TemplateId, out int[] map))
+            if (ReorderedItems.Contains(compoundItem.Id) && !IsSorted(__instance))
             {
-                GridView[] orderedGridView = new GridView[____presetGridViews.Length];
+                // This is a new context of a sorted Item, need to presort the GridViews
-                for (int i = 0; i < ____presetGridViews.Length; i++)
+                if (GridMaps.TryGetValue(compoundItem.TemplateId, out int[] map))
                 {
-                    orderedGridView[map[i]] = ____presetGridViews[i];
+                    GridView[] orderedGridView = new GridView[____presetGridViews.Length];
-                }
+                    for (int i = 0; i < ____presetGridViews.Length; i++)
-
-                ____presetGridViews = orderedGridView;
-
-                if (!ReorderedItems.Contains(compoundItem.Id))
-                {
-                    StashGridClass[] orderedGrids = new StashGridClass[compoundItem.Grids.Length];
-                    for (int i = 0; i < compoundItem.Grids.Length; i++)
                     {
-                        orderedGrids[map[i]] = compoundItem.Grids[i];
+                        orderedGridView[map[i]] = ____presetGridViews[i];
                     }
 
-                    compoundItem.Grids = orderedGrids;
+                    ____presetGridViews = orderedGridView;
-                    ReorderedItems.Add(compoundItem.Id);
+                    MarkSorted(__instance);
+                }
+                else
+                {
+                    Logger.LogError($"Item {compoundItem.Id}, tpl: {compoundItem.TemplateId} has sorted Grids but no map to sort GridViews!");
                 }
             }
         }
@@ -108,7 +128,38 @@ namespace UIFixes
 
             compoundItem.Grids = sorted.Select(pair => pair.Key).ToArray();
             ____presetGridViews = orderedGridViews;
+
             ReorderedItems.Add(compoundItem.Id);
+            MarkSorted(__instance);
         }
+
+        private static bool IsSorted(TemplatedGridsView view)
+        {
+            return view != null && view.GetComponent<SortedMarker>() != null;
+        }
+
+        private static void MarkSorted(TemplatedGridsView view, bool marked = true)
+        {
+            if (view == null)
+            {
+                return;
+            }
+
+
+            if (marked)
+            {
+                view.GetOrAddComponent<SortedMarker>();
+            }
+            else
+            {
+                SortedMarker marker = view.GetComponent<SortedMarker>();
+                if (marker != null)
+                {
+                    UnityEngine.Object.DestroyImmediate(marker);
+                }
+            }
+        }
+
+        public class SortedMarker : MonoBehaviour { }
     }
 }