From 9490a091b576155b8d60a9216a6d0192153143a3 Mon Sep 17 00:00:00 2001
From: AJ Ferguson <AJ-Ferguson@users.noreply.github.com>
Date: Tue, 23 Apr 2024 20:43:46 -0400
Subject: [PATCH] Calculate horse inventory size

---
 .../geysermc/geyser/entity/EntityDefinitions.java  |  2 +-
 .../type/living/animal/horse/LlamaEntity.java      | 13 +++++++++++++
 .../inventory/JavaHorseScreenOpenTranslator.java   | 14 ++++++++++++--
 3 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/core/src/main/java/org/geysermc/geyser/entity/EntityDefinitions.java b/core/src/main/java/org/geysermc/geyser/entity/EntityDefinitions.java
index 9563e6d37..f454bef88 100644
--- a/core/src/main/java/org/geysermc/geyser/entity/EntityDefinitions.java
+++ b/core/src/main/java/org/geysermc/geyser/entity/EntityDefinitions.java
@@ -943,7 +943,7 @@ public final class EntityDefinitions {
             LLAMA = EntityDefinition.inherited(LlamaEntity::new, chestedHorseEntityBase)
                     .type(EntityType.LLAMA)
                     .height(1.87f).width(0.9f)
-                    .addTranslator(MetadataType.INT, (entity, entityMetadata) -> entity.getDirtyMetadata().put(EntityDataTypes.STRENGTH, entityMetadata.getValue()))
+                    .addTranslator(MetadataType.INT, LlamaEntity::setStrength)
                     .addTranslator(MetadataType.INT, (entity, entityMetadata) -> entity.getDirtyMetadata().put(EntityDataTypes.VARIANT, entityMetadata.getValue()))
                     .build();
             TRADER_LLAMA = EntityDefinition.inherited(TraderLlamaEntity::new, LLAMA)
diff --git a/core/src/main/java/org/geysermc/geyser/entity/type/living/animal/horse/LlamaEntity.java b/core/src/main/java/org/geysermc/geyser/entity/type/living/animal/horse/LlamaEntity.java
index 7fcada504..29da2e0dc 100644
--- a/core/src/main/java/org/geysermc/geyser/entity/type/living/animal/horse/LlamaEntity.java
+++ b/core/src/main/java/org/geysermc/geyser/entity/type/living/animal/horse/LlamaEntity.java
@@ -25,16 +25,24 @@
 
 package org.geysermc.geyser.entity.type.living.animal.horse;
 
+import com.github.steveice10.mc.protocol.data.game.entity.metadata.type.IntEntityMetadata;
+import lombok.Getter;
 import org.cloudburstmc.math.vector.Vector3f;
 import org.cloudburstmc.protocol.bedrock.data.entity.EntityDataTypes;
 import org.geysermc.geyser.entity.EntityDefinition;
 import org.geysermc.geyser.item.Items;
 import org.geysermc.geyser.item.type.Item;
 import org.geysermc.geyser.session.GeyserSession;
+import org.geysermc.geyser.util.MathUtils;
 
 import java.util.UUID;
 
 public class LlamaEntity extends ChestedHorseEntity {
+    /**
+     * Used to calculate inventory size
+     */
+    @Getter
+    private int strength = 1;
 
     public LlamaEntity(GeyserSession session, int entityId, long geyserId, UUID uuid, EntityDefinition<?> definition, Vector3f position, Vector3f motion, float yaw, float pitch, float headYaw) {
         super(session, entityId, geyserId, uuid, definition, position, motion, yaw, pitch, headYaw);
@@ -42,6 +50,11 @@ public class LlamaEntity extends ChestedHorseEntity {
         dirtyMetadata.put(EntityDataTypes.CONTAINER_STRENGTH_MODIFIER, 3); // Presumably 3 slots for every 1 strength
     }
 
+    public void setStrength(IntEntityMetadata entityMetadata) {
+        strength = MathUtils.constrain(entityMetadata.getPrimitiveValue(), 1, 5);
+        this.dirtyMetadata.put(EntityDataTypes.STRENGTH, strength);
+    }
+
     @Override
     public boolean canEat(Item item) {
         return item == Items.WHEAT || item == Items.HAY_BLOCK;
diff --git a/core/src/main/java/org/geysermc/geyser/translator/protocol/java/inventory/JavaHorseScreenOpenTranslator.java b/core/src/main/java/org/geysermc/geyser/translator/protocol/java/inventory/JavaHorseScreenOpenTranslator.java
index 33cfeaf31..d2abcb5e3 100644
--- a/core/src/main/java/org/geysermc/geyser/translator/protocol/java/inventory/JavaHorseScreenOpenTranslator.java
+++ b/core/src/main/java/org/geysermc/geyser/translator/protocol/java/inventory/JavaHorseScreenOpenTranslator.java
@@ -29,6 +29,7 @@ import com.github.steveice10.mc.protocol.packet.ingame.clientbound.inventory.Cli
 import org.cloudburstmc.nbt.NbtMap;
 import org.cloudburstmc.nbt.NbtMapBuilder;
 import org.cloudburstmc.nbt.NbtType;
+import org.cloudburstmc.protocol.bedrock.data.entity.EntityFlag;
 import org.cloudburstmc.protocol.bedrock.data.inventory.ContainerType;
 import org.cloudburstmc.protocol.bedrock.packet.UpdateEquipPacket;
 import org.geysermc.geyser.entity.type.Entity;
@@ -114,16 +115,25 @@ public class JavaHorseScreenOpenTranslator extends PacketTranslator<ClientboundH
 
         // Since 1.20.5, the armor slot is not included in the container size,
         // but everything is still indexed the same.
-        int slotCount = packet.getNumberOfSlots() + 1;
+        int slotCount = 2; // Don't depend on slot count sent from server
 
         InventoryTranslator inventoryTranslator;
-        if (entity instanceof LlamaEntity) {
+        if (entity instanceof LlamaEntity llamaEntity) {
+            if (entity.getFlag(EntityFlag.CHESTED)) {
+                slotCount += llamaEntity.getStrength() * 3;
+            }
             inventoryTranslator = new LlamaInventoryTranslator(slotCount);
             slots.add(CARPET_SLOT);
         } else if (entity instanceof ChestedHorseEntity) {
+            if (entity.getFlag(EntityFlag.CHESTED)) {
+                slotCount += 15;
+            }
             inventoryTranslator = new DonkeyInventoryTranslator(slotCount);
             slots.add(SADDLE_SLOT);
         } else if (entity instanceof CamelEntity) {
+            if (entity.getFlag(EntityFlag.CHESTED)) {
+                slotCount += 15;
+            }
             // The camel has an invisible armor slot and needs special handling, same as the donkey
             inventoryTranslator = new DonkeyInventoryTranslator(slotCount);
             slots.add(SADDLE_SLOT);