2 files changed, 23 insertions, 23 deletions

diff --git a/src/main/java/at/hannibal2/skyhanni/mixins/hooks/RenderLivingEntityHelper.kt b/src/main/java/at/hannibal2/skyhanni/mixins/hooks/RenderLivingEntityHelper.kt
index 4d54aa90b..eb7456b4c 100644
--- a/src/main/java/at/hannibal2/skyhanni/mixins/hooks/RenderLivingEntityHelper.kt
+++ b/src/main/java/at/hannibal2/skyhanni/mixins/hooks/RenderLivingEntityHelper.kt
@@ -1,8 +1,6 @@
 package at.hannibal2.skyhanni.mixins.hooks
 
 import at.hannibal2.skyhanni.events.LorenzWorldChangeEvent
-import at.hannibal2.skyhanni.events.RenderMobColoredEvent
-import at.hannibal2.skyhanni.events.ResetEntityHurtEvent
 import at.hannibal2.skyhanni.test.SkyHanniDebugsAndTests
 import net.minecraft.entity.EntityLivingBase
 import net.minecraftforge.fml.common.eventhandler.SubscribeEvent
@@ -14,7 +12,6 @@ class RenderLivingEntityHelper {
         entityColorMap.clear()
         entityColorCondition.clear()
 
-        entityNoHurTime.clear()
         entityNoHurTimeCondition.clear()
     }
 
@@ -23,7 +20,6 @@ class RenderLivingEntityHelper {
         private val entityColorMap = mutableMapOf<EntityLivingBase, Int>()
         private val entityColorCondition = mutableMapOf<EntityLivingBase, () -> Boolean>()
 
-        private val entityNoHurTime = mutableListOf<EntityLivingBase>()
         private val entityNoHurTimeCondition = mutableMapOf<EntityLivingBase, () -> Boolean>()
 
         fun <T : EntityLivingBase> removeEntityColor(entity: T) {
@@ -37,11 +33,24 @@ class RenderLivingEntityHelper {
         }
 
         fun <T : EntityLivingBase> setNoHurtTime(entity: T, condition: () -> Boolean) {
-            entityNoHurTime.add(entity)
             entityNoHurTimeCondition[entity] = condition
         }
 
-        fun <T : EntityLivingBase> setColorMultiplier(entity: T): Int {
+        fun <T : EntityLivingBase> setEntityColorWithNoHurtTime(entity: T, color: Int, condition: () -> Boolean) {
+            setEntityColor(entity, color, condition)
+            setNoHurtTime(entity, condition)
+        }
+
+        fun <T : EntityLivingBase> removeNoHurtTime(entity: T) {
+            entityNoHurTimeCondition.remove(entity)
+        }
+
+        fun <T : EntityLivingBase> removeCustomRender(entity: T) {
+            removeEntityColor(entity)
+            removeNoHurtTime(entity)
+        }
+
+        fun <T : EntityLivingBase> internalSetColorMultiplier(entity: T): Int {
             if (!SkyHanniDebugsAndTests.globalRender) return 0
             if (entityColorMap.containsKey(entity)) {
                 val condition = entityColorCondition[entity]!!
@@ -49,27 +58,18 @@ class RenderLivingEntityHelper {
                     return entityColorMap[entity]!!
                 }
             }
-
-            // TODO remove event
-            if (!SkyHanniDebugsAndTests.globalRender) return 0
-            val event = RenderMobColoredEvent(entity, 0)
-            event.postAndCatch()
-            return event.color
+            return 0
         }
 
-        fun <T : EntityLivingBase> changeHurtTime(entity: T): Int {
-            if (!SkyHanniDebugsAndTests.globalRender) return 0
-            if (entityNoHurTime.contains(entity)) {
-                val condition = entityNoHurTimeCondition[entity]!!
+        fun <T : EntityLivingBase> internalChangeHurtTime(entity: T): Int {
+            if (!SkyHanniDebugsAndTests.globalRender) return entity.hurtTime
+            run {
+                val condition = entityNoHurTimeCondition[entity] ?: return@run
                 if (condition.invoke()) {
                     return 0
                 }
             }
-
-            // TODO remove event
-            val event = ResetEntityHurtEvent(entity, false)
-            event.postAndCatch()
-            return if (event.shouldReset) 0 else entity.hurtTime
+            return entity.hurtTime
         }
     }
 }
diff --git a/src/main/java/at/hannibal2/skyhanni/mixins/transformers/renderer/MixinRendererLivingEntity.java b/src/main/java/at/hannibal2/skyhanni/mixins/transformers/renderer/MixinRendererLivingEntity.java
index b74905564..eb1b5afda 100644
--- a/src/main/java/at/hannibal2/skyhanni/mixins/transformers/renderer/MixinRendererLivingEntity.java
+++ b/src/main/java/at/hannibal2/skyhanni/mixins/transformers/renderer/MixinRendererLivingEntity.java
@@ -21,11 +21,11 @@ public abstract class MixinRendererLivingEntity<T extends EntityLivingBase> exte
 
     @Inject(method = "getColorMultiplier", at = @At("HEAD"), cancellable = true)
     private void setColorMultiplier(T entity, float lightBrightness, float partialTickTime, CallbackInfoReturnable<Integer> cir) {
-        cir.setReturnValue(RenderLivingEntityHelper.Companion.setColorMultiplier(entity));
+        cir.setReturnValue(RenderLivingEntityHelper.Companion.internalSetColorMultiplier(entity));
     }
 
     @Redirect(method = "setBrightness", at = @At(value = "FIELD", target = "Lnet/minecraft/entity/EntityLivingBase;hurtTime:I", opcode = Opcodes.GETFIELD))
     private int changeHurtTime(EntityLivingBase entity) {
-        return RenderLivingEntityHelper.Companion.changeHurtTime(entity);
+        return RenderLivingEntityHelper.Companion.internalChangeHurtTime(entity);
     }
 }