From 9a19ed36c1ef9e3c2a42730800d6dec110169de9 Mon Sep 17 00:00:00 2001
From: Sollace <sollacea@gmail.com>
Date: Thu, 13 Oct 2022 21:18:56 +0200
Subject: [PATCH] Include max health in what's swapped by the mind swap spell

---
 .../unicopia/entity/behaviour/EntitySwap.java      | 14 +++++++++++++-
 .../unicopia/entity/player/PlayerAttributes.java   |  6 ++++++
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/src/main/java/com/minelittlepony/unicopia/entity/behaviour/EntitySwap.java b/src/main/java/com/minelittlepony/unicopia/entity/behaviour/EntitySwap.java
index c16f825e..bc411313 100644
--- a/src/main/java/com/minelittlepony/unicopia/entity/behaviour/EntitySwap.java
+++ b/src/main/java/com/minelittlepony/unicopia/entity/behaviour/EntitySwap.java
@@ -3,10 +3,12 @@ package com.minelittlepony.unicopia.entity.behaviour;
 import java.util.*;
 
 import com.minelittlepony.unicopia.entity.Living;
+import com.minelittlepony.unicopia.entity.player.PlayerAttributes;
 import com.minelittlepony.unicopia.util.Swap;
 
 import net.minecraft.entity.Entity;
 import net.minecraft.entity.LivingEntity;
+import net.minecraft.entity.attribute.EntityAttributes;
 import net.minecraft.entity.mob.PathAwareEntity;
 import net.minecraft.entity.player.PlayerEntity;
 import net.minecraft.network.packet.s2c.play.EntitySetHeadYawS2CPacket;
@@ -46,12 +48,22 @@ public interface EntitySwap {
             effects.forEach(entity::addStatusEffect);
         }
     );
+    Swap<LivingEntity> MAX_HEALTH = Swap.of(LivingEntity::getMaxHealth, (e, newMax) -> {
+        float oldHealthPercentage = e.getHealth() / e.getMaxHealth();
+        e.getAttributeInstance(EntityAttributes.GENERIC_MAX_HEALTH).removeModifier(PlayerAttributes.HEALTH_SWAPPING_MODIFIER_ID);
+
+        float change = newMax - e.getMaxHealth();
+        if (!MathHelper.approximatelyEquals(change, 0)) {
+            e.getAttributeInstance(EntityAttributes.GENERIC_MAX_HEALTH).addPersistentModifier(PlayerAttributes.healthChange(change));
+        }
+        e.setHealth(oldHealthPercentage * newMax);
+    });
     Swap<LivingEntity> HEALTH = Swap.of(LivingEntity::getHealth, LivingEntity::getMaxHealth, LivingEntity::setHealth, Number::floatValue);
     Swap<LivingEntity> AIR = Swap.of(LivingEntity::getAir, LivingEntity::getMaxAir, LivingEntity::setAir, Number::intValue);
 
     List<Swap<Entity>> REGISTRY = new ArrayList<>(List.of(
             Swap.union(POSITION, VELOCITY, PITCH, YAW, HEAD_YAW, BODY_YAW, FIRE_TICKS, PASSENGERS),
-            Swap.union(STATUS_EFFECTS, HEALTH, AIR).upcast(e -> e instanceof LivingEntity)
+            Swap.union(STATUS_EFFECTS, MAX_HEALTH, HEALTH, AIR).upcast(e -> e instanceof LivingEntity)
     ));
     Swap<Entity> ALL = Swap.union(REGISTRY);
 }
diff --git a/src/main/java/com/minelittlepony/unicopia/entity/player/PlayerAttributes.java b/src/main/java/com/minelittlepony/unicopia/entity/player/PlayerAttributes.java
index 172ac730..896a0906 100644
--- a/src/main/java/com/minelittlepony/unicopia/entity/player/PlayerAttributes.java
+++ b/src/main/java/com/minelittlepony/unicopia/entity/player/PlayerAttributes.java
@@ -29,6 +29,12 @@ public class PlayerAttributes implements Tickable {
     public static final EntityAttributeModifier BAT_HANGING =
             new EntityAttributeModifier(UUID.fromString("a54f2595-521e-480b-b9d5-6e750577a564"), "Bat Pony Hanging", -2, Operation.MULTIPLY_TOTAL);
 
+    public static final UUID HEALTH_SWAPPING_MODIFIER_ID = UUID.fromString("7b93803e-4b25-11ed-951e-00155d43e0a2");
+
+    public static EntityAttributeModifier healthChange(float addition) {
+        return new EntityAttributeModifier(HEALTH_SWAPPING_MODIFIER_ID, "Health Swap", addition, Operation.ADDITION);
+    }
+
     private final Pony pony;
 
     public PlayerAttributes(Pony pony) {