Butterfly 3

package com.mysticsbiomes.common.entity.animal;

import com.google.common.collect.Maps;
import com.mysticsbiomes.MysticsBiomes;
import com.mysticsbiomes.common.block.entity.ButterflyNestBlockEntity;
import com.mysticsbiomes.init.MysticPoiTypes;
import net.minecraft.Util;
import net.minecraft.core.BlockPos;
import net.minecraft.nbt.CompoundTag;
import net.minecraft.nbt.NbtUtils;
import net.minecraft.network.syncher.EntityDataAccessor;
import net.minecraft.network.syncher.EntityDataSerializers;
import net.minecraft.network.syncher.SynchedEntityData;
import net.minecraft.resources.ResourceLocation;
import net.minecraft.server.level.ServerLevel;
import net.minecraft.tags.ItemTags;
import net.minecraft.util.ByIdMap;
import net.minecraft.util.Mth;
import net.minecraft.util.RandomSource;
import net.minecraft.util.StringRepresentable;
import net.minecraft.world.DifficultyInstance;
import net.minecraft.world.InteractionHand;
import net.minecraft.world.InteractionResult;
import net.minecraft.world.entity.*;
import net.minecraft.world.entity.ai.attributes.AttributeSupplier;
import net.minecraft.world.entity.ai.attributes.Attributes;
import net.minecraft.world.entity.ai.control.FlyingMoveControl;
import net.minecraft.world.entity.ai.control.LookControl;
import net.minecraft.world.entity.ai.goal.BreedGoal;
import net.minecraft.world.entity.ai.goal.FloatGoal;
import net.minecraft.world.entity.ai.goal.Goal;
import net.minecraft.world.entity.ai.goal.TemptGoal;
import net.minecraft.world.entity.ai.navigation.FlyingPathNavigation;
import net.minecraft.world.entity.ai.navigation.PathNavigation;
import net.minecraft.world.entity.ai.util.AirAndWaterRandomPos;
import net.minecraft.world.entity.ai.util.AirRandomPos;
import net.minecraft.world.entity.ai.util.HoverRandomPos;
import net.minecraft.world.entity.ai.village.poi.PoiManager;
import net.minecraft.world.entity.ai.village.poi.PoiRecord;
import net.minecraft.world.entity.animal.Animal;
import net.minecraft.world.entity.animal.FlyingAnimal;
import net.minecraft.world.entity.player.Player;
import net.minecraft.world.item.crafting.Ingredient;
import net.minecraft.world.level.Level;
import net.minecraft.world.level.LevelReader;
import net.minecraft.world.level.ServerLevelAccessor;
import net.minecraft.world.level.block.entity.BlockEntity;
import net.minecraft.world.level.block.state.BlockState;
import net.minecraft.world.level.pathfinder.Path;
import net.minecraft.world.phys.Vec3;

import javax.annotation.Nullable;
import java.util.*;
import java.util.function.IntFunction;
import java.util.function.Predicate;
import java.util.stream.Collectors;
import java.util.stream.Stream;

/**
 * Butterflies are friendly ambient anthropoids, useful for growing flowers.
 *
 * When a butterfly is given a flower, they will save it and will look for that same flower in the world.
 * Under specific circumstances
 *      the butterfly has been out of their nest for more than x ticks.
 *      it has pollinated at least one flower beforehand.
 */
public class Butterfly extends Animal implements FlyingAnimal {
    public static final Map<Integer, Map<Integer, ResourceLocation>> TEXTURE_BY_TYPE = Util.make(Maps.newHashMap(), (map) -> {
        for (Type type : Type.values()) {
            map.put(type.getId(), addButterflyVariants(type.getSerializedName()));
        }
    });
    private static final EntityDataAccessor<Integer> DATA_TYPE_ID = SynchedEntityData.defineId(Butterfly.class, EntityDataSerializers.INT);
    private static final int TICKS_PER_FLAP = Mth.ceil(1.4959966F);
    private boolean wantsToEnterNest;
    private int stayOutOfNestCountdown;
    int remainingCooldownBeforeLocatingNewNest;
    @Nullable
    BlockPos nestPos;
    Butterfly.GoToNestGoal goToNestGoal;
    public final AnimationState flyingAnimation = new AnimationState();

