From 2ee488f429fba191670c9cb6c39ffe6f80fee526 Mon Sep 17 00:00:00 2001
From: AJ Ferguson <AJ-Ferguson@users.noreply.github.com>
Date: Fri, 19 Apr 2024 05:16:24 -0400
Subject: [PATCH] Fix default wolf and cat collar color

---
 .../living/animal/tameable/CatEntity.java     | 11 +++++-----
 .../living/animal/tameable/WolfEntity.java    | 22 ++++++++++++++-----
 2 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/core/src/main/java/org/geysermc/geyser/entity/type/living/animal/tameable/CatEntity.java b/core/src/main/java/org/geysermc/geyser/entity/type/living/animal/tameable/CatEntity.java
index 412157b5d..93cb92c7d 100644
--- a/core/src/main/java/org/geysermc/geyser/entity/type/living/animal/tameable/CatEntity.java
+++ b/core/src/main/java/org/geysermc/geyser/entity/type/living/animal/tameable/CatEntity.java
@@ -45,7 +45,7 @@ import java.util.UUID;
 
 public class CatEntity extends TameableEntity {
 
-    private byte collarColor;
+    private byte collarColor = 14; // Red - default
 
     public CatEntity(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);
@@ -76,10 +76,7 @@ public class CatEntity extends TameableEntity {
     @Override
     public void setTameableFlags(ByteEntityMetadata entityMetadata) {
         super.setTameableFlags(entityMetadata);
-        // Update collar color if tamed
-        if (getFlag(EntityFlag.TAMED)) {
-            dirtyMetadata.put(EntityDataTypes.COLOR, collarColor);
-        }
+        updateCollarColor();
     }
 
     public void setCatVariant(IntEntityMetadata entityMetadata) {
@@ -101,6 +98,10 @@ public class CatEntity extends TameableEntity {
 
     public void setCollarColor(IntEntityMetadata entityMetadata) {
         collarColor = (byte) entityMetadata.getPrimitiveValue();
+        updateCollarColor();
+    }
+
+    private void updateCollarColor() {
         // Needed or else wild cats are a red color
         if (getFlag(EntityFlag.TAMED)) {
             dirtyMetadata.put(EntityDataTypes.COLOR, collarColor);
diff --git a/core/src/main/java/org/geysermc/geyser/entity/type/living/animal/tameable/WolfEntity.java b/core/src/main/java/org/geysermc/geyser/entity/type/living/animal/tameable/WolfEntity.java
index 0f5b36ec3..c75247fdf 100644
--- a/core/src/main/java/org/geysermc/geyser/entity/type/living/animal/tameable/WolfEntity.java
+++ b/core/src/main/java/org/geysermc/geyser/entity/type/living/animal/tameable/WolfEntity.java
@@ -32,6 +32,7 @@ import org.checkerframework.checker.nullness.qual.NonNull;
 import org.cloudburstmc.math.vector.Vector3f;
 import org.cloudburstmc.protocol.bedrock.data.entity.EntityDataTypes;
 import org.cloudburstmc.protocol.bedrock.data.entity.EntityFlag;
+import org.cloudburstmc.protocol.bedrock.packet.UpdateAttributesPacket;
 import org.geysermc.geyser.entity.EntityDefinition;
 import org.geysermc.geyser.inventory.GeyserItemStack;
 import org.geysermc.geyser.item.Items;
@@ -41,6 +42,7 @@ import org.geysermc.geyser.session.GeyserSession;
 import org.geysermc.geyser.util.InteractionResult;
 import org.geysermc.geyser.util.InteractiveTag;
 
+import java.util.Collections;
 import java.util.Set;
 import java.util.UUID;
 
@@ -54,7 +56,7 @@ public class WolfEntity extends TameableEntity {
             Items.PORKCHOP, Items.BEEF, Items.RABBIT, Items.COOKED_PORKCHOP, Items.COOKED_BEEF, Items.ROTTEN_FLESH, Items.MUTTON, Items.COOKED_MUTTON,
             Items.COOKED_RABBIT);
 
-    private byte collarColor;
+    private byte collarColor = 14; // Red - default
 
     public WolfEntity(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);
@@ -64,19 +66,27 @@ public class WolfEntity extends TameableEntity {
     public void setTameableFlags(ByteEntityMetadata entityMetadata) {
         super.setTameableFlags(entityMetadata);
         // Reset wolf color
-        byte xd = entityMetadata.getPrimitiveValue();
-        boolean angry = (xd & 0x02) == 0x02;
-        if (angry) {
+        if (getFlag(EntityFlag.ANGRY)) {
             dirtyMetadata.put(EntityDataTypes.COLOR, (byte) 0);
+        } else if (getFlag(EntityFlag.TAMED)) {
+            updateCollarColor();
+
+            // This fixes tail angle when taming
+            UpdateAttributesPacket packet = new UpdateAttributesPacket();
+            packet.setRuntimeEntityId(geyserId);
+            packet.setAttributes(Collections.singletonList(createHealthAttribute()));
+            session.sendUpstreamPacket(packet);
         }
     }
 
     public void setCollarColor(IntEntityMetadata entityMetadata) {
         collarColor = (byte) entityMetadata.getPrimitiveValue();
-        if (getFlag(EntityFlag.ANGRY)) {
-            return;
+        if (!getFlag(EntityFlag.ANGRY) && getFlag(EntityFlag.TAMED)) {
+            updateCollarColor();
         }
+    }
 
+    private void updateCollarColor() {
         dirtyMetadata.put(EntityDataTypes.COLOR, collarColor);
         if (ownerBedrockId == 0) {
             // If a color is set and there is no owner entity ID, set one.