Hopefully fix everything

src/main/java/systems/brn/serverstorage/blockentities/DisplayBlockEntity.java

@@ -8,16 +8,21 @@ import net.minecraft.block.BlockState;
 import net.minecraft.block.entity.BlockEntity;
 import net.minecraft.item.ItemStack;
 import net.minecraft.item.Items;
+import net.minecraft.nbt.NbtCompound;
+import net.minecraft.registry.RegistryWrapper;
 import net.minecraft.server.world.ServerWorld;
 import net.minecraft.storage.ReadView;
 import net.minecraft.storage.WriteView;
 import net.minecraft.text.Text;
 import net.minecraft.util.math.BlockPos;
+import net.minecraft.util.math.Vec3d;
 import net.minecraft.world.World;
 import org.joml.Vector3f;
 import systems.brn.serverstorage.lib.SortMode;
 import systems.brn.serverstorage.lib.StorageNetwork;
 
+import java.util.Map;
+
 import static systems.brn.serverstorage.ServerStorage.DISPLAY_BLOCK_ENTITY;
 import static systems.brn.serverstorage.blocks.DisplayBlock.FACING;
 
@@ -50,21 +55,50 @@ public class DisplayBlockEntity extends BlockEntity {
         // You can offset to make it float in front of the block
         Vector3f vec = new Vector3f(0.5f, 0.5f, 0.5f); // mutable
 
-        itemElement.setOffset(state.get(FACING).getDoubleVector().multiply(0.5)); // adjust Z based on block face
+// Get the block's facing direction as a Vec3d
+        Vec3d facingVec = state.get(FACING).getDoubleVector(); // Example: NORTH = (0, 0, -1)
+
+// Calculate outward offset for the item
+        Vec3d itemOffset = facingVec.multiply(0.5);
+        itemElement.setOffset(itemOffset);
+
+// Define global "up"
+        Vec3d up = new Vec3d(0, 1, 0);
+
+// Compute perpendicular "down" relative to the block face
+        Vec3d down = facingVec.crossProduct(up).normalize().multiply(0.3);
+
+// Final text offset: outward from face + downward relative to face
+        Vec3d textOffset = facingVec.normalize().multiply(0.5).add(down);
+        textDisplayElement.setOffset(textOffset.add(0,0.25,0));
+
+        switch (state.get(FACING)) {
+            case NORTH -> textDisplayElement.setRotation(0, 180);
+            case SOUTH -> textDisplayElement.setRotation(0, 0);
+            case WEST  -> textDisplayElement.setRotation(0, 90);
+            case EAST  -> textDisplayElement.setRotation(0, 270);
+            case UP    -> textDisplayElement.setRotation(-90, 0); // text points down
+            case DOWN  -> textDisplayElement.setRotation(90, 0);  // text points up
+        }
+
         // assign to const interface
         itemElement.setScale(vec); // Optional
         holder.addElement(itemElement);
+        holder.addElement(textDisplayElement);
     }
 
     @Override
     protected void readData(ReadView view) {
         super.readData(view);
-        this.targetItem = view.read("TargetItem", ItemStack.CODEC).orElse(ItemStack.EMPTY);
+        this.targetItem = ItemStack.EMPTY;
+        this.itemCount = 0;
+        view.read("TargetItem", ItemStack.CODEC).ifPresent(itemStack -> this.targetItem = itemStack.copy());
         this.itemCount = view.getInt("TargetItemCount", 0);
-        reindexDrives();
+    }
-        itemCount = network.itemStackMap.getOrDefault(targetItem, 0);
+
-        itemElement.setItem(targetItem);
+    @Override
-        textDisplayElement.setText(Text.of(String.valueOf(itemCount)));
+    public NbtCompound toInitialChunkDataNbt(RegistryWrapper.WrapperLookup registries) {
+        return createNbt(registries);
     }
 
     @Override