    public Butterfly(EntityType<? extends Butterfly> type, Level level) {
        super(type, level);
        this.moveControl = new FlyingMoveControl(this, 20, true);
        this.lookControl = new LookControl(this);
    }

    public static AttributeSupplier.Builder createAttributes() {
        return Mob.createMobAttributes().add(Attributes.MAX_HEALTH, 6.0D).add(Attributes.FLYING_SPEED, 1.0D).add(Attributes.FOLLOW_RANGE, 16.0D);
    }

    protected void registerGoals() {
        this.goalSelector.addGoal(1, new EnterNestGoal());
        this.goalSelector.addGoal(2, new BreedGoal(this, 1.0D));
        this.goalSelector.addGoal(3, new TemptGoal(this, 1.25D, Ingredient.of(ItemTags.FLOWERS), false));
        this.goalSelector.addGoal(5, new LocateNestGoal());
        this.goToNestGoal = new GoToNestGoal();
        this.goalSelector.addGoal(8, new WanderGoal());
        this.goalSelector.addGoal(9, new FloatGoal(this));
    }

    protected void defineSynchedData() {
        super.defineSynchedData();
        this.entityData.define(DATA_TYPE_ID, this.random.nextInt(1, 2));
    }

    @Override
    public void addAdditionalSaveData(CompoundTag tag) {
        super.addAdditionalSaveData(tag);

        if (this.nestPos != null) {
            tag.put("NestPos", NbtUtils.writeBlockPos(this.nestPos));
        }

        tag.putString("Type", this.getVariant().name);
        tag.putInt("CannotEnterNestTicks", this.stayOutOfNestCountdown);
    }

    @Override
    public void readAdditionalSaveData(CompoundTag tag) {
        this.nestPos = null;
        if (tag.contains("NestPos")) {
            this.nestPos = NbtUtils.readBlockPos(tag.getCompound("NestPos"));
        }

        super.readAdditionalSaveData(tag);
        this.setVariant(Type.byName(tag.getString("Type")));
        this.stayOutOfNestCountdown = tag.getInt("CannotEnterNestTicks");
    }

    @Nullable
    @Override
    public AgeableMob getBreedOffspring(ServerLevel level, AgeableMob mob) {
        return null;
    }

    @Override
    public MobType getMobType() {
        return MobType.ARTHROPOD;
    }

    public Type getVariant() {
        return Type.byId(this.entityData.get(DATA_TYPE_ID));
    }

    public void setVariant(Type type) {
        this.entityData.set(DATA_TYPE_ID, type.getId());
    }

    public static Map<Integer, ResourceLocation> addButterflyVariants(String type) {
        Map<Integer, ResourceLocation> map = Maps.newHashMap();
        map.put(1, MysticsBiomes.modLoc("textures/entity/butterfly/" + type + ".png"));
        map.put(2, MysticsBiomes.modLoc("textures/entity/butterfly/" + type + "_still.png"));
        return map;
    }

    @Nullable
    @Override
    public SpawnGroupData finalizeSpawn(ServerLevelAccessor accessor, DifficultyInstance instance, MobSpawnType spawnType, @Nullable SpawnGroupData spawnData, @Nullable CompoundTag tag) {
        RandomSource source = accessor.getRandom();

        Butterfly.Type type;
        if (spawnData instanceof Butterfly.GroupData) {
            type = ((GroupData)spawnData).type;
        } else {
            type = Util.getRandom(Butterfly.Type.values(), source);
            spawnData = new Butterfly.GroupData(type);
        }

        this.setVariant(type);
        return super.finalizeSpawn(accessor, instance, spawnType, spawnData, tag);
    }

    //////////////////////////////////////////////////////////////////

    public void tick() {
        if (this.level().isClientSide()) {
            this.flyingAnimation.animateWhen(this.walkAnimation.isMoving(), this.tickCount);
        }
        super.tick();
    }

    public void aiStep() {
        super.aiStep();
        if (!this.level().isClientSide) {
            if (this.stayOutOfNestCountdown > 0) {
                --this.stayOutOfNestCountdown;
            }

            if (this.remainingCooldownBeforeLocatingNewNest > 0) {
                --this.remainingCooldownBeforeLocatingNewNest;
            }

            if (this.tickCount % 20 == 0 && !this.isNestValid()) {
                this.nestPos = null;
            }
        }
    }

    protected void customServerAiStep() {

    }

    //////////////////////////////////////////////////////////////////

