Twine/src/main/java/com/thebrokenrail/twine/entity/FollowPassiveEntityGoal.java

package com.thebrokenrail.twine.entity;

import com.thebrokenrail.twine.component.StageDataComponent;
import net.minecraft.entity.LivingEntity;
import net.minecraft.entity.ai.goal.FollowTargetGoal;
import net.minecraft.entity.mob.Angerable;
import net.minecraft.entity.mob.HostileEntity;
import net.minecraft.entity.mob.MobEntity;
import net.minecraft.entity.passive.PassiveEntity;
import net.minecraft.entity.player.PlayerEntity;
import net.minecraft.server.world.ServerWorld;
import net.minecraft.util.math.BlockPos;

public class FollowPassiveEntityGoal extends FollowTargetGoal<LivingEntity> {
    public FollowPassiveEntityGoal(MobEntity mob) {
        super(mob, LivingEntity.class, 10, false, false, entity -> {
            StageDataComponent component = StageDataComponent.getFromWorld((ServerWorld) mob.getEntityWorld());
            BlockPos pos = mob.getBlockPos();
            int stage = component.findEffectiveStageOfChunk((ServerWorld) mob.getEntityWorld(), pos);

            return entity instanceof PassiveEntity || (entity instanceof PlayerEntity && stage >= 4);
        });
    }

    public boolean canStart() {
        StageDataComponent component = StageDataComponent.getFromWorld((ServerWorld) mob.getEntityWorld());
        BlockPos pos = mob.getBlockPos();
        int stage = component.findEffectiveStageOfChunk((ServerWorld) mob.getEntityWorld(), pos);

        return mob instanceof HostileEntity && (!(mob instanceof Angerable) || stage >= 4) && stage >= 2 && super.canStart();
    }
}