@@ -77,15 +111,32 @@ public class DisplayBlockEntity extends BlockEntity {
     @Override
     protected void writeData(WriteView view) {
         super.writeData(view);
-        view.put("TargetItem", ItemStack.CODEC, targetItem);
+
         view.putInt("TargetItemCount", itemCount);
+        view.putNullable("TargetItem", ItemStack.CODEC, targetItem.copy());
     }
 
     public static <T extends BlockEntity> void tick(World world, BlockPos blockPos, BlockState blockState, T t) {
         if (t instanceof DisplayBlockEntity displayBlockEntity) {
-            if (!displayBlockEntity.attached && !world.isClient()) {
+            if (!world.isClient()) {
-                displayBlockEntity.attached = true;
+                if (displayBlockEntity.network == null) {
-                ChunkAttachment.ofTicking(displayBlockEntity.holder, (ServerWorld) world, blockPos.toCenterPos());
+                    displayBlockEntity.reindexDrives();
+                    displayBlockEntity.itemCount = 0;
+                    for(Map.Entry<ItemStack, Integer> entry : displayBlockEntity.network.itemStackMap.entrySet()) {
+                        ItemStack key = entry.getKey();
+                        Integer value = entry.getValue();
+                        if (ItemStack.areItemsAndComponentsEqual(key, displayBlockEntity.targetItem)) {
+                            displayBlockEntity.itemCount = value;
+                            break;
+                        }
+                    }
+                    displayBlockEntity.itemElement.setItem(displayBlockEntity.targetItem);
+                    displayBlockEntity.textDisplayElement.setText(Text.of(String.valueOf(displayBlockEntity.itemCount)));
+                }
+                if (!displayBlockEntity.attached) {
+                    displayBlockEntity.attached = true;
+                    ChunkAttachment.ofTicking(displayBlockEntity.holder, (ServerWorld) world, blockPos.toCenterPos());
+                }
             }
         }
     }

src/main/java/systems/brn/serverstorage/blocks/DisplayBlock.java

@@ -34,6 +34,7 @@ import systems.brn.serverstorage.screens.DisplayBlockMangementScreen;
 import xyz.nucleoid.packettweaker.PacketContext;
 
 import java.util.HashMap;
+import java.util.Map;
 
 import static systems.brn.serverstorage.ServerStorage.*;
 import static systems.brn.serverstorage.lib.StorageOperations.*;
@@ -80,6 +81,7 @@ public class DisplayBlock extends ConnectedBlock implements PolymerTexturedBlock
                 modId,
                 FabricBlockEntityTypeBuilder.create(DisplayBlockEntity::new, DISPLAY_BLOCK).build(null)
         );
+
         DISPLAY_BLOCK_ENTITY.addSupportedBlock(DISPLAY_BLOCK);
         AttackBlockCallback.EVENT.register(DisplayBlock::onAttack);
         PolymerBlockUtils.registerBlockEntity(DISPLAY_BLOCK_ENTITY);
@@ -94,9 +96,17 @@ public class DisplayBlock extends ConnectedBlock implements PolymerTexturedBlock
             ItemStack insertedStack = displayBlockEntity.targetItem.copy();
             insertedStack.setCount(finalCount);
 
-            displayBlockEntity.itemCount = displayBlockEntity.network.itemStackMap.getOrDefault(displayBlockEntity.targetItem, 0);
+            displayBlockEntity.itemCount = 0;
+            for (Map.Entry<ItemStack, Integer> entry : displayBlockEntity.network.itemStackMap.entrySet()) {
+                ItemStack key = entry.getKey();
+                Integer value = entry.getValue();
+                if (ItemStack.areItemsAndComponentsEqual(key, displayBlockEntity.targetItem)) {
+                    displayBlockEntity.itemCount = value;
+                    break;
+                }
+            }
 
-            displayBlockEntity.network.removeItemStack(insertedStack);
+            finalCount -= displayBlockEntity.network.removeItemStack(insertedStack).getCount();
             displayBlockEntity.network.updateDisplays();
 
             displayBlockEntity.itemCount -= finalCount;
@@ -104,8 +114,6 @@ public class DisplayBlock extends ConnectedBlock implements PolymerTexturedBlock
             displayBlockEntity.itemElement.setItem(displayBlockEntity.targetItem);
             displayBlockEntity.textDisplayElement.setText(Text.of(String.valueOf(displayBlockEntity.itemCount)));
 
-            //TODO update the count
-
             int remainingToInsert = finalCount;
             for (int i = 0; i < Math.ceilDivExact(finalCount, displayBlockEntity.targetItem.getMaxCount()); i++) {
                 ItemStack cappedStack = insertedStack.copy();
@@ -137,25 +145,36 @@ public class DisplayBlock extends ConnectedBlock implements PolymerTexturedBlock
                 BlockEntity storageBlockEntity = world.getBlockEntity(pos);
                 if (storageBlockEntity instanceof DisplayBlockEntity displayBlockEntity) {
                     displayBlockEntity.reindexDrives();
-                    displayBlockEntity.itemCount = displayBlockEntity.network.itemStackMap.getOrDefault(displayBlockEntity.targetItem, 0);
+                    displayBlockEntity.itemCount = 0;
-                    displayBlockEntity.itemElement.setItem(displayBlockEntity.targetItem);
+                    for (Map.Entry<ItemStack, Integer> entry : displayBlockEntity.network.itemStackMap.entrySet()) {
-                    displayBlockEntity.textDisplayElement.setText(Text.of(String.valueOf(displayBlockEntity.itemCount)));
+                        ItemStack key = entry.getKey();
-                    if (player.isSneaking()) {
+                        Integer value = entry.getValue();
+                        if (ItemStack.areItemsAndComponentsEqual(key, displayBlockEntity.targetItem)) {
+                            displayBlockEntity.itemCount = value;
+                            break;
+                        }
+                    }
+
+                    ItemStack stack = player.getStackInHand(hand);
+                    if (stack.isEmpty() && !player.isSneaking()) {
                         DisplayBlockMangementScreen displayBlockMangementScreen = new DisplayBlockMangementScreen(player, displayBlockEntity);
                         displayBlockMangementScreen.updateDisplay();
                         displayBlockMangementScreen.open();
                     } else {
-                        ItemStack stack = player.getStackInHand(hand);
+                        ItemStack stack1 = stack.copy();
-                        int canPutIn = stack.getCount() - displayBlockEntity.network.putItemStackRemainder(stack);
+                        stack1.setCount(player.isSneaking() ? stack.getCount() : 1);
+                        int canPutIn = stack1.getCount() - displayBlockEntity.network.putItemStackRemainder(stack1);
                         if (canPutIn > 0) {
                             removeFromInventory(player.getInventory(), stack, canPutIn);
                         }
-                        //TODO add update the item count
+                        displayBlockEntity.itemElement.setItem(displayBlockEntity.targetItem);
+                        displayBlockEntity.itemCount += canPutIn;
+                        displayBlockEntity.textDisplayElement.setText(Text.of(String.valueOf(displayBlockEntity.itemCount)));
                     }
 
                 }
             }
-        return ActionResult.SUCCESS;
+            return ActionResult.SUCCESS;
         }
         return ActionResult.PASS;
     }
