From 3656395ce11b028ae3c1c5b8fbe4f783664b846d Mon Sep 17 00:00:00 2001
From: AJ Ferguson <AJ-Ferguson@users.noreply.github.com>
Date: Thu, 25 Apr 2024 21:03:17 -0400
Subject: [PATCH] Armadillo states

---
 .../entity/properties/type/EnumProperty.java  |  2 +-
 .../type/living/animal/ArmadilloEntity.java   | 25 ++++++++++++++++---
 .../entity/JavaEntityEventTranslator.java     |  6 +++++
 gradle/libs.versions.toml                     |  2 +-
 4 files changed, 30 insertions(+), 5 deletions(-)

diff --git a/core/src/main/java/org/geysermc/geyser/entity/properties/type/EnumProperty.java b/core/src/main/java/org/geysermc/geyser/entity/properties/type/EnumProperty.java
index 9bc45f560..05e12ba61 100644
--- a/core/src/main/java/org/geysermc/geyser/entity/properties/type/EnumProperty.java
+++ b/core/src/main/java/org/geysermc/geyser/entity/properties/type/EnumProperty.java
@@ -50,7 +50,7 @@ public class EnumProperty implements PropertyType {
     public NbtMap nbtMap() {
         return NbtMap.builder()
                 .putString("name", name)
-                .putList("values", NbtType.STRING, values)
+                .putList("enum", NbtType.STRING, values)
                 .putInt("type", 3)
                 .build();
     }
diff --git a/core/src/main/java/org/geysermc/geyser/entity/type/living/animal/ArmadilloEntity.java b/core/src/main/java/org/geysermc/geyser/entity/type/living/animal/ArmadilloEntity.java
index 205c8cbd9..51fe09383 100644
--- a/core/src/main/java/org/geysermc/geyser/entity/type/living/animal/ArmadilloEntity.java
+++ b/core/src/main/java/org/geysermc/geyser/entity/type/living/animal/ArmadilloEntity.java
@@ -32,23 +32,42 @@ import org.geysermc.mcprotocollib.protocol.data.game.entity.metadata.ArmadilloSt
 import org.geysermc.mcprotocollib.protocol.data.game.entity.metadata.type.ObjectEntityMetadata;
 
 import java.util.UUID;
+import java.util.concurrent.TimeUnit;
 
 public class ArmadilloEntity extends AnimalEntity {
+    private ArmadilloState armadilloState = ArmadilloState.IDLE;
+
     public ArmadilloEntity(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);
     }
 
-    // TODO: This is completely wrong; probably need to store the previous IDLE/ROLLING/SCARED state and check for transitions (pain)
     public void setArmadilloState(ObjectEntityMetadata<ArmadilloState> entityMetadata) {
-        ArmadilloState armadilloState = entityMetadata.getValue();
+        armadilloState = entityMetadata.getValue();
 
         switch (armadilloState) {
             case IDLE -> propertyManager.add("minecraft:armadillo_state", "unrolled");
             case ROLLING -> propertyManager.add("minecraft:armadillo_state", "rolled_up");
-            case SCARED -> propertyManager.add("minecraft:armadillo_state", "rolled_up_peeking");
+            case SCARED -> propertyManager.add("minecraft:armadillo_state", "rolled_up_relaxing");
+            case UNROLLING -> propertyManager.add("minecraft:armadillo_state", "rolled_up_unrolling");
         }
 
         updateBedrockEntityProperties();
     }
+
+    public void onPeeking() {
+        // Technically we should wait if not currently scared
+        if (armadilloState == ArmadilloState.SCARED) {
+            propertyManager.add("minecraft:armadillo_state", "rolled_up_peeking");
+            updateBedrockEntityProperties();
+
+            // Needed for consecutive peeks
+            session.scheduleInEventLoop(() -> {
+                if (armadilloState == ArmadilloState.SCARED) {
+                    propertyManager.add("minecraft:armadillo_state", "rolled_up_relaxing");
+                    updateBedrockEntityProperties();
+                }
+            }, 250, TimeUnit.MILLISECONDS);
+        }
+    }
 }
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 db9874af4..e119d39ce 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
@@ -40,6 +40,7 @@ import org.geysermc.geyser.entity.type.Entity;
 import org.geysermc.geyser.entity.type.EvokerFangsEntity;
 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.session.GeyserSession;
 import org.geysermc.geyser.translator.protocol.PacketTranslator;
@@ -257,6 +258,11 @@ public class JavaEntityEventTranslator extends PacketTranslator<ClientboundEntit
                     wardenEntity.onSonicBoom();
                 }
                 break;
+            case ARMADILLO_PEEKING:
+                if (entity instanceof ArmadilloEntity armadilloEntity) {
+                    armadilloEntity.onPeeking();
+                }
+                break;
         }
 
         if (entityEventPacket.getType() != null) {
diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml
index 3cb136683..22ea061f6 100644
--- a/gradle/libs.versions.toml
+++ b/gradle/libs.versions.toml
@@ -15,7 +15,7 @@ protocol-connection = "3.0.0.Beta1-20240411.165033-128"
 raknet = "1.0.0.CR3-20240416.144209-1"
 blockstateupdater="1.20.80-20240411.142413-1"
 mcauthlib = "d9d773e"
-mcprotocollib = "1ca8808" # Revert from jitpack after release
+mcprotocollib = "1e053f0" # Revert from jitpack after release
 adventure = "4.14.0"
 adventure-platform = "4.3.0"
 junit = "5.9.2"