Tarkov-UIFixes/Patches/WeaponModdingPatches.cs

using Diz.LanguageExtensions;
using EFT.InventoryLogic;
using HarmonyLib;
using SPT.Reflection.Patching;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Reflection;

namespace UIFixes;

public static class WeaponModdingPatches
{
    public static void Enable()
    {
        new ResizePatch().Enable();
        new ResizeHelperPatch().Enable();
        new ResizeOperationRollbackPatch().Enable();
        new MoveBeforeNetworkTransactionPatch().Enable();

        //new ModdingMoveToSortingTablePatch().Enable();
        //new PresetMoveToSortingTablePatch().Enable();
    }

    public class ResizePatch : ModulePatch
    {
        public static MoveOperation NecessaryMoveOperation = null;

        private static bool InPatch = false;
        protected override MethodBase GetTargetMethod()
        {
            return AccessTools.Method(typeof(StashGridClass), nameof(StashGridClass.Resize));
        }

        [PatchPostfix]
        public static void Postfix(StashGridClass __instance, Item item, XYCellSizeStruct oldSize, XYCellSizeStruct newSize, bool simulate, ref bool __result)
        {
            if (__result || InPatch)
            {
                return;
            }

            if (item.Owner is not InventoryControllerClass inventoryController)
            {
                return;
            }

            LocationInGrid itemLocation = __instance.GetItemLocation(item);

            // The sizes passed in are the template sizes, need to make match the item's rotation
            XYCellSizeStruct actualOldSize = itemLocation.r.Rotate(oldSize);
            XYCellSizeStruct actualNewSize = itemLocation.r.Rotate(newSize);

            // Figure out which direction(s) its growing
            int horizontalGrowth = actualNewSize.X - actualOldSize.X;
            int verticalGrowth = actualNewSize.Y - actualOldSize.Y;

            // Can't move up/left more than the position
            horizontalGrowth = Math.Min(horizontalGrowth, itemLocation.x);
            verticalGrowth = Math.Min(verticalGrowth, itemLocation.y);

            // Try moving it
            try
            {
                InPatch = true;
                for (int x = 0; x <= horizontalGrowth; x++)
                {
                    for (int y = 0; y <= verticalGrowth; y++)
                    {
                        if (x + y == 0)
                        {
                            continue;
                        }

                        LocationInGrid newLocation = new(itemLocation.x - x, itemLocation.y - y, itemLocation.r);
                        ItemAddress newAddress = new GridItemAddress(__instance, newLocation);

                        var moveOperation = InteractionsHandlerClass.Move(item, newAddress, inventoryController, false);
                        if (moveOperation.Failed || moveOperation.Value == null)
                        {
                            continue;
                        }

                        bool resizeResult = __instance.Resize(item, oldSize, newSize, simulate);

                        // If simulating, rollback. Note that for some reason, only the Fold case even uses simulate
                        // The other cases (adding a mod, etc) never simulate, and then rollback later. Likely because there is normally
                        // no server side-effect of a resize - the only effect is updating the grid's free/used map. 
                        if (simulate || !resizeResult)
                        {
                            moveOperation.Value.RollBack();
                        }

                        if (resizeResult)
                        {
                            // Stash the move operation so it can be executed or rolled back later
                            NecessaryMoveOperation = moveOperation.Value;

                            __result = true;
                            return;
                        }
                    }
                }
            }
            finally
            {
                InPatch = false;
            }
        }
    }

    public class ResizeHelperPatch : ModulePatch
    {
        protected override MethodBase GetTargetMethod()
        {
            return AccessTools.Method(typeof(InteractionsHandlerClass), nameof(InteractionsHandlerClass.Resize_Helper));
        }

        [PatchPostfix]
        public static void Postfix(ref GStruct414<ResizeOperation> __result)
        {
            if (__result.Failed || __result.Value == null)
            {
                return;
            }

            if (ResizePatch.NecessaryMoveOperation != null)
            {
                __result.Value.SetMoveOperation(ResizePatch.NecessaryMoveOperation);
                ResizePatch.NecessaryMoveOperation = null;
            }
        }
    }

    public class ResizeOperationRollbackPatch : ModulePatch
    {
        protected override MethodBase GetTargetMethod()
        {
            return AccessTools.Method(typeof(ResizeOperation), nameof(ResizeOperation.RollBack));
        }

        [PatchPostfix]
        public static void Postfix(ResizeOperation __instance)
        {
            MoveOperation moveOperation = __instance.GetMoveOperation();
            if (moveOperation != null)
            {
                moveOperation.RollBack();
            }
        }
    }

    public class MoveBeforeNetworkTransactionPatch : ModulePatch
    {
        private static bool InPatch = false;

        protected override MethodBase GetTargetMethod()
        {
            return AccessTools.Method(typeof(TraderControllerClass), nameof(TraderControllerClass.RunNetworkTransaction));
        }