    protected void updateWalkAnimation(float f) {
        this.walkAnimation.update(Math.min(f * 25.0F, 1.0F), 0.4F);
    }

    //////////////////////////////////////////////////////////////////

    public float getWalkTargetValue(BlockPos pos, LevelReader reader) {
        return reader.getBlockState(pos).isAir() ? 6.0F : 0.0F;
    }

    protected void checkFallDamage(double distance, boolean b, BlockState state, BlockPos pos) {
    }

    protected PathNavigation createNavigation(Level level) {
        FlyingPathNavigation navigation = new FlyingPathNavigation(this, level) {
            public boolean isStableDestination(BlockPos pos) {
                return !this.level.getBlockState(pos.below()).isAir();
            }
        };
        navigation.setCanOpenDoors(false);
        navigation.setCanFloat(false);
        navigation.setCanPassDoors(true);
        return navigation;
    }

    private Optional<BlockPos> findNearbyFlower() {
        return Butterfly.this.findNearestBlock(null, 5.0D);
    }

    private Optional<BlockPos> findNearestBlock(Predicate<BlockState> predicate, double distance) {
        BlockPos pos = Butterfly.this.blockPosition();
        BlockPos.MutableBlockPos mutablePos = new BlockPos.MutableBlockPos();

        for (int i = 0; (double)i <= distance; i = i > 0 ? -i : 1 - i) {
            for (int j = 0; (double)j < distance; ++j) {
                for (int k = 0; k <= j; k = k > 0 ? -k : 1 - k) {
                    for (int l = k < j && k > -j ? j : 0; l <= j; l = l > 0 ? -l : 1 - l) {
                        mutablePos.setWithOffset(pos, k, i - 1, l);
                        if (pos.closerThan(mutablePos, distance) && predicate.test(Butterfly.this.level().getBlockState(mutablePos))) {
                            return Optional.of(mutablePos);
                        }
                    }
                }
            }
        }
        return Optional.empty();
    }

    private void pathfindRandomlyTowards(BlockPos pos) {
        Vec3 vec3 = Vec3.atBottomCenterOf(pos);
        int i = 0;
        BlockPos pos1 = this.blockPosition();
        int j = (int)vec3.y - pos1.getY();
        if (j > 2) {
            i = 4;
        } else if (j < -2) {
            i = -4;
        }

        int k = 6;
        int l = 8;
        int i1 = pos1.distManhattan(pos);
        if (i1 < 15) {
            k = i1 / 2;
            l = i1 / 2;
        }

        Vec3 vec31 = AirRandomPos.getPosTowards(this, k, l, i, vec3, ((float)Math.PI / 10F));
        if (vec31 != null) {
            this.navigation.setMaxVisitedNodesMultiplier(0.5F);
            this.navigation.moveTo(vec31.x, vec31.y, vec31.z, 1.0D);
        }
    }

    private boolean pathfindDirectlyTowards(BlockPos pos) {
        Butterfly.this.navigation.setMaxVisitedNodesMultiplier(10.0F);
        Butterfly.this.navigation.moveTo(pos.getX(), pos.getY(), pos.getZ(), 1.0D);
        return Butterfly.this.navigation.getPath() != null && Butterfly.this.navigation.getPath().canReach();
    }

    private boolean closerThan(BlockPos pos, int distance) {
        return pos.closerThan(this.blockPosition(), distance);
    }

    private boolean isTooFarAway(BlockPos pos) {
        return !this.closerThan(pos, 32);
    }

    public boolean isFlying() {
        return !this.level().isNight() || !this.isInWaterRainOrBubble();
    }

    public boolean isFlapping() {
        return this.isFlying() && this.tickCount % TICKS_PER_FLAP == 0;
    }

    //////////////////////////////////////////////////////////////////

    public void setStayOutOfNestCountdown(int ticks) {
        this.stayOutOfNestCountdown = ticks;
    }

    public boolean isNestValid() {
        if (this.nestPos != null) {
            if (this.isTooFarAway(this.nestPos)) {
                return false;
            } else {
                BlockEntity blockEntity = this.level().getBlockEntity(this.nestPos);
                return blockEntity instanceof ButterflyNestBlockEntity;
            }
        } else {
            return false;
        }
    }

