Sea Otter 2

package com.mysticsbiomes.common.entity.animal;

import com.mysticsbiomes.init.MysticEntities;
import com.mysticsbiomes.init.MysticItems;
import net.minecraft.core.BlockPos;
import net.minecraft.nbt.CompoundTag;
import net.minecraft.network.syncher.EntityDataAccessor;
import net.minecraft.network.syncher.EntityDataSerializers;
import net.minecraft.network.syncher.SynchedEntityData;
import net.minecraft.server.level.ServerLevel;
import net.minecraft.util.Mth;
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.behavior.BehaviorUtils;
import net.minecraft.world.entity.ai.control.MoveControl;
import net.minecraft.world.entity.ai.control.SmoothSwimmingLookControl;
import net.minecraft.world.entity.ai.control.SmoothSwimmingMoveControl;
import net.minecraft.world.entity.ai.goal.*;
import net.minecraft.world.entity.ai.navigation.PathNavigation;
import net.minecraft.world.entity.ai.navigation.WaterBoundPathNavigation;
import net.minecraft.world.entity.ai.util.DefaultRandomPos;
import net.minecraft.world.entity.animal.Animal;
import net.minecraft.world.entity.player.Player;
import net.minecraft.world.item.Items;
import net.minecraft.world.item.crafting.Ingredient;
import net.minecraft.world.level.Level;
import net.minecraft.world.level.LevelReader;
import net.minecraft.world.level.block.state.BlockState;
import net.minecraft.world.level.pathfinder.BlockPathTypes;
import net.minecraft.world.level.pathfinder.PathComputationType;
import net.minecraft.world.level.pathfinder.PathFinder;
import net.minecraft.world.level.pathfinder.SwimNodeEvaluator;
import net.minecraft.world.phys.Vec3;

import javax.annotation.Nullable;
import java.util.EnumSet;

/**
 * swims on the surface.
 * only swims underwater fully when playing or searching for food.
 */
public class SeaOtter extends Animal {
    private static final EntityDataAccessor<Boolean> DATA_FLOATING_ID = SynchedEntityData.defineId(SeaOtter.class, EntityDataSerializers.BOOLEAN);
    private static final EntityDataAccessor<Boolean> DATA_SWIMMING_ID = SynchedEntityData.defineId(SeaOtter.class, EntityDataSerializers.BOOLEAN);
    int ticksSinceSearchedForFood;
    int cooldownBeforeFindingFoodAgain;
    boolean needsToSurface;

    public final AnimationState idleAnimationState = new AnimationState();

    public SeaOtter(EntityType<? extends SeaOtter> type, Level level) {
        super(type, level);
        this.setPathfindingMalus(BlockPathTypes.WATER, 0.0F);
        this.moveControl = new SeaOtter.SeaOtterMoveControl(this);
        this.lookControl = new SmoothSwimmingLookControl(this, 10);
    }

    protected void defineSynchedData() {
        super.defineSynchedData();
        this.entityData.define(DATA_FLOATING_ID, false);
        this.entityData.define(DATA_SWIMMING_ID, false);
    }

    protected void registerGoals() {
        this.goalSelector.addGoal(1, new SeaOtter.SearchForFoodGoal());
        this.goalSelector.addGoal(2, new SeaOtter.FloatAtSurfaceGoal());
        this.goalSelector.addGoal(2, new SeaOtter.SwimToSurfaceGoal(this, 1.0D, 16));
        this.goalSelector.addGoal(3, new SeaOtter.SwimAroundGoal(this, 1.0D, 10));
        this.goalSelector.addGoal(4, new LookAtPlayerGoal(this, Player.class, 6.0F));
        this.goalSelector.addGoal(5, new RandomLookAroundGoal(this));
    }

    public static AttributeSupplier.Builder createAttributes() {
        return Mob.createMobAttributes().add(Attributes.MAX_HEALTH, 16.0F).add(Attributes.MOVEMENT_SPEED, 0.25D);
    }

    public void addAdditionalSaveData(CompoundTag tag) {
        super.addAdditionalSaveData(tag);
        tag.putBoolean("Floating", this.isFloating());
        tag.putBoolean("Swimming", this.isSwimming());
    }

    public void readAdditionalSaveData(CompoundTag tag) {
        super.readAdditionalSaveData(tag);
        this.setFloating(tag.getBoolean("Floating"));
        this.setSwimming(tag.getBoolean("Swimming"));
    }