@@ -166,6 +185,7 @@ public class DisplayBlock extends ConnectedBlock implements PolymerTexturedBlock
         return new DisplayBlockEntity(pos, state);
     }
 
+
     @Nullable
     @Override
     public <T extends BlockEntity> BlockEntityTicker<T> getTicker(World world, BlockState state, BlockEntityType<T> type) {

src/main/java/systems/brn/serverstorage/screens/DisplayBlockMangementScreen.java

@@ -1,12 +1,15 @@
 package systems.brn.serverstorage.screens;
 
 import eu.pb4.sgui.api.elements.GuiElementBuilder;
+import net.minecraft.item.ItemStack;
 import net.minecraft.item.Items;
 import net.minecraft.server.network.ServerPlayerEntity;
 import net.minecraft.text.Text;
 import systems.brn.serverstorage.blockentities.DisplayBlockEntity;
 import systems.brn.serverstorage.lib.PagedGui;
 
+import java.util.Map;
+
 public class DisplayBlockMangementScreen extends PagedGui {
     public final DisplayBlockEntity displayBlockEntity;
 
@@ -26,11 +29,20 @@ public class DisplayBlockMangementScreen extends PagedGui {
         var builder = new GuiElementBuilder(displayBlockEntity.targetItem.getCount() > 0 ? displayBlockEntity.targetItem.getItem() : Items.AIR)
                 .setName(displayBlockEntity.targetItem.getItemName())
                 .setCallback((clickIndex, clickType, slotActionType) -> {
-                    displayBlockEntity.targetItem = getPlayer().currentScreenHandler.getCursorStack();
+                    displayBlockEntity.targetItem = getPlayer().currentScreenHandler.getCursorStack().copy();
 
                     displayBlockEntity.reindexDrives();
-                    displayBlockEntity.itemCount = displayBlockEntity.network.itemStackMap.getOrDefault(displayBlockEntity.targetItem, 0);
+
-                    displayBlockEntity.itemElement.setItem(displayBlockEntity.targetItem);
+                    displayBlockEntity.itemCount = 0;
+                    for(Map.Entry<ItemStack, Integer> entry : displayBlockEntity.network.itemStackMap.entrySet()) {
+                        ItemStack key = entry.getKey();
+                        Integer value = entry.getValue();
+                        if (ItemStack.areItemsAndComponentsEqual(key, displayBlockEntity.targetItem)) {
+                            displayBlockEntity.itemCount = value;
+                            break;
+                        }
+                    }
+                    displayBlockEntity.itemElement.setItem(displayBlockEntity.targetItem.copy());
                     displayBlockEntity.textDisplayElement.setText(Text.of(String.valueOf(displayBlockEntity.itemCount)));
                     displayBlockEntity.markDirty();
                     playClickSound(getPlayer());

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

@@ -173,6 +173,7 @@ public class StorageScreen extends PagedGui implements Searchable {
                 return false;
             }
         } else {
+            blockEntity.reindexDrives();
             blockEntity.openStorageScreens.add(this);
             blockEntity.updateDisplays();
         }