From 1cea2bf3917429518a465a7398006aecc5495fa1 Mon Sep 17 00:00:00 2001
From: Sollace <sollacea@gmail.com>
Date: Sun, 11 Sep 2022 16:47:06 +0200
Subject: [PATCH] Exclude spellbooks from shields' effects and exclude/include
 classes of entities for shields based on certain traits added to the spell

---
 .../magic/spell/effect/ShieldSpell.java       | 23 +++++++++++++++----
 .../magic/spell/effect/TargetSelecter.java    |  2 ++
 2 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/src/main/java/com/minelittlepony/unicopia/ability/magic/spell/effect/ShieldSpell.java b/src/main/java/com/minelittlepony/unicopia/ability/magic/spell/effect/ShieldSpell.java
index 866c55ba..3cdbe28f 100644
--- a/src/main/java/com/minelittlepony/unicopia/ability/magic/spell/effect/ShieldSpell.java
+++ b/src/main/java/com/minelittlepony/unicopia/ability/magic/spell/effect/ShieldSpell.java
@@ -21,6 +21,9 @@ import net.minecraft.entity.LivingEntity;
 import net.minecraft.entity.TntEntity;
 import net.minecraft.entity.Entity.RemovalReason;
 import net.minecraft.entity.decoration.ArmorStandEntity;
+import net.minecraft.entity.mob.HostileEntity;
+import net.minecraft.entity.passive.PassiveEntity;
+import net.minecraft.entity.player.PlayerEntity;
 import net.minecraft.entity.vehicle.AbstractMinecartEntity;
 import net.minecraft.entity.vehicle.BoatEntity;
 import net.minecraft.util.math.Vec3d;
@@ -111,14 +114,26 @@ public class ShieldSpell extends AbstractSpell {
     }
 
     protected boolean isValidTarget(Caster<?> source, Entity entity) {
-        return (entity instanceof LivingEntity
+        boolean valid = (entity instanceof LivingEntity
                 || entity instanceof TntEntity
                 || entity instanceof FallingBlockEntity
-                || entity instanceof EyeOfEnderEntity
-                || entity instanceof BoatEntity
                 || ProjectileUtil.isFlyingProjectile(entity)
                 || entity instanceof AbstractMinecartEntity)
-            && !(entity instanceof ArmorStandEntity);
+            && !(  entity instanceof ArmorStandEntity
+                || entity instanceof EyeOfEnderEntity
+                || entity instanceof BoatEntity
+        );
+
+        if (getTraits().get(Trait.LIFE) > 0) {
+            valid &= !(entity instanceof PassiveEntity);
+        }
+        if (getTraits().get(Trait.BLOOD) > 0) {
+            valid &= entity instanceof HostileEntity;
+        }
+        if (getTraits().get(Trait.ICE) > 0) {
+            valid &= entity instanceof PlayerEntity;
+        }
+        return valid;
     }
 
     protected long applyEntities(Caster<?> source) {
diff --git a/src/main/java/com/minelittlepony/unicopia/ability/magic/spell/effect/TargetSelecter.java b/src/main/java/com/minelittlepony/unicopia/ability/magic/spell/effect/TargetSelecter.java
index 8936bb9e..279599f2 100644
--- a/src/main/java/com/minelittlepony/unicopia/ability/magic/spell/effect/TargetSelecter.java
+++ b/src/main/java/com/minelittlepony/unicopia/ability/magic/spell/effect/TargetSelecter.java
@@ -12,6 +12,7 @@ import com.minelittlepony.unicopia.ability.magic.Affine;
 import com.minelittlepony.unicopia.ability.magic.Caster;
 import com.minelittlepony.unicopia.ability.magic.SpellPredicate;
 import com.minelittlepony.unicopia.ability.magic.spell.Spell;
+import com.minelittlepony.unicopia.entity.SpellbookEntity;
 import com.minelittlepony.unicopia.entity.player.Pony;
 import com.minelittlepony.unicopia.item.FriendshipBraceletItem;
 import net.minecraft.entity.Entity;
@@ -33,6 +34,7 @@ public class TargetSelecter {
 
         return source.findAllEntitiesInRange(radius)
             .filter(entity -> entity.isAlive() && !entity.isRemoved() && !ownerCheck.test(entity) && !SpellPredicate.IS_SHIELD_LIKE.isOn(entity))
+            .filter(entity -> !(entity instanceof SpellbookEntity))
             .filter(e -> filter.test(source, e))
             .map(i -> {
                 targets.computeIfAbsent(i.getUuid(), Target::new);