    /** @return true if it's raining, night or their nest is on fire. */
    public boolean wantsToEnterNest() {
        if (this.stayOutOfNestCountdown <= 0) {
            return this.wantsToEnterNest || this.level().isRaining() || this.level().isNight() && !this.isNestNearFire();
        } else {
            return false;
        }
    }

    private boolean doesNestHaveSpace(BlockPos pos) {
        BlockEntity blockEntity = this.level().getBlockEntity(pos);

        if (blockEntity instanceof ButterflyNestBlockEntity entity) {
            return !entity.isFull();
        } else {
            return false;
        }
    }

    private boolean isNestNearFire() {
        if (this.nestPos == null) {
            return false;
        } else {
            BlockEntity blockEntity = this.level().getBlockEntity(this.nestPos);
            return blockEntity instanceof ButterflyNestBlockEntity && ((ButterflyNestBlockEntity)blockEntity).isFireNearby();
        }
    }

    //////////////////////////////////////////////////////////////////

    @Override
    public InteractionResult mobInteract(Player player, InteractionHand hand) {
        if (player.isCreative()) {
            this.remainingCooldownBeforeLocatingNewNest = 0;
            this.stayOutOfNestCountdown = 0;
            return InteractionResult.SUCCESS;
        }
        return super.mobInteract(player, hand);
    }

    //////////////////////////////////////////////////////////////////

    class LocateNestGoal extends Goal {

        public boolean canUse() {
            return Butterfly.this.remainingCooldownBeforeLocatingNewNest == 0 && Butterfly.this.nestPos != null && Butterfly.this.wantsToEnterNest();
        }

        public boolean canContinueToUse() {
            return false;
        }

        public void start() {
            Butterfly.this.remainingCooldownBeforeLocatingNewNest = 200;

            List<BlockPos> list = this.findNearbyNestWithSpace();
            if (!list.isEmpty()) {
                for (BlockPos pos : list) {
                    Butterfly.this.nestPos = pos;
                    return;
                }
                Butterfly.this.nestPos = list.get(0);
            }
        }

        private List<BlockPos> findNearbyNestWithSpace() {
            BlockPos pos = Butterfly.this.blockPosition();
            PoiManager manager = ((ServerLevel)Butterfly.this.level()).getPoiManager();

            Stream<PoiRecord> stream = manager.getInRange((holder) -> holder.is(MysticPoiTypes.BUTTERFLY_NEST.getId()), pos, 26, PoiManager.Occupancy.ANY);
            return stream.map(PoiRecord::getPos).filter(Butterfly.this::doesNestHaveSpace).sorted(Comparator.comparingDouble((pos1) -> pos1.distSqr(pos))).collect(Collectors.toList());
        }
    }

    public class GoToNestGoal extends Goal {
        int travellingTicks = Butterfly.this.level().random.nextInt(10);
        @Nullable
        private Path lastPath;

        GoToNestGoal() {
            this.setFlags(EnumSet.of(Goal.Flag.MOVE));
        }

        public boolean canUse() {
            return Butterfly.this.nestPos != null && !Butterfly.this.hasRestriction() && Butterfly.this.wantsToEnterNest() && !this.hasReachedTarget(Butterfly.this.nestPos);
        }

        public boolean canContinueToUse() {
            return this.canUse();
        }

        public void start() {
            this.travellingTicks = 0;
            super.start();
        }

        public void stop() {
            this.travellingTicks = 0;
            Butterfly.this.navigation.stop();
            Butterfly.this.navigation.resetMaxVisitedNodesMultiplier();
        }

        public void tick() {
            if (Butterfly.this.nestPos != null) {
                ++this.travellingTicks;

                if (!Butterfly.this.navigation.isInProgress()) {
                    if (!Butterfly.this.closerThan(Butterfly.this.nestPos, 16)) {
                        if (!Butterfly.this.isTooFarAway(Butterfly.this.nestPos)) {
                            Butterfly.this.pathfindRandomlyTowards(Butterfly.this.nestPos);
                        }
                    } else {
                        boolean flag = Butterfly.this.pathfindDirectlyTowards(Butterfly.this.nestPos);

                        if (flag && this.lastPath != null) {
                            this.lastPath = Butterfly.this.navigation.getPath();
                        }
                    }
                }
            }
        }

        private boolean hasReachedTarget(BlockPos pos) {
            if (Butterfly.this.closerThan(pos, 2)) {
                return true;
            } else {
                Path path = Butterfly.this.navigation.getPath();
                return path != null && path.getTarget().equals(pos) && path.canReach() && path.isDone();
            }
        }
    }

