2 files changed, 31 insertions, 19 deletions

diff --git a/src/main/java/at/hannibal2/skyhanni/data/mob/Mob.kt b/src/main/java/at/hannibal2/skyhanni/data/mob/Mob.kt
index 3d8b00b70..8b9c0509e 100644
--- a/src/main/java/at/hannibal2/skyhanni/data/mob/Mob.kt
+++ b/src/main/java/at/hannibal2/skyhanni/data/mob/Mob.kt
@@ -5,7 +5,7 @@ import at.hannibal2.skyhanni.data.mob.MobFilter.summonOwnerPattern
 import at.hannibal2.skyhanni.events.MobEvent
 import at.hannibal2.skyhanni.mixins.hooks.RenderLivingEntityHelper
 import at.hannibal2.skyhanni.utils.CollectionUtils.toSingletonListOrEmpty
-import at.hannibal2.skyhanni.utils.ColorUtils.withAlpha
+import at.hannibal2.skyhanni.utils.ColorUtils.addAlpha
 import at.hannibal2.skyhanni.utils.EntityUtils.canBeSeen
 import at.hannibal2.skyhanni.utils.EntityUtils.cleanName
 import at.hannibal2.skyhanni.utils.EntityUtils.isCorrupted
@@ -114,20 +114,28 @@ class Mob(
     fun isInvisible() = if (baseEntity !is EntityZombie) baseEntity.isInvisible else false
 
     private var highlightColor: Color? = null
+
+    /** If no alpha is set or alpha is set to 255 it will set the alpha to 127 */
     fun highlight(color: Color) {
-        highlightColor = color
+        highlightColor = color.takeIf { it.alpha == 255 }?.addAlpha(127) ?: color
         internalHighlight()
     }
 
     private fun internalHighlight() {
-        highlightColor?.let {
-            RenderLivingEntityHelper.setEntityColorWithNoHurtTime(baseEntity, it.withAlpha(127)) { true }
+        highlightColor?.let { color ->
+            RenderLivingEntityHelper.setEntityColorWithNoHurtTime(baseEntity, color.rgb) { true }
+            extraEntities.forEach {
+                RenderLivingEntityHelper.setEntityColorWithNoHurtTime(it, color.rgb) { true }
+            }
         }
     }
 
     private fun internalRemoveColor() {
         if (highlightColor == null) return
         RenderLivingEntityHelper.removeCustomRender(baseEntity)
+        extraEntities.forEach {
+            RenderLivingEntityHelper.removeCustomRender(it)
+        }
     }
 
     val boundingBox: AxisAlignedBB
@@ -166,22 +174,22 @@ class Mob(
         fullEntityList().associateWith { this }
 
     internal fun internalAddEntity(entity: EntityLivingBase) {
+        internalRemoveColor()
         if (baseEntity.entityId > entity.entityId) {
             extraEntitiesList.add(0, baseEntity)
-            internalRemoveColor()
             baseEntity = entity
-            internalHighlight()
         } else {
             extraEntitiesList.add(extraEntitiesList.lastIndex + 1, entity)
         }
+        internalHighlight()
         updateBoundingBox()
         MobData.entityToMob[entity] = this
     }
 
     internal fun internalAddEntity(entities: Collection<EntityLivingBase>) {
         val list = entities.drop(1).toMutableList().apply { add(baseEntity) }
-        extraEntitiesList.addAll(0, list)
         internalRemoveColor()
+        extraEntitiesList.addAll(0, list)
         baseEntity = entities.first()
         internalHighlight()
         updateBoundingBox()
@@ -189,19 +197,21 @@ class Mob(
         MobData.entityToMob.putAll(entities.associateWith { this })
     }
 
-    internal fun internalUpdateOfEntity(entity: EntityLivingBase) = when (entity.entityId) {
-        baseEntity.entityId -> {
-            internalRemoveColor()
-            baseEntity = entity
-            internalHighlight()
-        }
-
-        armorStand?.entityId ?: Int.MIN_VALUE -> armorStand = entity as EntityArmorStand
-        else -> {
-            extraEntitiesList.remove(entity)
-            extraEntitiesList.add(entity)
-            Unit // To make return type of this branch Unit
+    internal fun internalUpdateOfEntity(entity: EntityLivingBase) {
+        internalRemoveColor()
+        when (entity.entityId) {
+            baseEntity.entityId -> {
+                baseEntity = entity
+            }
+
+            armorStand?.entityId ?: Int.MIN_VALUE -> armorStand = entity as EntityArmorStand
+            else -> {
+                extraEntitiesList.remove(entity)
+                extraEntitiesList.add(entity)
+                Unit // To make return type of this branch Unit
+            }
         }
+        internalHighlight()
     }
 
     override fun hashCode() = id.hashCode()
diff --git a/src/main/java/at/hannibal2/skyhanni/utils/ColorUtils.kt b/src/main/java/at/hannibal2/skyhanni/utils/ColorUtils.kt
index e5ff47073..b7dbce754 100644
--- a/src/main/java/at/hannibal2/skyhanni/utils/ColorUtils.kt
+++ b/src/main/java/at/hannibal2/skyhanni/utils/ColorUtils.kt
@@ -35,4 +35,6 @@ object ColorUtils {
     }
 
     fun Color.withAlpha(alpha: Int): Int = (alpha.coerceIn(0, 255) shl 24) or (this.rgb and 0x00ffffff)
+
+    fun Color.addAlpha(alpha: Int): Color = Color(red, green, blue, alpha)
 }