From 9c191849106369e98debeb9c647f2e37f0217835 Mon Sep 17 00:00:00 2001
From: Sollace <sollacea@gmail.com>
Date: Sat, 6 Mar 2021 14:53:40 +0200
Subject: [PATCH] Keep track of entities summoned and kill then when the spell
 is cancelled

---
 .../ability/magic/spell/NecromancySpell.java  | 51 +++++++++++++++++++
 .../unicopia/entity/EntityReference.java      |  7 +++
 2 files changed, 58 insertions(+)

diff --git a/src/main/java/com/minelittlepony/unicopia/ability/magic/spell/NecromancySpell.java b/src/main/java/com/minelittlepony/unicopia/ability/magic/spell/NecromancySpell.java
index c4038930..d17a0dc6 100644
--- a/src/main/java/com/minelittlepony/unicopia/ability/magic/spell/NecromancySpell.java
+++ b/src/main/java/com/minelittlepony/unicopia/ability/magic/spell/NecromancySpell.java
@@ -1,9 +1,11 @@
 package com.minelittlepony.unicopia.ability.magic.spell;
 
+import java.util.ArrayList;
 import java.util.List;
 
 import com.google.common.collect.Lists;
 import com.minelittlepony.unicopia.ability.magic.Caster;
+import com.minelittlepony.unicopia.entity.EntityReference;
 import com.minelittlepony.unicopia.util.WorldEvent;
 import com.minelittlepony.unicopia.util.shape.Shape;
 import com.minelittlepony.unicopia.util.shape.Sphere;
@@ -11,6 +13,8 @@ import com.minelittlepony.unicopia.util.shape.Sphere;
 import net.minecraft.entity.EntityType;
 import net.minecraft.entity.LivingEntity;
 import net.minecraft.entity.mob.ZombieEntity;
+import net.minecraft.nbt.CompoundTag;
+import net.minecraft.nbt.ListTag;
 import net.minecraft.particle.ParticleTypes;
 import net.minecraft.util.math.BlockPos;
 import net.minecraft.util.math.MathHelper;
@@ -25,6 +29,8 @@ public class NecromancySpell extends AbstractPlacedSpell {
             EntityType.ZOMBIFIED_PIGLIN
     );
 
+    private final List<EntityReference<LivingEntity>> summonedEntities = new ArrayList<>();
+
     protected NecromancySpell(SpellType<?> type) {
         super(type);
     }
@@ -70,9 +76,27 @@ public class NecromancySpell extends AbstractPlacedSpell {
             }
         }
 
+        summonedEntities.removeIf(ref -> !ref.isPresent(source.getWorld()));
+
         return true;
     }
 
+    @Override
+    public void onDestroyed(Caster<?> caster) {
+        super.onDestroyed(caster);
+        LivingEntity master = caster.getMaster();
+        summonedEntities.forEach(ref -> {
+            ref.ifPresent(caster.getWorld(), e -> {
+                if (master != null) {
+                    master.dealDamage(master, e);
+                }
+                if (caster.getWorld().random.nextInt(2000) != 0) {
+                    e.setHealth(0);
+                }
+            });
+        });
+    }
+
     protected void spawnMonster(Caster<?> source, Vec3d pos) {
         int index = (int)MathHelper.nextDouble(source.getWorld().random, 0, spawns.size());
         LivingEntity zombie = spawns.get(index).create(source.getWorld());
@@ -85,5 +109,32 @@ public class NecromancySpell extends AbstractPlacedSpell {
         source.getWorld().syncWorldEvent(WorldEvent.ZOMBIE_BREAK_WOODEN_DOOR, zombie.getBlockPos(), 0);
 
         source.getWorld().spawnEntity(zombie);
+
+        EntityReference<LivingEntity> ref = new EntityReference<>();
+        ref.set(zombie);
+        summonedEntities.add(ref);
+    }
+
+    @Override
+    public void toNBT(CompoundTag compound) {
+        super.toNBT(compound);
+        if (summonedEntities.size() > 0) {
+            ListTag list = new ListTag();
+            summonedEntities.forEach(ref -> list.add(ref.toNBT()));
+            compound.put("summonedEntities", list);
+        }
+    }
+
+    @Override
+    public void fromNBT(CompoundTag compound) {
+        super.fromNBT(compound);
+        if (compound.contains("summonedEntities")) {
+            summonedEntities.clear();
+            compound.getList("summonedEntities", 10).forEach(tag -> {
+                EntityReference<LivingEntity> ref = new EntityReference<>();
+                ref.fromNBT((CompoundTag)tag);
+                summonedEntities.add(ref);
+            });
+        }
     }
 }
diff --git a/src/main/java/com/minelittlepony/unicopia/entity/EntityReference.java b/src/main/java/com/minelittlepony/unicopia/entity/EntityReference.java
index 41b9af56..b92ff655 100644
--- a/src/main/java/com/minelittlepony/unicopia/entity/EntityReference.java
+++ b/src/main/java/com/minelittlepony/unicopia/entity/EntityReference.java
@@ -2,6 +2,7 @@ package com.minelittlepony.unicopia.entity;
 
 import java.util.Optional;
 import java.util.UUID;
+import java.util.function.Consumer;
 
 import javax.annotation.Nullable;
 
@@ -37,6 +38,12 @@ public class EntityReference<T extends Entity> implements NbtSerialisable {
         return entity != null && !entity.removed;
     }
 
+    public void ifPresent(World world, Consumer<T> consumer) {
+        if (isPresent(world)) {
+            consumer.accept(get(world));
+        }
+    }
+
     @SuppressWarnings("unchecked")
     @Nullable
     public T get(World world) {