minecraft-kajetans-mod/src/main/java/me/km/snuviscript/commands/EntityCommands.java

package me.km.snuviscript.commands;

import me.hammerle.snuviscript.code.ScriptManager;
import me.km.entities.EntityItemProjectile;
import me.km.scheduler.SnuviScheduler;
import static me.km.snuviscript.commands.CommandUtils.getNamedClass;
import me.km.utils.Location;
import me.km.utils.Mapper;
import me.km.utils.Utils;
import net.minecraft.enchantment.EnchantmentHelper;
import net.minecraft.entity.*;
import net.minecraft.entity.ai.attributes.Attributes;
import net.minecraft.entity.effect.LightningBoltEntity;
import net.minecraft.entity.item.*;
import net.minecraft.entity.monster.CreeperEntity;
import net.minecraft.entity.passive.SheepEntity;
import net.minecraft.entity.player.ServerPlayerEntity;
import net.minecraft.entity.projectile.*;
import net.minecraft.inventory.EquipmentSlotType;
import net.minecraft.item.ItemStack;
import net.minecraft.nbt.CompoundNBT;
import net.minecraft.nbt.JsonToNBT;
import net.minecraft.potion.Effect;
import net.minecraft.potion.EffectInstance;
import net.minecraft.potion.PotionUtils;
import net.minecraft.util.DamageSource;
import net.minecraft.util.Direction;
import net.minecraft.util.ResourceLocation;
import net.minecraft.util.math.BlockPos;
import net.minecraft.util.math.vector.Vector3d;
import net.minecraft.util.text.StringTextComponent;
import net.minecraft.world.World;
import net.minecraft.world.server.ServerWorld;

