Redirect lightning into a jar when the player is struck

src/main/java/com/minelittlepony/unicopia/entity/Living.java

@@ -1,18 +1,24 @@
 package com.minelittlepony.unicopia.entity;
 
+import java.util.Optional;
+import java.util.stream.Stream;
+
 import javax.annotation.Nullable;
 
 import com.minelittlepony.unicopia.ability.magic.Attached;
 import com.minelittlepony.unicopia.ability.magic.Caster;
 import com.minelittlepony.unicopia.ability.magic.Spell;
 import com.minelittlepony.unicopia.ability.magic.spell.DisguiseSpell;
+import com.minelittlepony.unicopia.item.UItems;
 import com.minelittlepony.unicopia.network.EffectSync;
 
 import net.minecraft.entity.LivingEntity;
 import net.minecraft.entity.damage.DamageSource;
 import net.minecraft.entity.data.TrackedData;
 import net.minecraft.entity.projectile.ProjectileEntity;
+import net.minecraft.item.ItemStack;
 import net.minecraft.nbt.CompoundTag;
+import net.minecraft.util.Hand;
 
 public abstract class Living<T extends LivingEntity> implements Equine<T>, Caster<T> {
 
@@ -26,6 +32,8 @@ public abstract class Living<T extends LivingEntity> implements Equine<T>, Caste
     @Nullable
     private Runnable landEvent;
 
+    private int invinsibilityTicks;
+
     protected Living(T entity, TrackedData<CompoundTag> effect) {
         this.entity = entity;
         this.effectDelegate = new EffectSync(this, effect);
@@ -75,6 +83,10 @@ public abstract class Living<T extends LivingEntity> implements Equine<T>, Caste
             }
         }
 
+        if (invinsibilityTicks > 0) {
+            invinsibilityTicks--;
+        }
+
         prevSneaking = entity.isSneaking();
         prevLanded = entity.isOnGround();
 
@@ -90,6 +102,35 @@ public abstract class Living<T extends LivingEntity> implements Equine<T>, Caste
         }
     }
 
+    @Override
+    public Optional<Boolean> onDamage(DamageSource source, float amount) {
+
+        if (source == DamageSource.LIGHTNING_BOLT) {
+            if (invinsibilityTicks > 0 || tryCaptureLightning()) {
+                return Optional.of(false);
+            }
+        }
+
+        return Optional.empty();
+    }
+
+    private boolean tryCaptureLightning() {
+        return getInventoryStacks().filter(stack -> !stack.isEmpty() && stack.getItem() == UItems.EMPTY_JAR).findFirst().map(stack -> {
+            invinsibilityTicks = 20;
+            stack.split(1);
+            giveBackItem(UItems.LIGHTNING_JAR.getDefaultStack());
+            return stack;
+        }).isPresent();
+    }
+
+    protected Stream<ItemStack> getInventoryStacks() {
+        return Stream.of(entity.getStackInHand(Hand.MAIN_HAND), entity.getStackInHand(Hand.OFF_HAND));
+    }
+
+    protected void giveBackItem(ItemStack stack) {
+        entity.dropStack(stack);
+    }
+
     @Override
     public boolean onProjectileImpact(ProjectileEntity projectile) {
         if (hasSpell()) {

src/main/java/com/minelittlepony/unicopia/entity/player/Pony.java

@@ -3,6 +3,7 @@ package com.minelittlepony.unicopia.entity.player;
 import java.util.List;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.stream.Stream;
 
 import javax.annotation.Nullable;
 
@@ -30,6 +31,7 @@ import com.minelittlepony.unicopia.util.Copieable;
 import com.minelittlepony.unicopia.util.MagicalDamageSource;
 import com.minelittlepony.common.util.animation.LinearInterpolator;
 import com.google.common.collect.Lists;
+import com.google.common.collect.Streams;
 import com.minelittlepony.common.util.animation.Interpolator;
 import com.mojang.authlib.GameProfile;
 
@@ -344,6 +346,21 @@ public class Pony extends Living<PlayerEntity> implements Transmittable, Copieab
         return foodSubtract - lostLevels;
     }
 
+    @Override
+    protected Stream<ItemStack> getInventoryStacks() {
+        return Streams.concat(
+                super.getInventoryStacks(),
+                entity.inventory.main.stream()
+        );
+    }
+
+    @Override
+    protected void giveBackItem(ItemStack stack) {
+        if (!entity.giveItemStack(stack)) {
+            entity.dropItem(stack, false);
+        }
+    }
+
     public Optional<Text> trySleep(BlockPos pos) {
         return findAllSpellsInRange(10)
                 .filter(p -> p instanceof Pony && ((Pony)p).isEnemy(this))

src/main/java/com/minelittlepony/unicopia/mixin/MixinLivingEntity.java

@@ -22,6 +22,7 @@ import com.minelittlepony.unicopia.entity.ItemWielder;
 
 import net.minecraft.entity.Entity;
 import net.minecraft.entity.LivingEntity;
+import net.minecraft.entity.damage.DamageSource;
 import net.minecraft.item.ItemStack;
 import net.minecraft.nbt.CompoundTag;
 import net.minecraft.util.Hand;
@@ -111,6 +112,11 @@ abstract class MixinLivingEntity extends Entity implements PonyContainer<Equine<
         Creature.boostrap();
     }
 
+    @Inject(method = "damage(Lnet/minecraft/entity/damage/DamageSource;F)Z", at = @At("HEAD"), cancellable = true)
+    private void onDamage(DamageSource source, float amount, CallbackInfoReturnable<Boolean> info) {
+        get().onDamage(source, amount).ifPresent(info::setReturnValue);
+    }
+
     @Inject(method = "writeCustomDataToTag(Lnet/minecraft/nbt/CompoundTag;)V", at = @At("HEAD"))
     private void onWriteCustomDataToTag(CompoundTag tag, CallbackInfo info) {
         tag.put("unicopia_caster", get().toNBT());