Fix item energy bridge

energy-bridge/src/main/java/dev/su5ed/sinytra/connectorextras/energybridge/EnergyBridgeSetup.java

@@ -38,7 +38,7 @@ public static void init() {
             if (!stack.isEmpty() && !COMPUTING_CAPABILITY_LOCK.get()) {
                 COMPUTING_CAPABILITY_LOCK.set(true);
                 EnergyStorage storage = stack.getCapability(ForgeCapabilities.ENERGY, null)
-                    .map(ForgeEnergyStorageHandler::new)
+                    .map(s -> new ForgeEnergyStorageHandler(s, context, stack))
                     .orElse(null);
                 COMPUTING_CAPABILITY_LOCK.set(false);
                 return storage;
@@ -47,8 +47,7 @@ public static void init() {
         });
 
         MinecraftForge.EVENT_BUS.addGenericListener(BlockEntity.class, EnergyBridgeSetup::onAttachBlockEntityCapabilities);
-        // TODO Fix container item context
-//        MinecraftForge.EVENT_BUS.addGenericListener(ItemStack.class, EnergyBridgeSetup::onAttachItemStackCapabilities);
+        MinecraftForge.EVENT_BUS.addGenericListener(ItemStack.class, EnergyBridgeSetup::onAttachItemStackCapabilities);
     }
 
     private static void onAttachBlockEntityCapabilities(AttachCapabilitiesEvent<BlockEntity> event) {
@@ -76,10 +75,11 @@ private static void onAttachItemStackCapabilities(AttachCapabilitiesEvent<ItemSt
             public @NotNull <T> LazyOptional<T> getCapability(@NotNull Capability<T> cap, @Nullable Direction side) {
                 if (cap == ForgeCapabilities.ENERGY && !COMPUTING_CAPABILITY_LOCK.get()) {
                     COMPUTING_CAPABILITY_LOCK.set(true);
-                    EnergyStorage storage = EnergyStorage.ITEM.find(stack, ContainerItemContext.withConstant(stack));
+                    FabricEnergySlotHandler handler = new FabricEnergySlotHandler(stack);
+                    EnergyStorage storage = EnergyStorage.ITEM.find(stack, ContainerItemContext.ofSingleSlot(handler));
                     COMPUTING_CAPABILITY_LOCK.set(false);
                     if (storage != null) {
-                        return CAPS.computeIfAbsent(storage, b -> LazyOptional.of(() -> new FabricEnergyStorageHandler(storage))).cast();
+                        return CAPS.computeIfAbsent(storage, b -> LazyOptional.of(() -> new FabricEnergyStorageHandler(storage, stack, handler))).cast();
                     }
                 }
                 return LazyOptional.empty();

energy-bridge/src/main/java/dev/su5ed/sinytra/connectorextras/energybridge/FabricEnergySlotHandler.java

@@ -0,0 +1,56 @@
+package dev.su5ed.sinytra.connectorextras.energybridge;
+
+import net.fabricmc.fabric.api.transfer.v1.item.ItemVariant;
+import net.fabricmc.fabric.api.transfer.v1.storage.base.SingleSlotStorage;
+import net.fabricmc.fabric.api.transfer.v1.transaction.TransactionContext;
+import net.minecraft.world.item.ItemStack;
+
+public class FabricEnergySlotHandler implements SingleSlotStorage<ItemVariant> {
+
+    private ItemStack stack;
+
+    public FabricEnergySlotHandler(ItemStack stack) {
+        this.stack = stack;
+    }
+
+    @Override
+    public long insert(ItemVariant resource, long maxAmount, TransactionContext transaction) {
+        this.stack = resource.toStack();
+        return maxAmount;
+    }
+
+    @Override
+    public long extract(ItemVariant resource, long maxAmount, TransactionContext transaction) {
+        this.stack = resource.toStack();
+        return maxAmount;
+    }
+
+    @Override
+    public boolean isResourceBlank() {
+        return getResource().isBlank();
+    }
+
+    @Override
+    public ItemVariant getResource() {
+        return ItemVariant.of(stack);
+    }
+
+    public ItemStack getStack() {
+        return stack;
+    }
+
+    public void setStack(ItemStack stack) {
+        this.stack = stack;
+    }
+
+    @Override
+    public long getAmount() {
+        return stack.getCount();
+    }
+
+    @Override
+    public long getCapacity() {
+        return stack.getMaxStackSize();
+    }
+
+}

energy-bridge/src/main/java/dev/su5ed/sinytra/connectorextras/energybridge/FabricEnergyStorageHandler.java

@@ -1,14 +1,25 @@
 package dev.su5ed.sinytra.connectorextras.energybridge;
 
 import net.fabricmc.fabric.api.transfer.v1.transaction.Transaction;
+import net.minecraft.world.item.ItemStack;
 import net.minecraftforge.energy.IEnergyStorage;
 import team.reborn.energy.api.EnergyStorage;
 
 public class FabricEnergyStorageHandler implements IEnergyStorage {
     private final EnergyStorage storage;
+    private final ItemStack stack;
+    private final FabricEnergySlotHandler handler;
 
     public FabricEnergyStorageHandler(EnergyStorage storage) {
         this.storage = storage;
+        this.stack = null;
+        this.handler = null;
+    }
+
+    public FabricEnergyStorageHandler(EnergyStorage storage, ItemStack stack, FabricEnergySlotHandler handler) {
+        this.storage = storage;
+        this.stack = stack;
+        this.handler = handler;
     }
 
     @Override
@@ -18,9 +29,17 @@ public int receiveEnergy(int amount, boolean simulate) {
         }
         try (Transaction transaction = Transaction.openOuter()) {
             long e = EnergyBridge.convertForgeToFabricEnergy(amount);
+            long originalAmount = this.storage.getAmount();
             long inserted = (int) this.storage.insert(e, transaction);
             if (!simulate) {
                 transaction.commit();
+                if (stack != null) {
+                    stack.setTag(handler.getStack().getTag());
+                }
+            } else if (stack != null && originalAmount == (this.storage.getAmount() - inserted)) {
+                //Some implementations just commit the transaction by themselves, so it must be reverted afterwards
+                //We can't just extract the energy because the storage could be insert-only
+                handler.setStack(stack);
             }
             return EnergyBridge.unConvertForgeToFabricEnergy(inserted);
         }
@@ -36,6 +55,9 @@ public int extractEnergy(int amount, boolean simulate) {
             long extracted = (int) this.storage.extract(e, transaction);
             if (!simulate) {
                 transaction.commit();
+                if (stack != null) {
+                    stack.setTag(handler.getStack().getTag());
+                }
             }
             return EnergyBridge.convertFabricToForgeEnergy(extracted);
         }

energy-bridge/src/main/java/dev/su5ed/sinytra/connectorextras/energybridge/ForgeEnergyStorageHandler.java

@@ -1,14 +1,28 @@
 package dev.su5ed.sinytra.connectorextras.energybridge;
 
+import net.fabricmc.fabric.api.transfer.v1.context.ContainerItemContext;
+import net.fabricmc.fabric.api.transfer.v1.item.ItemVariant;
+import net.fabricmc.fabric.api.transfer.v1.transaction.Transaction;
 import net.fabricmc.fabric.api.transfer.v1.transaction.TransactionContext;
+import net.minecraft.world.item.ItemStack;
 import net.minecraftforge.energy.IEnergyStorage;
 import team.reborn.energy.api.EnergyStorage;
 
 public class ForgeEnergyStorageHandler implements EnergyStorage {
     private final IEnergyStorage storage;
+    private final ContainerItemContext context;
+    private final ItemStack stack;
 
     public ForgeEnergyStorageHandler(IEnergyStorage storage) {
         this.storage = storage;
+        this.context = null;
+        this.stack = null;
+    }
+
+    public ForgeEnergyStorageHandler(IEnergyStorage storage, ContainerItemContext context, ItemStack stack) {
+        this.storage = storage;
+        this.context = context;
+        this.stack = stack;
     }
 
     @Override
@@ -18,6 +32,7 @@ public long insert(long maxAmount, TransactionContext transaction) {
         transaction.addCloseCallback((context, result) -> {
             if (result.wasCommitted()) {
                 this.storage.receiveEnergy(fe, false);
+                updateStack();
             }
         });
         return EnergyBridge.unConvertFabricToForgeEnergy(inserted);
@@ -30,11 +45,27 @@ public long extract(long maxAmount, TransactionContext transaction) {
         transaction.addCloseCallback((context, result) -> {
             if (result.wasCommitted()) {
                 this.storage.extractEnergy(fe, false);
+                updateStack();
             }
         });
         return EnergyBridge.convertForgeToFabricEnergy(extracted);
     }
 
+    private void updateStack() {
+        if (stack != null) {
+            //We can't open a new Transaction while one is closing, so we make a new Thread
+            Thread thread = new Thread(() -> {
+                try (Transaction t = Transaction.openOuter()) {
+                    //Update the item (NBT).
+                    this.context.extract(this.context.getItemVariant(), 1, t);
+                    this.context.insert(ItemVariant.of(stack), 1, t);
+                    t.commit();
+                }
+            });
+            thread.start();
+        }
+    }
+
     @Override
     public long getAmount() {
         return EnergyBridge.convertForgeToFabricEnergy(this.storage.getEnergyStored());