    class EnterNestGoal extends Goal {

        public boolean canUse() {
            if (Butterfly.this.nestPos != null && Butterfly.this.wantsToEnterNest() && Butterfly.this.nestPos.closerToCenterThan(Butterfly.this.position(), 2.0D)) {
                BlockEntity blockEntity = Butterfly.this.level().getBlockEntity(Butterfly.this.nestPos);

                if (blockEntity instanceof ButterflyNestBlockEntity entity) {
                    if (!entity.isFull()) {
                        return true;
                    }
                    Butterfly.this.nestPos = null;
                }
            }
            return false;
        }

        public boolean canContinueToUse() {
            return false;
        }

        public void start() {
            if (Butterfly.this.nestPos != null) {
                BlockEntity blockEntity = Butterfly.this.level().getBlockEntity(Butterfly.this.nestPos);

                if (blockEntity instanceof ButterflyNestBlockEntity entity) {
                    entity.addOccupant(Butterfly.this);
                }
            }
        }

        public void stop() {
            Butterfly.this.wantsToEnterNest = false;
        }
    }

    class WanderGoal extends Goal {

        WanderGoal() {
            this.setFlags(EnumSet.of(Goal.Flag.MOVE));
        }

        public boolean canUse() {
            return Butterfly.this.navigation.isDone() && Butterfly.this.random.nextInt(10) == 0;
        }

        public boolean canContinueToUse() {
            return Butterfly.this.navigation.isInProgress();
        }

        public void start() {
            Vec3 vec3 = this.findPos();

            if (vec3 != null) {
                Butterfly.this.navigation.moveTo(Butterfly.this.navigation.createPath(BlockPos.containing(vec3), 1), 1.0D);
            }
        }

        private Vec3 findPos() {
            Vec3 vec3;
            if (Butterfly.this.nestPos != null && !Butterfly.this.closerThan(Butterfly.this.nestPos, 22)) {
                Vec3 vec31 = Vec3.atCenterOf(Butterfly.this.nestPos);
                vec3 = vec31.subtract(Butterfly.this.position()).normalize();
            } else {
                vec3 = Butterfly.this.getViewVector(0.0F);
            }

            Vec3 vec32 = HoverRandomPos.getPos(Butterfly.this, 8, 7, vec3.x, vec3.z, ((float)Math.PI / 2F), 3, 1);
            return vec32 != null ? vec32 : AirAndWaterRandomPos.getPos(Butterfly.this, 8, 4, -2, vec3.x, vec3.z, ((float)Math.PI / 2F));
        }
    }

    public static class GroupData extends AgeableMob.AgeableMobGroupData {
        public final Butterfly.Type type;

        GroupData(Butterfly.Type type) {
            super(true);
            this.type = type;
        }
    }

    public enum Type implements StringRepresentable {
        MONARCH(0, "monarch"),
        MORPHO(1, "morpho");

        public static final StringRepresentable.EnumCodec<Type> CODEC = StringRepresentable.fromEnum(Type::values);
        private static final IntFunction<Type> BY_ID = ByIdMap.continuous(Type::getId, values(), ByIdMap.OutOfBoundsStrategy.ZERO);
        private final int id;
        private final String name;

        Type(int id, String name) {
            this.id = id;
            this.name = name;
        }

        public String getSerializedName() {
            return this.name;
        }

        public int getId() {
            return this.id;
        }

        public static Type byName(String name) {
            return CODEC.byName(name, MONARCH);
        }

        public static Type byId(int id) {
            return BY_ID.apply(id);
        }
    }

    /**
     * TODO:
     *
     */
    static class Butterfly2 extends Animal implements FlyingAnimal {

        public Butterfly2(EntityType<? extends Butterfly2> type, Level level) {
            super(type, level);
        }

        public static AttributeSupplier.Builder createAttributes() {
            return Mob.createMobAttributes().add(Attributes.MAX_HEALTH, 6.0D).add(Attributes.FLYING_SPEED, 1.0D).add(Attributes.FOLLOW_RANGE, 16.0D);
        }

        @Override
        public AgeableMob getBreedOffspring(ServerLevel level, AgeableMob mob) {
            return null;
        }

        @Override
        public boolean isFlying() {
            return false;
        }
    }

}