    public AgeableMob getBreedOffspring(ServerLevel level, AgeableMob mob) {
        return MysticEntities.SEA_OTTER.get().create(level);
    }

    public MobType getMobType() {
        return MobType.WATER;
    }

    // TICKS & A.I.

    public void tick() {
        super.tick();

        if (this.level().isClientSide) {
            this.idleAnimationState.animateWhen(this.isFloating(), this.tickCount);
        }
    }

    public void aiStep() {
        super.aiStep();
        ++this.ticksSinceSearchedForFood;

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

        if (this.isFloating()) {
            this.setDeltaMovement(this.getDeltaMovement().multiply(1.0D, 0.0D, 1.0D));
            this.setYya(0.0F);
        }

        if (this.isUnderWater() && (this.getAirSupply() < 200 || this.random.nextFloat() <= 0.001F)) {
            this.setNeedsToSurface(true);
        }
    }

    // NAVIGATION & MOVEMENT

    protected PathNavigation createNavigation(Level level) {
        return new SeaOtter.SeaOtterPathNavigation(this, level);
    }

    public void travel(Vec3 vec3) {
        if (this.isEffectiveAi() && this.isInWater()) {
            this.moveRelative(this.getSpeed(), vec3);
            this.move(MoverType.SELF, this.getDeltaMovement());
            this.setDeltaMovement(this.getDeltaMovement().scale(0.9D));
        } else {
            super.travel(vec3);
        }
    }

    public int getMaxAirSupply() {
        return 6000;
    }

    public boolean isFloating() {
        return this.entityData.get(DATA_FLOATING_ID);
    }

    public void setFloating(boolean value) {
        this.entityData.set(DATA_FLOATING_ID, value);
        this.entityData.set(DATA_SWIMMING_ID, !value);
    }

    public boolean isSwimming() {
        return this.entityData.get(DATA_SWIMMING_ID);
    }

    public void setSwimming(boolean value) {
        this.entityData.set(DATA_SWIMMING_ID, value);
        this.entityData.set(DATA_FLOATING_ID, !value);
    }

    public boolean wantsToSearchForFood() {
        return this.ticksSinceSearchedForFood > 200;
    }

    /**
     * @return when they need to breath above water, to eat, or to sleep.
     */
    public boolean needsToSurface() {
        return this.needsToSurface;
    }

    public void setNeedsToSurface(boolean needsToSurface) {
        this.needsToSurface = needsToSurface;
    }

    @Override
    public InteractionResult mobInteract(Player player, InteractionHand hand) {
        boolean flag = player.getItemInHand(hand).is(Items.STICK);
        this.setFloating(flag);
        this.setSwimming(!flag);
        return super.mobInteract(player, hand);
    }

    // TODO: GOALS

    class FloatAtSurfaceGoal extends Goal {

        public boolean canUse() {
            return !SeaOtter.this.isSwimming() || !SeaOtter.this.wantsToSearchForFood();
        }

        public void start() {
            SeaOtter.this.setFloating(true);
        }
    }

    class SwimToSurfaceGoal extends MoveToBlockGoal {

        SwimToSurfaceGoal(PathfinderMob mob, double speed, int range) {
            super(mob, speed, range);
            this.setFlags(EnumSet.of(Flag.MOVE, Flag.LOOK));
        }

        public boolean canUse() {
            return super.canUse() && SeaOtter.this.needsToSurface() && !SeaOtter.this.isFloating() && !SeaOtter.this.onGround();
        }

        public void stop() {
            SeaOtter.this.setNeedsToSurface(false);
            SeaOtter.this.getNavigation().stop();
        }

        protected boolean isValidTarget(LevelReader reader, BlockPos pos) {
            return reader.getBlockState(pos.above()).isAir();
        }
    }

    class SwimAroundGoal extends RandomSwimmingGoal {

        SwimAroundGoal(PathfinderMob mob, double speed, int interval) {
            super(mob, speed, interval);
        }

        public boolean canUse() {
            return super.canUse() && !SeaOtter.this.isFloating();
        }

        public boolean canContinueToUse() {
            return super.canContinueToUse() && !SeaOtter.this.isFloating();
        }

        public void start() {
            super.start();
            SeaOtter.this.setSwimming(true);
        }
    }

    class SearchForFoodGoal extends Goal {
        int ticks;