public class EntityCommands {
    public static void registerFunctions(ScriptManager sm, SnuviScheduler scheduler) {
        sm.registerConsumer("entity.setnopickup", (sc, in) -> {
            ((AbstractArrowEntity) in[0].get(sc)).pickupStatus = AbstractArrowEntity.PickupStatus.DISALLOWED;
        });
        sm.registerFunction("entity.shootprojectile", (sc, in) -> launchProjectile((LivingEntity) in[0].get(sc),
                getNamedClass(in[1].getString(sc)), in[2].getDouble(sc), in.length >= 4 ? in[3].get(sc) : null));
        sm.registerFunction("entity.isblocking", (sc, in) -> ((LivingEntity) in[0].get(sc)).isActiveItemStackBlocking());
        sm.registerFunction("entity.getarmorthoughness", (sc, in) -> ((LivingEntity) in[0].get(sc)).getAttribute(Attributes.ARMOR_TOUGHNESS).getValue());
        sm.registerFunction("entity.getarmor", (sc, in) -> (double) ((LivingEntity) in[0].get(sc)).getTotalArmorValue());
        sm.registerFunction("entity.getenchantmentmodifier", (sc, in) -> EnchantmentHelper.getEnchantmentModifierDamage(((LivingEntity) in[0].get(sc)).getArmorInventoryList(), (DamageSource) in[1].get(sc)));
        sm.registerConsumer("entity.setburning", (sc, in) -> ((Entity) in[0].get(sc)).setFire(in[1].getInt(sc)));
        sm.registerFunction("entity.isburning", (sc, in) -> ((Entity) in[0].get(sc)).isBurning());
        sm.registerFunction("entity.getlook", (sc, in) -> {
            Object[] o = new Object[3];
            Vector3d v = ((Entity) in[0].get(sc)).getLookVec();
            o[0] = v.x;
            o[1] = v.y;
            o[2] = v.z;
            return o;
        });
        sm.registerFunction("entity.getlocation", (sc, in) -> new Location((Entity) in[0].get(sc)));
        sm.registerConsumer("entity.heal", (sc, in) -> {
            LivingEntity liv = (LivingEntity) in[0].get(sc);
            float heal = in[1].getFloat(sc);
            scheduler.scheduleTask(() -> liv.heal(heal));
        });
        sm.registerConsumer("entity.damage", (sc, in) -> {
            LivingEntity liv = (LivingEntity) in[0].get(sc);
            float damage = in[1].getFloat(sc);
            DamageSource damageSource = (in.length >= 3) ? (DamageSource) in[2].get(sc) : DamageSource.GENERIC;
            scheduler.scheduleTask(() -> liv.attackEntityFrom(damageSource, damage));
        });
        sm.registerFunction("entity.fromsource", (sc, in) -> {
            DamageSource ds = (DamageSource) in[0].get(sc);
            Entity ent = ds.getTrueSource();
            if(ent == null) {
                return ds.getImmediateSource();
            }
            return ent;
        });
        sm.registerAlias("damage.get", "entity.getdamagesource");
        sm.registerFunction("entity.getmaxhealth", (sc, in) -> (double) ((LivingEntity) in[0].get(sc)).getMaxHealth());
        sm.registerFunction("entity.gethealth", (sc, in) -> (double) ((LivingEntity) in[0].get(sc)).getHealth());
        sm.registerConsumer("entity.sethealth", (sc, in) -> ((LivingEntity) in[0].get(sc)).setHealth(in[1].getFloat(sc)));
        sm.registerConsumer("entity.setname", (sc, in) -> {
            Entity ent = (Entity) in[0].get(sc);
            ent.setCustomName(new StringTextComponent(in[1].getString(sc)));
            if(in.length >= 3) {
                ent.setCustomNameVisible(in[2].getBoolean(sc));
                return;
            }
            ent.setCustomNameVisible(false);
        });
        sm.registerFunction("entity.getname", (sc, in) -> ((Entity) in[0].get(sc)).getDisplayName().getString());
        sm.registerConsumer("entity.throw", (sc, in) -> {
            Entity ent = (Entity) in[0].get(sc);
            ent.setMotion(in[1].getDouble(sc), in[2].getDouble(sc), in[3].getDouble(sc));
            ent.velocityChanged = true;
        });
        sm.registerConsumer("entity.teleport", (sc, in) -> {
            Entity ent = (Entity) in[0].get(sc);
            Location l = (Location) in[1].get(sc);
            if(l.getWorld() == null) {
                throw new IllegalArgumentException("world must not be null");
            }
            if(ent instanceof ServerPlayerEntity) {
                ServerPlayerEntity p = (ServerPlayerEntity) ent;

                p.stopRiding();
                if(p.isSleeping()) {
                    p.stopSleepInBed(true, true);
                }

                float yaw = l.getYaw() != 0.0f ? l.getYaw() : ent.rotationYaw;
                float pitch = l.getPitch() != 0.0f ? l.getPitch() : ent.rotationPitch;
                p.teleport((ServerWorld) l.getWorld(), l.getX(), l.getY(), l.getZ(), yaw, pitch);
            } else {
                if(ent.world != l.getWorld()) {
                    ServerWorld ws = (ServerWorld) l.getWorld();
                    ent.changeDimension(ws);
                }
                float yaw = l.getYaw() != 0.0f ? l.getYaw() : ent.rotationYaw;
                float pitch = l.getPitch() != 0.0f ? l.getPitch() : ent.rotationPitch;
                ent.setLocationAndAngles(l.getX(), l.getY(), l.getZ(), yaw, pitch);
            }
        });
        sm.registerConsumer("entity.setequip", (sc, in) -> {
            LivingEntity liv = (LivingEntity) in[0].get(sc);
            ItemStack stack = ((ItemStack) in[2].get(sc)).copy();
            switch(in[1].getString(sc)) {
                case "hand":
                    liv.setItemStackToSlot(EquipmentSlotType.MAINHAND, stack);
                    return;
                case "head":
                    liv.setItemStackToSlot(EquipmentSlotType.HEAD, stack);
                    return;
                case "chest":
                    liv.setItemStackToSlot(EquipmentSlotType.CHEST, stack);
                    return;
                case "legs":
                    liv.setItemStackToSlot(EquipmentSlotType.LEGS, stack);
                    return;
                case "feet":
                    liv.setItemStackToSlot(EquipmentSlotType.FEET, stack);
                    return;
                case "offhand":
                    liv.setItemStackToSlot(EquipmentSlotType.OFFHAND, stack);
            }
        });
        sm.registerFunction("entity.getequip", (sc, in) -> {
            LivingEntity liv = (LivingEntity) in[0].get(sc);
            switch(in[1].getString(sc)) {
                case "hand":
                    return liv.getItemStackFromSlot(EquipmentSlotType.MAINHAND);
                case "head":
                    return liv.getItemStackFromSlot(EquipmentSlotType.HEAD);
                case "chest":
                    return liv.getItemStackFromSlot(EquipmentSlotType.CHEST);
                case "legs":
                    return liv.getItemStackFromSlot(EquipmentSlotType.LEGS);
                case "feet":
                    return liv.getItemStackFromSlot(EquipmentSlotType.FEET);
                case "offhand":
                    return liv.getItemStackFromSlot(EquipmentSlotType.OFFHAND);
            }
            return ItemStack.EMPTY;
        });
        sm.registerConsumer("entity.removeall", (sc, in) -> {
            Class<? extends Entity> c = (Class<? extends Entity>) getNamedClass(in[0].getString(sc));
            if(c == Entity.class) {
                return;
            }
            Location l = (Location) in[1].get(sc);
            Utils.getEntities(l.getWorld(), l.getX(), l.getY(), l.getZ(), in[2].getDouble(sc), c).stream().forEach(ent -> {
                ent.remove();
            });
        });
        sm.registerConsumer("entity.remove", (sc, in) -> ((Entity) in[0].get(sc)).remove());
        sm.registerConsumer("entity.setinvulnerable", (sc, in) -> {
            ((Entity) in[0].get(sc)).setInvulnerable(in[1].getBoolean(sc));
        });
        sm.registerConsumer("entity.setsilent", (sc, in) -> {
            ((Entity) in[0].get(sc)).setSilent(in[1].getBoolean(sc));
        });
        sm.registerConsumer("entity.setinvisible", (sc, in) -> {
            ((Entity) in[0].get(sc)).setInvisible(in[1].getBoolean(sc));
        });
        sm.registerConsumer("entity.mount", (sc, in) -> {
            ((Entity) in[0].get(sc)).startRiding(((Entity) in[1].get(sc)));
        });
        sm.registerConsumer("entity.unmount", (sc, in) -> {
            ((Entity) in[0].get(sc)).stopRiding();
        });
        sm.registerConsumer("entity.addeffect", (sc, in) -> {
            LivingEntity base = (LivingEntity) in[0].get(sc);
            Effect potion = Mapper.getPotion(in[1].getString(sc));
            if(potion == null) { // doing this only to prevent EffectInstance doing shit
                throw new IllegalArgumentException("potion does not exist");
            }
            if(base.isPotionActive(potion)) {
                base.removePotionEffect(potion);
            }
            boolean showParticles = in.length >= 5 ? in[4].getBoolean(sc) : true;
            base.addPotionEffect(new EffectInstance(potion, in[2].getInt(sc), in[3].getInt(sc), false, showParticles));
        });
        sm.registerConsumer("entity.cleareffects", (sc, in) -> {
            ((LivingEntity) in[0].get(sc)).clearActivePotions();
        });
        sm.registerFunction("entity.geteffectamplifier", (sc, in) -> {
            EffectInstance effect = ((LivingEntity) in[0].get(sc)).getActivePotionEffect(Mapper.getPotion(in[1].getString(sc)));
            return effect == null ? 0 : effect.getAmplifier() + 1;
        });
        sm.registerConsumer("entity.spawnitemframe", (sc, in) -> {
            Location l = ((Location) in[0].get(sc));
            ItemFrameEntity frame = new ItemFrameEntity(l.getWorld(), l.getBlockPos(), Direction.byName(in[1].getString(sc)));
            frame.setDisplayedItem(((ItemStack) in[2].get(sc))); // copy happens in internals
            l.getWorld().addEntity(frame);
        });
        sm.registerFunction("entity.getitemfromframe", (sc, in) -> ((ItemFrameEntity) in[0].get(sc)).getDisplayedItem());
        sm.registerFunction("entity.get", (sc, in) -> {
            Location l = (Location) in[0].get(sc);
            return Utils.getEntity(l.getWorld(), l.getX(), l.getY(), l.getZ(), in[1].getDouble(sc), getNamedClass(in[2].getString(sc)));
        });
        sm.registerFunction("entity.getpotiontype", (sc, in) -> PotionUtils.getPotionFromItem(((PotionEntity) in[0].get(sc)).getItem()).getRegistryName().toString());
        sm.registerConsumer("entity.setgravity", (sc, in) -> {
            ((Entity) in[0].get(sc)).setNoGravity(!in[1].getBoolean(sc));
        });
        sm.registerFunction("entity.iswet", (sc, in) -> ((Entity) in[0].get(sc)).isWet());
        sm.registerConsumer("entity.setpickupdelay", (sc, in) -> {
            ((ItemEntity) in[0].get(sc)).setPickupDelay(in[1].getInt(sc));
        });
        sm.registerFunction("entity.spawn", (sc, in) -> {
            ResourceLocation rl = new ResourceLocation(in[0].getString(sc));
            Location l = (Location) in[1].get(sc);
            ServerWorld sw = (ServerWorld) l.getWorld();

            CompoundNBT compoundnbt = in.length >= 3 ? JsonToNBT.getTagFromJson(in[2].getString(sc)) : new CompoundNBT();
            compoundnbt.putString("id", rl.toString());
            if(EntityType.getKey(EntityType.LIGHTNING_BOLT).equals(rl)) {
                LightningBoltEntity ent = EntityType.LIGHTNING_BOLT.create(sw);
                ent.moveForced(l.getPos());
                sw.addEntity(ent);
                return ent;
            }

            Entity entity = EntityType.loadEntityAndExecute(compoundnbt, sw, (ent) -> {
                ent.setLocationAndAngles(l.getX(), l.getY(), l.getZ(), ent.rotationYaw, ent.rotationPitch);
                return sw.summonEntity(ent) ? ent : null;
            });
            if(entity != null && entity instanceof MobEntity) {
                ((MobEntity) entity).onInitialSpawn(sw, sw.getDifficultyForLocation(entity.getPosition()), SpawnReason.COMMAND, null, null);
            }
            return entity;
        });
        sm.registerFunction("entity.near", (sc, in) -> {
            Object o = in[0].get(sc);
            if(o instanceof Location) {
                return Utils.getEntities((Location) o, in[1].getDouble(sc));
            }
            return Utils.getEntities((Entity) o, in[1].getDouble(sc));
        });
        sm.registerConsumer("entity.setspeed", (sc, in) -> {
            ((LivingEntity) in[0].get(sc)).getAttribute(Attributes.MOVEMENT_SPEED).setBaseValue(in[1].getDouble(sc));
        });
        sm.registerConsumer("entity.setgrowingage", (sc, in) -> {
            ((AgeableEntity) in[0].get(sc)).setGrowingAge(in[1].getInt(sc));
        });
        sm.registerFunction("entity.gettype", (sc, in) -> ((Entity) in[0].get(sc)).getType().getRegistryName().getPath());
        sm.registerFunction("entity.issneaking", (sc, in) -> ((Entity) in[0].get(sc)).isCrouching());
        sm.registerFunction("entity.issneaking", (sc, in) -> ((Entity) in[0].get(sc)).isCrouching());
        sm.registerFunction("sheep.issheared", (sc, in) -> ((SheepEntity) in[0].get(sc)).getSheared());
        sm.registerFunction("sheep.getcolor", (sc, in) -> ((SheepEntity) in[0].get(sc)).getFleeceColor().toString());
        sm.registerConsumer("creeper.explode", (sc, in) -> ((CreeperEntity) in[0].get(sc)).ignite());
    }

