Add persistence in BlockEntity, refactor and make tha amount of pages correct when searching

gradle.properties

@@ -8,7 +8,7 @@ yarn_mappings=1.20.6+build.3
 loader_version=0.15.11
 
 # Mod Properties
-mod_version=1.9
+mod_version=2.0
 maven_group=systems.brn
 archives_base_name=Server_storage
 

src/main/java/systems/brn/server_storage/blockentities/StorageBlockEntity.java

@@ -6,6 +6,8 @@ import net.minecraft.entity.player.PlayerEntity;
 import net.minecraft.inventory.Inventory;
 import net.minecraft.inventory.SidedInventory;
 import net.minecraft.item.ItemStack;
+import net.minecraft.nbt.NbtCompound;
+import net.minecraft.registry.RegistryWrapper;
 import net.minecraft.util.math.BlockPos;
 import net.minecraft.util.math.Direction;
 import systems.brn.server_storage.lib.ConnectedChests;
@@ -14,7 +16,6 @@ import systems.brn.server_storage.screens.StorageScreen;
 import java.util.ArrayList;
 
 import static systems.brn.server_storage.ServerStorage.STORAGE_BLOCK_ENTITY;
-import static systems.brn.server_storage.lib.StorageOperations.*;
 
 public class StorageBlockEntity extends BlockEntity implements Inventory, SidedInventory, Runnable {
 
@@ -26,6 +27,10 @@ public class StorageBlockEntity extends BlockEntity implements Inventory, SidedI
 
     public StorageScreen storageScreenToDelete;
 
+    public Boolean sortAlphabetically = false;
+    public String searchString = "";
+    public int page = 0;
+
     public ArrayList<StorageScreen> openStorageScreens = new ArrayList<>();
 
     public StorageBlockEntity(BlockPos pos, BlockState state) {
@@ -60,8 +65,8 @@ public class StorageBlockEntity extends BlockEntity implements Inventory, SidedI
     @Override
     public void setStack(int slot, ItemStack stack) {
         markDirty();
-        ConnectedChests chests = getConnectedChests(world, this.pos);
+        ConnectedChests chests = new ConnectedChests(world, this.pos, false);
-        tryPutItemStackIntoChests(chests.getInventories(), stack);
+        chests.tryPutItemStack(stack);
         for (StorageScreen screen : openStorageScreens) {
             screen.updateDisplay();
         }
@@ -85,12 +90,33 @@ public class StorageBlockEntity extends BlockEntity implements Inventory, SidedI
 
     @Override
     public boolean canInsert(int slot, ItemStack stack, Direction dir) {
-        ConnectedChests chests = getConnectedChests(world, this.pos);
+        ConnectedChests chests = new ConnectedChests(world, this.pos, false);
-        return canPutItemStackIntoChests(chests.getInventories(), stack);
+        return chests.canAddItemStack(stack);
     }
 
     @Override
     public boolean canExtract(int slot, ItemStack stack, Direction dir) {
         return false; // Prevent extraction
     }
+
+    // Serialize the BlockEntity
+    @Override
+    public void writeNbt(NbtCompound nbt, RegistryWrapper.WrapperLookup wrapperLookup) {
+        // Save the current value of the number to the nbt
+        nbt.putInt("page", page);
+        nbt.putBoolean("sortAlphabetically", sortAlphabetically);
+        nbt.putString("searchString", searchString);
+
+        super.writeNbt(nbt, wrapperLookup);
+    }
+
+    // Deserialize the BlockEntity
+    @Override
+    public void readNbt(NbtCompound nbt, RegistryWrapper.WrapperLookup wrapperLookup) {
+        super.readNbt(nbt, wrapperLookup);
+
+        page = nbt.getInt("page");
+        sortAlphabetically = nbt.getBoolean("sortAlphabetically");
+        searchString = nbt.getString("searchString");
+    }
 }

src/main/java/systems/brn/server_storage/blocks/StorageBlock.java

@@ -12,7 +12,6 @@ import net.minecraft.block.entity.BlockEntityType;
 import net.minecraft.component.DataComponentTypes;
 import net.minecraft.component.type.WrittenBookContentComponent;
 import net.minecraft.entity.player.PlayerEntity;
-import net.minecraft.inventory.Inventory;
 import net.minecraft.item.ItemStack;
 import net.minecraft.item.Items;
 import net.minecraft.registry.Registries;
@@ -35,8 +34,6 @@ import systems.brn.server_storage.screens.StorageScreen;
 import java.util.List;
 
 import static systems.brn.server_storage.ServerStorage.*;
-import static systems.brn.server_storage.lib.StorageOperations.getCombinedInventoryFromChests;
-import static systems.brn.server_storage.lib.StorageOperations.getConnectedChests;
 import static systems.brn.server_storage.lib.Util.generateBookContent;
 
 public class StorageBlock extends Block implements PolymerTexturedBlock, BlockEntityProvider {
@@ -84,10 +81,8 @@ public class StorageBlock extends Block implements PolymerTexturedBlock, BlockEn
                 }
                 else if(player.getStackInHand(hand).getItem() == Items.WRITTEN_BOOK) {
                     ItemStack book = player.getStackInHand(hand);
-                    ConnectedChests chests = getConnectedChests(world, pos);
+                    ConnectedChests chests = new ConnectedChests(world, pos, true);
-                    List<Inventory> inventories = chests.getInventories();
+                    List<RawFilteredPair<Text>> generatedContent = generateBookContent(chests.inventory);
-                    Inventory inventory = getCombinedInventoryFromChests(inventories, true);
-                    List<RawFilteredPair<Text>> generatedContent = generateBookContent(inventory);
                     book.set(DataComponentTypes.WRITTEN_BOOK_CONTENT, new WrittenBookContentComponent(
                             RawFilteredPair.of("Item Listing"),
                             player.getGameProfile().getName(),

src/main/java/systems/brn/server_storage/lib/ConnectedChests.java

@@ -1,48 +1,150 @@
 package systems.brn.server_storage.lib;
 
+import net.minecraft.block.entity.BlockEntity;
+import net.minecraft.block.entity.HopperBlockEntity;
 import net.minecraft.inventory.Inventory;
+import net.minecraft.inventory.SimpleInventory;
+import net.minecraft.item.ItemStack;
+import net.minecraft.util.math.BlockPos;
+import net.minecraft.util.math.Direction;
+import net.minecraft.world.World;
+import systems.brn.server_storage.blocks.StorageBlock;
 
-import java.util.List;
+import java.util.*;
+
+import static systems.brn.server_storage.lib.StorageOperations.*;
 
 public class ConnectedChests {
-    private final List<Inventory> inventories;
+    public final List<Inventory> inventories;
+    public final Inventory inventory;
 
-    private final int containerCount;
+    public final int containerCount;
-    private int containerSlots = 0;
+    public final int containerSlots;
-    private int containerUsedSlots = 0;
+    public final int containerUsedSlots;
-    private int containerFreeSlots = 0;
+    public final int containerFreeSlots;
+
+    public ConnectedChests(World world, BlockPos startPos, boolean sortAlphabetically) {
+        List<Inventory> inventories = new ArrayList<>();
+        Set<BlockPos> visited = new HashSet<>();
+
+        getConnectedChestsHelper(world, startPos, inventories, visited);
 
-    ConnectedChests(List<Inventory> inventories){
         this.inventories = inventories;
         this.containerCount = inventories.size();
+        this.inventory = getCombinedInventory(sortAlphabetically);
+
+        int tempContainerSlots = 0;
+        int tempContainerUsedSlots = 0;
+        int tempContainerFreeSlots = 0;
         for (Inventory inventory : inventories) {
-            containerSlots += inventory.size();
+            tempContainerSlots += inventory.size();
             for (int slot = 0; slot < inventory.size(); slot++) {
                 if (inventory.getStack(slot).isEmpty()) {
-                    containerFreeSlots++;
+                    tempContainerFreeSlots++;
-                }
+                } else {
-                else {
+                    tempContainerUsedSlots++;
-                    containerUsedSlots++;
                 }
             }
-            assert containerSlots == containerUsedSlots + containerFreeSlots;
+            assert tempContainerSlots == tempContainerUsedSlots + tempContainerFreeSlots;
+        }
+        containerSlots = tempContainerSlots;
+        containerUsedSlots = tempContainerUsedSlots;
+        containerFreeSlots = tempContainerFreeSlots;
+    }
+
+    private static void getConnectedChestsHelper(World world, BlockPos pos, List<Inventory> inventories, Set<BlockPos> visited) {
+        if (visited.contains(pos)) {
+            return;
+        }
+        visited.add(pos);
+
+        BlockEntity blockEntity = world.getBlockEntity(pos);
+        if (blockEntity instanceof Inventory || world.getBlockState(pos).getBlock() instanceof StorageBlock) {
+            if (blockEntity instanceof Inventory && !(world.getBlockState(pos).getBlock() instanceof StorageBlock)) {
+                inventories.add((Inventory) blockEntity);
+            }
+
+            for (Direction direction : Direction.values()) {
+                BlockPos nextPos = pos.offset(direction);
+                getConnectedChestsHelper(world, nextPos, inventories, visited);
+            }
         }
     }
 
-    public int getContainerCount() {
+    // Modify getCombinedInventoryFromChests to include item metadata and combine stacks
-        return containerCount;
+    private SimpleInventory getCombinedInventory(boolean sortAlphabetically) {
+        Map<ItemStack, Integer> itemStackMap = new HashMap<>();
+        for (Inventory inventory : inventories) {
+            for (int i = 0; i < inventory.size(); i++) {
+                ItemStack stack = inventory.getStack(i);
+                if (!stack.isEmpty()) {
+                    // Check if there's an existing ItemStack with the same item and metadata
+                    boolean found = false;
+                    for (Map.Entry<ItemStack, Integer> entry : itemStackMap.entrySet()) {
+                        ItemStack existingStack = entry.getKey();
+                        if (ItemStack.areItemsAndComponentsEqual(stack, existingStack)) {
+                            int count = entry.getValue() + stack.getCount();
+                            itemStackMap.put(existingStack, count);
+                            found = true;
+                            break;
                         }
-    public int getContainerSlots() {
-        return containerSlots;
                     }
-    public int getContainerUsedSlots() {
+                    // If no existing stack with the same item and metadata, add a new entry
-        return containerUsedSlots;
+                    if (!found) {
+                        ItemStack copiedStack = stack.copy();
+                        itemStackMap.put(copiedStack, stack.getCount());
+                    }
+                }
+            }
+        }
+        return getSimpleInventory(itemStackMap.size(), itemStackMap, sortAlphabetically);
+    }
+
+    public boolean tryPutItemStack(ItemStack stack) {
+
+        // Iterate over each chest to try and insert the ItemStack
+        for (Inventory chest : inventories) {
+            if (!(chest instanceof HopperBlockEntity)) {
+                stack = insertStackIntoInventory(chest, stack);
+                if (stack.isEmpty()) {
+                    return true;
+                }
+            }
+        }
+
+        // If we still have remaining items, return false
+        return stack.isEmpty();
+    }
+
+    public boolean canAddItemStack(ItemStack stack) {
+        // Iterate over each chest to check if it's possible to insert the ItemStack
+        for (Inventory chest : inventories) {
+            if (!(chest instanceof HopperBlockEntity)) {
+                // Attempt to insert the ItemStack into the current chest
+                for (int i = 0; i < chest.size(); i++) {
+                    ItemStack slotStack = chest.getStack(i);
+                    if (slotStack.isEmpty() || canCombine(slotStack, stack)) {
+                        // If the slot is empty or can be combined with the ItemStack, return true
+                        return true;
+                    }
+                }
+            }
+        }
+
+        // If it's not possible to insert into any chest, return false
+        return false;
+    }
+
+    public void removeItemStack(ItemStack stackToRemove) {
+
+        int remainingToRemove = stackToRemove.getCount();
+
+        for (Inventory chest : inventories) {
+            remainingToRemove = removeFromInventory(chest, stackToRemove, remainingToRemove);
+            if (remainingToRemove <= 0) {
+                return;
             }
-    public int getContainerFreeSlots() {
-        return containerFreeSlots;
         }
 
-    public List<Inventory> getInventories() {
-        return inventories;
     }
 }

src/main/java/systems/brn/server_storage/lib/PagedGui.java

@@ -34,20 +34,15 @@ public abstract class PagedGui extends SimpleGui {
 
     public static final int PAGE_SIZE = 9 * 5;
     protected final Runnable closeCallback;
-    protected int page = 0;
     public boolean ignoreCloseCallback;
 
+    public int page = 0;
+
     public PagedGui(ServerPlayerEntity player, @Nullable Runnable closeCallback) {
         super(ScreenHandlerType.GENERIC_9X6, player, false);
         this.closeCallback = closeCallback;
     }
 
-
-    public void refreshOpen() {
-        this.updateDisplay();
-        this.open();
-    }
-
     @Override
     public void onClose() {
         if (this.closeCallback != null && !ignoreCloseCallback) {
@@ -74,7 +69,7 @@ public abstract class PagedGui extends SimpleGui {
     }
 
     protected void updateDisplay() {
-        var offset = this.page * PAGE_SIZE;
+        var offset = page * PAGE_SIZE;
 
         for (int i = 0; i < PAGE_SIZE; i++) {
             var element = this.getElement(offset + i);
@@ -118,16 +113,7 @@ public abstract class PagedGui extends SimpleGui {
             case 0 -> DisplayElement.previousPage(this);
             case 2 -> this.search();
             case 3 -> this.sorting();
-            case 4 -> DisplayElement.of(
+            case 4 -> this.refresh();
-                    new GuiElementBuilder(Items.PLAYER_HEAD)
-                            .setSkullOwner(GUI_REFRESH)
-                            .setName(Text.translatable("selectServer.refresh").formatted(Formatting.WHITE))
-                            .hideDefaultTooltip().noDefaults()
-                            .setCallback((x, y, z) -> {
-                                playClickSound(this.player);
-                                this.updateDisplay();
-                            })
-            );
             case 5 -> this.storeAll();
             case 6 -> DisplayElement.nextPage(this);
             case 8 -> DisplayElement.of(
@@ -136,13 +122,15 @@ public abstract class PagedGui extends SimpleGui {
                             .hideDefaultTooltip().noDefaults()
                             .setCallback((x, y, z) -> {
                                 playClickSound(this.player);
-                                this.close(this.closeCallback != null);
+                                this.close();
                             })
             );
             default -> DisplayElement.filler();
         };
     }
 
+    protected abstract DisplayElement refresh();
+
     protected abstract DisplayElement search();
 
     protected abstract DisplayElement sorting();

src/main/java/systems/brn/server_storage/lib/StorageOperations.java

@@ -1,199 +1,19 @@
 package systems.brn.server_storage.lib;
 
-import net.minecraft.block.entity.BlockEntity;
-import net.minecraft.block.entity.HopperBlockEntity;
 import net.minecraft.entity.player.PlayerEntity;
 import net.minecraft.entity.player.PlayerInventory;
 import net.minecraft.inventory.Inventory;
 import net.minecraft.inventory.SimpleInventory;
 import net.minecraft.item.Item;
 import net.minecraft.item.ItemStack;
-import net.minecraft.util.math.BlockPos;
-import net.minecraft.util.math.Direction;
-import net.minecraft.world.World;
 import org.jetbrains.annotations.NotNull;
-import systems.brn.server_storage.blocks.StorageBlock;
 
-import java.util.*;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.TreeMap;
 
 public class StorageOperations {
-    public static ConnectedChests getConnectedChests(World world, BlockPos startPos) {
-        List<Inventory> inventories = new ArrayList<>();
-        Set<BlockPos> visited = new HashSet<>();
-
-        getConnectedChestsHelper(world, startPos, inventories, visited);
-
-        return new ConnectedChests(inventories);
-    }
-
-    private static void getConnectedChestsHelper(World world, BlockPos pos, List<Inventory> inventories, Set<BlockPos> visited) {
-        if (visited.contains(pos)) {
-            return;
-        }
-        visited.add(pos);
-
-        BlockEntity blockEntity = world.getBlockEntity(pos);
-        if (blockEntity instanceof Inventory || world.getBlockState(pos).getBlock() instanceof StorageBlock) {
-            if (blockEntity instanceof Inventory && !(world.getBlockState(pos).getBlock() instanceof StorageBlock)) {
-                inventories.add((Inventory) blockEntity);
-            }
-
-            for (Direction direction : Direction.values()) {
-                BlockPos nextPos = pos.offset(direction);
-                getConnectedChestsHelper(world, nextPos, inventories, visited);
-            }
-        }
-    }
-
-    public static boolean tryPutItemStackIntoChests(List<Inventory> inventories, ItemStack stack) {
-
-        // Iterate over each chest to try and insert the ItemStack
-        for (Inventory chest : inventories) {
-            if(!(chest instanceof HopperBlockEntity)) {
-                stack = insertStackIntoInventory(chest, stack);
-                if (stack.isEmpty()) {
-                    return true;
-                }
-            }
-        }
-
-        // If we still have remaining items, return false
-        return stack.isEmpty();
-    }
-
-    public static boolean canPutItemStackIntoChests(List<Inventory> inventories, ItemStack stack) {
-        // Iterate over each chest to check if it's possible to insert the ItemStack
-        for (Inventory chest : inventories) {
-            if (!(chest instanceof HopperBlockEntity)) {
-                // Attempt to insert the ItemStack into the current chest
-                for (int i = 0; i < chest.size(); i++) {
-                    ItemStack slotStack = chest.getStack(i);
-                    if (slotStack.isEmpty() || canCombine(slotStack, stack)) {
-                        // If the slot is empty or can be combined with the ItemStack, return true
-                        return true;
-                    }
-                }
-            }
-        }
-
-        // If it's not possible to insert into any chest, return false
-        return false;
-    }
-
-    private static ItemStack insertStackIntoInventory(Inventory inventory, ItemStack stack) {
-        // First, try to merge with existing stacks
-        for (int i = 0; i < inventory.size(); i++) {
-            ItemStack slotStack = inventory.getStack(i);
-            if (canCombine(slotStack, stack)) {
-                int transferAmount = Math.min(stack.getCount(), slotStack.getMaxCount() - slotStack.getCount());
-                if (transferAmount > 0) {
-                    slotStack.increment(transferAmount);
-                    stack.decrement(transferAmount);
-                    inventory.markDirty();
-                    if (stack.isEmpty()) {
-                        return ItemStack.EMPTY;
-                    }
-                }
-            }
-        }
-
-        // Next, try to find an empty slot
-        for (int i = 0; i < inventory.size(); i++) {
-            ItemStack slotStack = inventory.getStack(i);
-            if (slotStack.isEmpty()) {
-                inventory.setStack(i, stack.copy());
-                stack.setCount(0);
-                inventory.markDirty();
-                return ItemStack.EMPTY;
-            }
-        }
-
-        return stack;
-    }
-
-    public static boolean canRemoveFromInventory(Inventory inventory, ItemStack stackToRemove) {
-        int remainingToRemove = stackToRemove.getCount();
-
-        for (int i = 0; i < inventory.size(); i++) {
-            ItemStack slotStack = inventory.getStack(i);
-            if (canCombine(slotStack, stackToRemove)) {
-                // If the slot contains the same item type
-                if (slotStack.getCount() >= remainingToRemove) {
-                    // If the count in the slot is sufficient to remove the requested amount
-                    return true;
-                } else {
-                    // If the count in the slot is not sufficient, update remainingToRemove
-                    remainingToRemove -= slotStack.getCount();
-                }
-            }
-        }
-        // If no matching stack with sufficient count is found, return false
-        return false;
-    }
-
-    public static void removeItemStackFromChests(List<Inventory> inventories, ItemStack stackToRemove) {
-
-        int remainingToRemove = stackToRemove.getCount();
-
-        for (Inventory chest : inventories) {
-            remainingToRemove = removeFromInventory(chest, stackToRemove, remainingToRemove);
-            if (remainingToRemove <= 0) {
-                return;
-            }
-        }
-
-    }
-
-    public static int removeFromInventory(Inventory inventory, ItemStack stackToRemove, int remainingToRemove) {
-        for (int i = 0; i < inventory.size(); i++) {
-            ItemStack slotStack = inventory.getStack(i);
-            if (canCombine(slotStack, stackToRemove)) {
-                int removeAmount = Math.min(slotStack.getCount(), remainingToRemove);
-                slotStack.decrement(removeAmount);
-                remainingToRemove -= removeAmount;
-                inventory.markDirty();
-
-                if (slotStack.isEmpty()) {
-                    inventory.setStack(i, ItemStack.EMPTY);
-                }
-
-                if (remainingToRemove <= 0) {
-                    return remainingToRemove;
-                }
-            }
-        }
-        return remainingToRemove;
-    }
-
-    // Modify getCombinedInventoryFromChests to include item metadata and combine stacks
-    public static SimpleInventory getCombinedInventoryFromChests(Iterable<Inventory> inventories, boolean sortAlphabetically) {
-        Map<ItemStack, Integer> itemStackMap = new HashMap<>();
-        for (Inventory inventory : inventories) {
-            for (int i = 0; i < inventory.size(); i++) {
-                ItemStack stack = inventory.getStack(i);
-                if (!stack.isEmpty()) {
-                    // Check if there's an existing ItemStack with the same item and metadata
-                    boolean found = false;
-                    for (Map.Entry<ItemStack, Integer> entry : itemStackMap.entrySet()) {
-                        ItemStack existingStack = entry.getKey();
-                        if (ItemStack.areItemsAndComponentsEqual(stack, existingStack)) {
-                            int count = entry.getValue() + stack.getCount();
-                            itemStackMap.put(existingStack, count);
-                            found = true;
-                            break;
-                        }
-                    }
-                    // If no existing stack with the same item and metadata, add a new entry
-                    if (!found) {
-                        ItemStack copiedStack = stack.copy();
-                        itemStackMap.put(copiedStack, stack.getCount());
-                    }
-                }
-            }
-        }
-        return getSimpleInventory(itemStackMap.size(), itemStackMap, sortAlphabetically);
-    }
-
     // Modify getSimpleInventory to include item metadata and sort conditionally
     public static @NotNull SimpleInventory getSimpleInventory(int size, Map<ItemStack, Integer> itemStackMap, boolean sortAlphabetically) {
         TreeMap<ItemStack, Integer> sortedMap;
@@ -235,7 +55,10 @@ public class StorageOperations {
     }
 
     // Modify filterInventory to include item metadata
-    public static SimpleInventory filterInventory(Inventory inventory, String query, boolean sortAlphabetically) {
+    public static Inventory filterInventory(Inventory inventory, String query, boolean sortAlphabetically) {
+        if(query == null || query.isEmpty()) {
+            return inventory;
+        }
         Map<ItemStack, Integer> itemStackMap = new HashMap<>();
         int itemCount = 0;
         for (int slot = 0; slot < inventory.size(); slot++) {
@@ -255,7 +78,7 @@ public class StorageOperations {
         Item item = stack.getItem();
         if (item != null) {
             String itemName = String.valueOf(item);
-            if (itemName != null && query != null && !query.isEmpty() && !itemName.contains(query)) {
+            if (itemName != null && !itemName.contains(query)) {
                 return ItemStack.EMPTY;
             }
         }
@@ -278,7 +101,7 @@ public class StorageOperations {
                 maxInsert += remainingSpace;
                 // If the maximum insertion count is greater than or equal to the item count, return the item count
                 if (maxInsert >= itemStack.getCount()) {
-                    return Math.min(itemStack.getCount(), maxInsert);
+                    return itemStack.getCount();
                 }
             }
         }
@@ -286,6 +109,78 @@ public class StorageOperations {
         return maxInsert; // Return the maximum insertion count
     }
 
+    public static boolean canRemoveFromInventory(Inventory inventory, ItemStack stackToRemove) {
+        int remainingToRemove = stackToRemove.getCount();
+
+        for (int i = 0; i < inventory.size(); i++) {
+            ItemStack slotStack = inventory.getStack(i);
+            if (canCombine(slotStack, stackToRemove)) {
+                // If the slot contains the same item type
+                if (slotStack.getCount() >= remainingToRemove) {
+                    // If the count in the slot is sufficient to remove the requested amount
+                    return true;
+                } else {
+                    // If the count in the slot is not sufficient, update remainingToRemove
+                    remainingToRemove -= slotStack.getCount();
+                }
+            }
+        }
+        // If no matching stack with sufficient count is found, return false
+        return false;
+    }
+
+    public static ItemStack insertStackIntoInventory(Inventory inventory, ItemStack stack) {
+        // First, try to merge with existing stacks
+        for (int i = 0; i < inventory.size(); i++) {
+            ItemStack slotStack = inventory.getStack(i);
+            if (canCombine(slotStack, stack)) {
+                int transferAmount = Math.min(stack.getCount(), slotStack.getMaxCount() - slotStack.getCount());
+                if (transferAmount > 0) {
+                    slotStack.increment(transferAmount);
+                    stack.decrement(transferAmount);
+                    inventory.markDirty();
+                    if (stack.isEmpty()) {
+                        return ItemStack.EMPTY;
+                    }
+                }
+            }
+        }
+
+        // Next, try to find an empty slot
+        for (int i = 0; i < inventory.size(); i++) {
+            ItemStack slotStack = inventory.getStack(i);
+            if (slotStack.isEmpty()) {
+                inventory.setStack(i, stack.copy());
+                stack.setCount(0);
+                inventory.markDirty();
+                return ItemStack.EMPTY;
+            }
+        }
+
+        return stack;
+    }
+
+    public static int removeFromInventory(Inventory inventory, ItemStack stackToRemove, int remainingToRemove) {
+        for (int i = 0; i < inventory.size(); i++) {
+            ItemStack slotStack = inventory.getStack(i);
+            if (canCombine(slotStack, stackToRemove)) {
+                int removeAmount = Math.min(slotStack.getCount(), remainingToRemove);
+                slotStack.decrement(removeAmount);
+                remainingToRemove -= removeAmount;
+                inventory.markDirty();
+
+                if (slotStack.isEmpty()) {
+                    inventory.setStack(i, ItemStack.EMPTY);
+                }
+
+                if (remainingToRemove <= 0) {
+                    return remainingToRemove;
+                }
+            }
+        }
+        return remainingToRemove;
+    }
+
     public static boolean canCombine(ItemStack stack1, ItemStack stack2) {
         return !stack1.isEmpty() && stack1.getItem() == stack2.getItem() && ItemStack.areItemsAndComponentsEqual(stack1, stack2);
     }

src/main/java/systems/brn/server_storage/screens/StorageScreen.java

@@ -13,23 +13,22 @@ import net.minecraft.util.Formatting;
 import net.minecraft.util.math.BlockPos;
 import net.minecraft.world.World;
 import org.jetbrains.annotations.Nullable;
+import systems.brn.server_storage.blockentities.StorageBlockEntity;
 import systems.brn.server_storage.lib.ConnectedChests;
 import systems.brn.server_storage.lib.PagedGui;
 
-import java.util.List;
-
 import static systems.brn.server_storage.lib.StorageOperations.*;
 import static systems.brn.server_storage.lib.Util.addCountToLore;
 import static systems.brn.server_storage.lib.Util.removeCountFromLore;
 
 public class StorageScreen extends PagedGui {
-    private Inventory inventory;
-    List<Inventory> inventories;
     private final BlockPos pos;
     private final ServerPlayerEntity player;
     private final World world;
-    private String query;
+    private ConnectedChests chests;
-    private boolean sortAlphabetically;
+    private final StorageBlockEntity blockEntity;
+    private String searchString;
+    private Boolean sortAlphabetically;
 
     public StorageScreen(ServerPlayerEntity player, BlockPos pos, @Nullable Runnable closeCallback) {
         super(player, closeCallback);
@@ -37,26 +36,34 @@ public class StorageScreen extends PagedGui {
         this.pos = pos;
         this.world = player.getWorld();
         this.setLockPlayerInventory(false);
+        this.blockEntity = (StorageBlockEntity) player.getWorld().getBlockEntity(pos);
+        assert blockEntity != null;
+    }
+
+    @Override
+    public boolean open() {
+        page = blockEntity.page;
+        searchString = blockEntity.searchString;
+        sortAlphabetically = blockEntity.sortAlphabetically;
         this.updateDisplay();
+        return super.open();
     }
 
     @Override
     public void updateDisplay() {
-        ConnectedChests chests = getConnectedChests(world, this.pos);
+        this.chests = new ConnectedChests(world, pos, sortAlphabetically);
-        this.inventories = chests.getInventories();
-        this.inventory = getCombinedInventoryFromChests(inventories, sortAlphabetically);
         String title = "Storage: " +
-                chests.getContainerUsedSlots() +
+                chests.containerUsedSlots +
                 "/" +
-                chests.getContainerSlots() +
+                chests.containerSlots +
                 "(" +
-                chests.getContainerFreeSlots() +
+                chests.containerFreeSlots +
                 " free)" +
                 "[" +
-                chests.getContainerCount() +
+                chests.containerCount +
                 "]";
 
-        this.setTitle(Text.of(title));
+        setTitle(Text.of(title));
         super.updateDisplay();
     }
 
@@ -64,13 +71,13 @@ public class StorageScreen extends PagedGui {
     protected int getPageAmount() {
         int sizeX = 9;
         int sizeY = 6;
-        return Math.ceilDivExact(inventory.size(), sizeX * sizeY);
+        return Math.ceilDivExact(filterInventory(chests.inventory, searchString, sortAlphabetically).size(), sizeX * sizeY);
     }
 
     @Override
     protected DisplayElement getElement(int id) {
-        if (this.inventory.size() > id) {
+        if (chests.inventory.size() > id) {
-            Inventory filteredInventory = filterInventory(this.inventory, query, this.sortAlphabetically);
+            Inventory filteredInventory = filterInventory(chests.inventory, searchString, sortAlphabetically);
             ItemStack itemStack = filteredInventory.getStack(id);
             ItemStack newStack = addCountToLore(itemStack.getCount(), itemStack);
             GuiElementBuilder guiElement = new GuiElementBuilder(newStack);
@@ -81,7 +88,7 @@ public class StorageScreen extends PagedGui {
 
     @Override
     public boolean onClick(int index, ClickType type, SlotActionType action, GuiElementInterface element) {
-        GuiElementInterface clickedElement = this.getSlot(index);
+        GuiElementInterface clickedElement = getSlot(index);
         if (clickedElement != null) {
             ItemStack clickedItem = clickedElement.getItemStack();
             ItemStack noLoreStack = removeCountFromLore(clickedItem);
@@ -92,7 +99,7 @@ public class StorageScreen extends PagedGui {
                     for (int i = 0; i < player.getInventory().main.size(); i++) {
                         ItemStack stack = player.getInventory().main.get(i);
                         if (ItemStack.areItemsAndComponentsEqual(stack, noLoreStack)) {
-                            insertItem(stack, 0, this.getVirtualSize(), false);
+                            insertItem(stack, 0, getVirtualSize(), false);
                         }
                     }
                 }
@@ -106,10 +113,10 @@ public class StorageScreen extends PagedGui {
                 int insertCount = canInsertItemIntoPlayerInventory(player, noLoreStack);
                 ItemStack insertingStack = noLoreStack.copy();
                 insertingStack.setCount(insertCount);
-                if (canRemoveFromInventory(inventory, noLoreStack) && insertCount > 0) {
+                if (canRemoveFromInventory(chests.inventory, noLoreStack) && insertCount > 0) {
                     player.getInventory().insertStack(insertingStack.copy());
-                    removeItemStackFromChests(inventories, insertingStack);
+                    chests.removeItemStack(insertingStack);
-                    this.updateDisplay();
+                    updateDisplay();
                 }
             }
         }
@@ -118,10 +125,10 @@ public class StorageScreen extends PagedGui {
 
     @Override
     public boolean insertItem(ItemStack stack, int startIndex, int endIndex, boolean fromLast) {
-        if (tryPutItemStackIntoChests(inventories, stack)) {
+        if (chests.tryPutItemStack(stack)) {
             removeFromInventory(player.getInventory(), stack, stack.getCount());
         }
-        this.updateDisplay();
+        updateDisplay();
         return super.insertItem(stack, startIndex, endIndex, fromLast);
     }
 
@@ -133,8 +140,8 @@ public class StorageScreen extends PagedGui {
                         .hideDefaultTooltip().noDefaults()
                         .setSkullOwner(GUI_QUESTION_MARK)
                         .setCallback((x, y, z) -> {
-                            SearchScreen searchScreen = new SearchScreen(this, "Search");
+                            SearchScreen searchScreen = new SearchScreen(this, "");
-                            playClickSound(this.getPlayer());
+                            playClickSound(getPlayer());
                             searchScreen.open();
                         })
         );
@@ -144,12 +151,12 @@ public class StorageScreen extends PagedGui {
     protected DisplayElement sorting() {
         return DisplayElement.of(
                 new GuiElementBuilder(Items.PLAYER_HEAD)
-                        .setName(Text.translatable(this.sortAlphabetically ? "A->Z" : "9->1").formatted(Formatting.WHITE))
+                        .setName(Text.translatable(sortAlphabetically ? "A->Z" : "9->1").formatted(Formatting.WHITE))
                         .hideDefaultTooltip().noDefaults()
-                        .setSkullOwner(this.sortAlphabetically ? GUI_A : GUI_1)
+                        .setSkullOwner(sortAlphabetically ? GUI_A : GUI_1)
                         .setCallback((x, y, z) -> {
                             this.sortAlphabetically ^= true;
-                            playClickSound(this.getPlayer());
+                            playClickSound(getPlayer());
                             updateDisplay();
                         })
         );
@@ -163,18 +170,43 @@ public class StorageScreen extends PagedGui {
                         .hideDefaultTooltip().noDefaults()
                         .setSkullOwner(GUI_STORE_ALL)
                         .setCallback((x, y, z) -> {
-                            playClickSound(this.player);
+                            playClickSound(player);
                             for (int i = 0; i < player.getInventory().main.size(); i++) {
                                 ItemStack stack = player.getInventory().main.get(i);
-                                insertItem(stack, 0, this.getVirtualSize(), false);
+                                insertItem(stack, 0, getVirtualSize(), false);
                             }
                         })
         );
     }
 
     public void doSearch(String query) {
-        this.query = query;
+        this.searchString = query;
         this.page = 0;
-        this.updateDisplay();
+        updateDisplay();
+    }
+
+    @Override
+    protected DisplayElement refresh() {
+        return DisplayElement.of(
+                new GuiElementBuilder(Items.PLAYER_HEAD)
+                        .setSkullOwner(GUI_REFRESH)
+                        .setName(Text.translatable("selectServer.refresh").formatted(Formatting.WHITE))
+                        .hideDefaultTooltip().noDefaults()
+                        .setCallback((x, y, z) -> {
+                            playClickSound(player);
+                            this.page = 0;
+                            this.searchString = "";
+                            updateDisplay();
+                        })
+        );
+    }
+
+    @Override
+    public void onClose() {
+        this.blockEntity.page = page;
+        this.blockEntity.searchString = searchString;
+        this.blockEntity.sortAlphabetically = sortAlphabetically;
+        this.blockEntity.markDirty();
+        super.onClose();
     }
 }