From a9e734b147cb552fff695fd543eb8bf0a1297e71 Mon Sep 17 00:00:00 2001
From: NichtStudioCode <51272202+NichtStudioCode@users.noreply.github.com>
Date: Sun, 10 Nov 2024 17:39:52 +0100
Subject: [PATCH] Fix xenondevs/Nova#499

---
 .../addon/logistics/tileentity/StorageUnit.kt | 19 ++++++++++++-------
 1 file changed, 12 insertions(+), 7 deletions(-)

diff --git a/logistics/src/main/kotlin/xyz/xenondevs/nova/addon/logistics/tileentity/StorageUnit.kt b/logistics/src/main/kotlin/xyz/xenondevs/nova/addon/logistics/tileentity/StorageUnit.kt
index 5665c22..e1ddb51 100644
--- a/logistics/src/main/kotlin/xyz/xenondevs/nova/addon/logistics/tileentity/StorageUnit.kt
+++ b/logistics/src/main/kotlin/xyz/xenondevs/nova/addon/logistics/tileentity/StorageUnit.kt
@@ -11,6 +11,7 @@ import xyz.xenondevs.cbf.Compound
 import xyz.xenondevs.commons.provider.mutable.MutableProvider
 import xyz.xenondevs.invui.gui.Gui
 import xyz.xenondevs.invui.inventory.VirtualInventory
+import xyz.xenondevs.invui.inventory.event.ItemPostUpdateEvent
 import xyz.xenondevs.invui.inventory.event.ItemPreUpdateEvent
 import xyz.xenondevs.invui.item.ItemProvider
 import xyz.xenondevs.invui.item.builder.ItemBuilder
@@ -20,7 +21,6 @@ import xyz.xenondevs.nova.addon.logistics.registry.Blocks.STORAGE_UNIT
 import xyz.xenondevs.nova.ui.menu.sideconfig.OpenSideConfigItem
 import xyz.xenondevs.nova.ui.menu.sideconfig.SideConfigMenu
 import xyz.xenondevs.nova.util.item.takeUnlessEmpty
-import xyz.xenondevs.nova.util.runTaskLater
 import xyz.xenondevs.nova.world.BlockPos
 import xyz.xenondevs.nova.world.block.state.NovaBlockState
 import xyz.xenondevs.nova.world.block.tileentity.NetworkedTileEntity
@@ -35,7 +35,7 @@ class StorageUnit(pos: BlockPos, state: NovaBlockState, data: Compound) : Networ
     
     private val inventory = StorageUnitInventory(storedValue("type", true, ItemStack::empty), storedValue("amount", true) { 0 })
     private val inputInventory = VirtualInventory(null, 1).apply { setPreUpdateHandler(::handleInputInventoryUpdate) }
-    private val outputInventory = VirtualInventory(null, 1).apply { setPreUpdateHandler(::handleOutputInventoryUpdate) }
+    private val outputInventory = VirtualInventory(null, 1).apply { setPreUpdateHandler(::handlePreOutputInventoryUpdate); setPostUpdateHandler(::handlePostOutputInventoryUpdate) }
     
     init {
         storedItemHolder(inventory to NetworkConnectionType.BUFFER)
@@ -46,18 +46,23 @@ class StorageUnit(pos: BlockPos, state: NovaBlockState, data: Compound) : Networ
             event.isCancelled = true
     }
     
-    private fun handleOutputInventoryUpdate(event: ItemPreUpdateEvent) {
+    private fun handlePreOutputInventoryUpdate(event: ItemPreUpdateEvent) {
         if (event.updateReason == SELF_UPDATE_REASON)
             return
         
-        if (event.isAdd || event.isSwap) {
+        if (!event.isRemove) {
             event.isCancelled = true
-        } else if (event.isRemove && !inventory.type.isEmpty) {
-            inventory.take(0, event.removedAmount)
-            runTaskLater(1) { menuContainer.forEachMenu(StorageUnitMenu::update) }
         }
     }
     
+    private fun handlePostOutputInventoryUpdate(event: ItemPostUpdateEvent) {
+        if (event.updateReason == SELF_UPDATE_REASON)
+            return
+        
+        // preUpdateHandler enforces that only remove is possible
+        inventory.take(0, event.removedAmount)
+    }
+    
     private fun updateOutputSlot() {
         if (inventory.type.isEmpty) {
             outputInventory.setItem(SELF_UPDATE_REASON, 0, null)