    private static <T> T launchProjectile(LivingEntity liv, Class<? extends T> projectile, double scale, Object data) {
        World w = liv.world;
        Entity launch = null;

        if(EntityItemProjectile.class == projectile) {
            if(data == null) {
                throw new NullPointerException("Data musn't be null for EntityItemProjectile");
            }
            ItemStack stack = (ItemStack) data;
            if(stack.isEmpty()) {
                throw new IllegalArgumentException("Empty ItemStack not allowed here");
            }
            launch = new EntityItemProjectile(liv, stack.copy());
            ((EntityItemProjectile) launch).setHeadingFromThrower(liv, liv.rotationPitch, liv.rotationYaw, 0.0f, 1.5f, 1.0f);
        } else if(SnowballEntity.class == projectile) {
            launch = new SnowballEntity(w, liv);
            ((SnowballEntity) launch).shoot(liv.rotationPitch, liv.rotationYaw, 0.0f, 1.5f, 1.0f);
        } else if(EggEntity.class == projectile) {
            launch = new EggEntity(w, liv);
            ((EggEntity) launch).shoot(liv.rotationPitch, liv.rotationYaw, 0.0f, 1.5f, 1.0f);
        } else if(EnderPearlEntity.class == projectile) {
            launch = new EnderPearlEntity(w, liv);
            ((EnderPearlEntity) launch).shoot(liv.rotationPitch, liv.rotationYaw, 0.0f, 1.5f, 1.0f);
        } else if(PotionEntity.class == projectile) {
            launch = new PotionEntity(w, liv);
            ((PotionEntity) launch).setItem((ItemStack) data);
            ((PotionEntity) launch).shoot(liv.rotationPitch, liv.rotationYaw, -20.0f, 0.5f, 1.0f);
        } else if(ExperienceBottleEntity.class == projectile) {
            launch = new ExperienceBottleEntity(w, liv);
            ((ExperienceBottleEntity) launch).shoot(liv.rotationPitch, liv.rotationYaw, -20.0f, 0.7f, 1.0f);
        } else if(AbstractArrowEntity.class.isAssignableFrom(projectile)) {
            if(SpectralArrowEntity.class == projectile) {
                launch = new SpectralArrowEntity(w, liv);
            } else {
                launch = new ArrowEntity(w, liv);
                if(data != null) {
                    ((ArrowEntity) launch).setPotionEffect((ItemStack) data);
                }
            }
            ((AbstractArrowEntity) launch).shoot(liv.rotationPitch, liv.rotationYaw, 0.0F, 3.0F, 1.0F);
        } else if(DamagingProjectileEntity.class.isAssignableFrom(projectile)) {
            Vector3d v = liv.getLookVec().scale(10);
            if(SmallFireballEntity.class == projectile) {
                launch = new SmallFireballEntity(w, liv, v.x, v.y, v.z);
            } else if(WitherSkullEntity.class == projectile) {
                launch = new WitherSkullEntity(w, liv, v.x, v.y, v.z);
            } else if(DragonFireballEntity.class == projectile) {
                launch = new DragonFireballEntity(w, liv, v.x, v.y, v.z);
            } else {
                launch = new FireballEntity(w, liv, v.x, v.y, v.z);
            }
        } else {
            return null;
        }

        launch.setMotion(launch.getMotion().scale(scale));
        w.addEntity(launch);
        return (T) launch;
    }
}