From 06890504a2bfcdbf218a9d2c2874fa1262d91cca Mon Sep 17 00:00:00 2001
From: chris <github@onechris.mozmail.com>
Date: Sun, 14 Jul 2024 21:55:57 +0200
Subject: [PATCH] Fix: Totem animation when playing totem effects manually
 (#4860)

* Fix: Totem animation for manually played totem effects

* Ensure we always reset the offhand correctly
---
 .../geyser/inventory/PlayerInventory.java     |  5 ++++
 .../inventory/item/StoredItemMappings.java    |  2 ++
 .../entity/JavaEntityEventTranslator.java     | 27 ++++++++++++++++++-
 .../geysermc/geyser/util/InventoryUtils.java  |  6 +++++
 4 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/core/src/main/java/org/geysermc/geyser/inventory/PlayerInventory.java b/core/src/main/java/org/geysermc/geyser/inventory/PlayerInventory.java
index 9bef4b08e..c3756d663 100644
--- a/core/src/main/java/org/geysermc/geyser/inventory/PlayerInventory.java
+++ b/core/src/main/java/org/geysermc/geyser/inventory/PlayerInventory.java
@@ -29,6 +29,7 @@ import lombok.Getter;
 import lombok.Setter;
 import org.checkerframework.checker.nullness.qual.NonNull;
 import org.geysermc.geyser.GeyserImpl;
+import org.geysermc.geyser.item.type.Item;
 import org.geysermc.geyser.session.GeyserSession;
 import org.geysermc.mcprotocollib.protocol.data.game.entity.player.Hand;
 import org.jetbrains.annotations.Range;
@@ -73,6 +74,10 @@ public class PlayerInventory extends Inventory {
         return items[36 + heldItemSlot];
     }
 
+    public boolean eitherHandMatchesItem(@NonNull Item item) {
+        return getItemInHand().asItem() == item || getItemInHand(Hand.OFF_HAND).asItem() == item;
+    }
+
     public void setItemInHand(@NonNull GeyserItemStack item) {
         if (36 + heldItemSlot > this.size) {
             GeyserImpl.getInstance().getLogger().debug("Held item slot was larger than expected!");
diff --git a/core/src/main/java/org/geysermc/geyser/inventory/item/StoredItemMappings.java b/core/src/main/java/org/geysermc/geyser/inventory/item/StoredItemMappings.java
index 05f6ba6cc..475a3e588 100644
--- a/core/src/main/java/org/geysermc/geyser/inventory/item/StoredItemMappings.java
+++ b/core/src/main/java/org/geysermc/geyser/inventory/item/StoredItemMappings.java
@@ -50,6 +50,7 @@ public class StoredItemMappings {
     private final ItemMapping milkBucket;
     private final ItemMapping powderSnowBucket;
     private final ItemMapping shield;
+    private final ItemMapping totem;
     private final ItemMapping upgradeTemplate;
     private final ItemMapping wheat;
     private final ItemMapping writableBook;
@@ -66,6 +67,7 @@ public class StoredItemMappings {
         this.milkBucket = load(itemMappings, Items.MILK_BUCKET);
         this.powderSnowBucket = load(itemMappings, Items.POWDER_SNOW_BUCKET);
         this.shield = load(itemMappings, Items.SHIELD);
+        this.totem = load(itemMappings, Items.TOTEM_OF_UNDYING);
         this.upgradeTemplate = load(itemMappings, Items.NETHERITE_UPGRADE_SMITHING_TEMPLATE);
         this.wheat = load(itemMappings, Items.WHEAT);
         this.writableBook = load(itemMappings, Items.WRITABLE_BOOK);
diff --git a/core/src/main/java/org/geysermc/geyser/translator/protocol/java/entity/JavaEntityEventTranslator.java b/core/src/main/java/org/geysermc/geyser/translator/protocol/java/entity/JavaEntityEventTranslator.java
index e119d39ce..6c2e02cd3 100644
--- a/core/src/main/java/org/geysermc/geyser/translator/protocol/java/entity/JavaEntityEventTranslator.java
+++ b/core/src/main/java/org/geysermc/geyser/translator/protocol/java/entity/JavaEntityEventTranslator.java
@@ -29,7 +29,9 @@ import org.cloudburstmc.protocol.bedrock.data.ParticleType;
 import org.cloudburstmc.protocol.bedrock.data.SoundEvent;
 import org.cloudburstmc.protocol.bedrock.data.entity.EntityDataTypes;
 import org.cloudburstmc.protocol.bedrock.data.entity.EntityEventType;
+import org.cloudburstmc.protocol.bedrock.data.inventory.ContainerId;
 import org.cloudburstmc.protocol.bedrock.packet.EntityEventPacket;
+import org.cloudburstmc.protocol.bedrock.packet.InventoryContentPacket;
 import org.cloudburstmc.protocol.bedrock.packet.LevelEventPacket;
 import org.cloudburstmc.protocol.bedrock.packet.LevelSoundEvent2Packet;
 import org.cloudburstmc.protocol.bedrock.packet.PlaySoundPacket;
@@ -42,11 +44,15 @@ import org.geysermc.geyser.entity.type.FishingHookEntity;
 import org.geysermc.geyser.entity.type.LivingEntity;
 import org.geysermc.geyser.entity.type.living.animal.ArmadilloEntity;
 import org.geysermc.geyser.entity.type.living.monster.WardenEntity;
+import org.geysermc.geyser.item.Items;
 import org.geysermc.geyser.session.GeyserSession;
+import org.geysermc.geyser.translator.inventory.InventoryTranslator;
 import org.geysermc.geyser.translator.protocol.PacketTranslator;
 import org.geysermc.geyser.translator.protocol.Translator;
+import org.geysermc.geyser.util.InventoryUtils;
 import org.geysermc.mcprotocollib.protocol.packet.ingame.clientbound.entity.ClientboundEntityEventPacket;
 
+import java.util.Collections;
 import java.util.concurrent.ThreadLocalRandom;
 
 @Translator(packet = ClientboundEntityEventPacket.class)
@@ -154,6 +160,16 @@ public class JavaEntityEventTranslator extends PacketTranslator<ClientboundEntit
                 entityEventPacket.setType(EntityEventType.WITCH_HAT_MAGIC); //TODO: CHECK
                 break;
             case TOTEM_OF_UNDYING_MAKE_SOUND:
+                // Bedrock will not play the spinning animation without the item in the hand o.o
+                // Fixes https://github.com/GeyserMC/Geyser/issues/2446
+                boolean totemItemWorkaround = !session.getPlayerInventory().eitherHandMatchesItem(Items.TOTEM_OF_UNDYING);
+                if (totemItemWorkaround) {
+                    InventoryContentPacket offhandPacket = new InventoryContentPacket();
+                    offhandPacket.setContainerId(ContainerId.OFFHAND);
+                    offhandPacket.setContents(Collections.singletonList(InventoryUtils.getTotemOfUndying().apply(session.getUpstream().getProtocolVersion())));
+                    session.sendUpstreamPacket(offhandPacket);
+                }
+
                 entityEventPacket.setType(EntityEventType.CONSUME_TOTEM);
 
                 PlaySoundPacket playSoundPacket = new PlaySoundPacket();
@@ -162,7 +178,16 @@ public class JavaEntityEventTranslator extends PacketTranslator<ClientboundEntit
                 playSoundPacket.setVolume(1.0F);
                 playSoundPacket.setPitch(1.0F + (ThreadLocalRandom.current().nextFloat() * 0.1F) - 0.05F);
                 session.sendUpstreamPacket(playSoundPacket);
-                break;
+
+                // Sent here early to ensure we have the totem in our hand
+                session.sendUpstreamPacket(entityEventPacket);
+
+                if (totemItemWorkaround) {
+                    // Reset the item again
+                    InventoryTranslator.PLAYER_INVENTORY_TRANSLATOR.updateSlot(session, session.getPlayerInventory(), 45);
+                }
+
+                return;
             case SHEEP_GRAZE_OR_TNT_CART_EXPLODE:
                 if (entity.getDefinition() == EntityDefinitions.SHEEP) {
                     entityEventPacket.setType(EntityEventType.EAT_GRASS);
diff --git a/core/src/main/java/org/geysermc/geyser/util/InventoryUtils.java b/core/src/main/java/org/geysermc/geyser/util/InventoryUtils.java
index a07e8deac..b0bfffc19 100644
--- a/core/src/main/java/org/geysermc/geyser/util/InventoryUtils.java
+++ b/core/src/main/java/org/geysermc/geyser/util/InventoryUtils.java
@@ -253,6 +253,12 @@ public class InventoryUtils {
                 .count(1).build();
     }
 
+    public static IntFunction<ItemData> getTotemOfUndying() {
+        return protocolVersion -> ItemData.builder()
+            .definition(Registries.ITEMS.forVersion(protocolVersion).getStoredItems().totem().getBedrockDefinition())
+            .count(1).build();
+    }
+
     /**
      * See {@link #findOrCreateItem(GeyserSession, String)}. This is for finding a specified {@link ItemStack}.
      *