        public boolean canUse() {
            return SeaOtter.this.wantsToSearchForFood() && SeaOtter.this.cooldownBeforeFindingFoodAgain <= 0;
        }

        public boolean canContinueToUse() {
            return !this.hasSearchedLongEnough();
        }

        public void start() {
            this.ticks = 0;
            SeaOtter.this.setSwimming(true);
        }

        public void stop() {
            if (this.hasSearchedLongEnough()) {
                SeaOtter.this.cooldownBeforeFindingFoodAgain = 200;
                SeaOtter.this.setFloating(true);
            }

            this.ticks = 0;
        }

        private boolean hasSearchedLongEnough() {
            return this.ticks > 200;
        }

        public void tick() {
            ++this.ticks;
        }
    }

    // TODO: NAVIGATION & CONTROLS

    class SeaOtterMoveControl extends MoveControl {

        public SeaOtterMoveControl(Mob mob) {
            super(mob);
        }

        @Override
        public void tick() {
            if (SeaOtter.this.isInWater()) {
                SeaOtter.this.setDeltaMovement(SeaOtter.this.getDeltaMovement().add(SeaOtter.this.getLookAngle().scale(SeaOtter.this.isFloating() ? 0.002F : 0.005F)));

                if (!SeaOtter.this.isFloating()) {
                    if (this.operation == Operation.MOVE_TO && !this.mob.getNavigation().isDone()) {
                        double d0 = this.wantedX - this.mob.getX();
                        double d1 = this.wantedY - this.mob.getY();
                        double d2 = this.wantedZ - this.mob.getZ();
                        double distanceSqr = d0 * d0 + d1 * d1 + d2 * d2;

                        if (distanceSqr < (double) 2.5000003E-7F) {
                            this.mob.setZza(0.0F);
                        } else {
                            float yRot = (float) (Mth.atan2(d2, d0) * (double) (180F / (float) Math.PI)) - 90.0F;
                            this.mob.setYRot(this.rotlerp(this.mob.getYRot(), yRot, 40.0F));
                            this.mob.yBodyRot = this.mob.getYRot();
                            this.mob.yHeadRot = this.mob.getYRot();
                            float speed = (float) (this.speedModifier * this.mob.getAttributeValue(Attributes.MOVEMENT_SPEED));
                            this.mob.setSpeed(speed * 0.2F);

                            double horizontalDistance = Math.sqrt(d0 * d0 + d2 * d2);
                            if (Math.abs(d1) > (double) 1.0E-5F || Math.abs(horizontalDistance) > (double) 1.0E-5F) {
                                float xRot = -((float) (Mth.atan2(d1, horizontalDistance) * (double) (180F / (float) Math.PI)));
                                xRot = Mth.clamp(Mth.wrapDegrees(xRot), -180.0F, 180.0F);
                                this.mob.setXRot(this.rotlerp(this.mob.getXRot(), xRot, 45.0F));
                            }

                            BlockPos wantedPos = new BlockPos((int) this.wantedX, (int) this.wantedY, (int) this.wantedZ);
                            BlockState wantedBlockState = this.mob.level().getBlockState(wantedPos);

                            if (d1 > (double) this.mob.maxUpStep() && d0 * d0 + d2 * d2 < 4.0F && d1 <= 1.0D && wantedBlockState.getFluidState().isEmpty()) {
                                this.mob.getJumpControl().jump();
                                this.mob.setSpeed(speed);
                            }

                            float f0 = Mth.cos(this.mob.getXRot() * ((float) Math.PI / 180F));
                            float f1 = Mth.sin(this.mob.getXRot() * ((float) Math.PI / 180F));
                            this.mob.zza = f0 * speed;
                            this.mob.yya = -f1 * (speed);
                        }
                    } else {
                        this.mob.setSpeed(0.0F);
                        this.mob.setXxa(0.0F);
                        this.mob.setYya(0.0F);
                        this.mob.setZza(0.0F);
                    }
                }
            } else {
                super.tick();
            }
        }
    }

    static class SeaOtterPathNavigation extends WaterBoundPathNavigation {

        public SeaOtterPathNavigation(Mob mob, Level level) {
            super(mob, level);
        }

        protected PathFinder createPathFinder(int nodes) {
            this.nodeEvaluator = new SwimNodeEvaluator(true);
            return new PathFinder(this.nodeEvaluator, nodes);
        }

        protected boolean canUpdatePath() {
            return true;
        }
    }

}