        [PatchPrefix]
        public static void Prefix(TraderControllerClass __instance, IRaiseEvents operationResult)
        {
            if (InPatch)
            {
                return;
            }

            MoveOperation extraOperation = null;
            if (operationResult is MoveOperation moveOperation)
            {
                extraOperation = moveOperation.R().AddOperation?.R().ResizeOperation?.GetMoveOperation();
            }
            else if (operationResult is FoldOperation foldOperation)
            {
                extraOperation = foldOperation.ResizeResult?.GetMoveOperation();
            }

            if (extraOperation != null)
            {
                try
                {
                    InPatch = true;
                    __instance.RunNetworkTransaction(extraOperation);
                }
                finally
                {
                    InPatch = false;
                }
            }
        }
    }

    // public class ModdingMoveToSortingTablePatch : ModulePatch
    // {
    //     private static bool InPatch = false;

    //     protected override MethodBase GetTargetMethod()
    //     {
    //         return AccessTools.DeclaredMethod(typeof(GClass2848), nameof(GClass2848.Select));
    //     }

    //     [PatchPostfix]
    //     public static void Postfix(GClass2848 __instance, Item item, ItemAddress itemAddress, bool simulate, ref Error error, ref bool __result)
    //     {
    //         if (!Settings.MoveBuildsToSortingTable.Value || InPatch || __result || error is not InteractionsHandlerClass.GClass3363)
    //         {
    //             return;
    //         }

    //         // get top level item (weapon)
    //         Item rootItem = itemAddress.Container.ParentItem.GetRootMergedItem();
    //         if (rootItem == null)
    //         {
    //             return;
    //         }

    //         // move it to sorting table
    //         SortingTableClass sortingTable = __instance.InventoryControllerClass.Inventory.SortingTable;
    //         if (sortingTable == null)
    //         {
    //             return;
    //         }

    //         ItemAddressClass sortingTableAddress = sortingTable.Grid.FindLocationForItem(rootItem);
    //         if (sortingTableAddress == null)
    //         {
    //             return;
    //         }

    //         var sortingTableMove = InteractionsHandlerClass.Move(rootItem, sortingTableAddress, __instance.InventoryControllerClass, simulate);
    //         if (sortingTableMove.Failed || sortingTableMove.Value == null)
    //         {
    //             return;
    //         }

    //         if (simulate)
    //         {
    //             // Just testing, and it was moveable to sorting table, so assume everything is fine.
    //             error = null;
    //             __result = true;
    //             return;
    //         }

    //         // Actually selecting it, so do it and then redo the select
    //         __instance.InventoryControllerClass.RunNetworkTransaction(sortingTableMove.Value);

    //         InPatch = true;
    //         __result = __instance.Select(item, itemAddress, simulate, out error);
    //         InPatch = false;
    //     }
    // }

    // public class PresetMoveToSortingTablePatch : ModulePatch
    // {
    //     private static bool InPatch = false;

    //     protected override MethodBase GetTargetMethod()
    //     {
    //         return AccessTools.Method(typeof(GClass2849), nameof(GClass2849.method_2));
    //     }

    //     [PatchPostfix]
    //     public static void Postfix(
    //         GClass2849 __instance,
    //         Item item,
    //         List<GClass2849.Class2174> modsWithSlots,
    //         TraderControllerClass itemController,
    //         bool simulate,
    //         ref GStruct416<List<GStruct414<GClass2802>>> __result)
    //     {
    //         if (!Settings.MoveBuildsToSortingTable.Value ||
    //             InPatch ||
    //             __result.Succeeded ||
    //             __result.Error is not InteractionsHandlerClass.GClass3363 ||
    //             itemController is not InventoryControllerClass inventoryController)
    //         {
    //             return;
    //         }

    //         // move it to sorting table
    //         SortingTableClass sortingTable = inventoryController.Inventory.SortingTable;
    //         if (sortingTable == null)
    //         {
    //             return;
    //         }

    //         ItemAddressClass sortingTableAddress = sortingTable.Grid.FindLocationForItem(item);
    //         if (sortingTableAddress == null)
    //         {
    //             return;
    //         }

    //         var sortingTableMove = InteractionsHandlerClass.Move(item, sortingTableAddress, inventoryController, simulate); // only called with simulate = false
    //         if (sortingTableMove.Failed || sortingTableMove.Value == null)
    //         {
    //             return;
    //         }

    //         InPatch = true;
    //         __result = __instance.method_2(item, modsWithSlots, itemController, simulate);
    //         InPatch = false;

    //         if (__result.Succeeded)
    //         {
    //             __result.Value.Prepend(sortingTableMove);
    //         }
    //         else
    //         {
    //             sortingTableMove.Value.RollBack();
    //         }
    //     }
    // }
}