package net.minecraft.world.entity.ai.goal;

import org.apache.logging.log4j.LogManager;
import java.util.stream.Stream;
import java.util.function.Predicate;
import com.google.common.collect.Sets;
import java.util.EnumMap;
import java.util.EnumSet;
import net.minecraft.util.profiling.ProfilerFiller;
import java.util.function.Supplier;
import java.util.Set;
import java.util.Map;
import org.apache.logging.log4j.Logger;

public class GoalSelector {
    private static final Logger LOGGER;
    private static final WrappedGoal NO_GOAL;
    private final Map<Goal.Flag, WrappedGoal> lockedFlags;
    private final Set<WrappedGoal> availableGoals;
    private final Supplier<ProfilerFiller> profiler;
    private final EnumSet<Goal.Flag> disabledFlags;
    private int newGoalRate;
    
    public GoalSelector(final Supplier<ProfilerFiller> supplier) {
        this.lockedFlags = new EnumMap<Goal.Flag, WrappedGoal>(Goal.Flag.class);
        this.availableGoals = Sets.newLinkedHashSet();
        this.disabledFlags = EnumSet.<Goal.Flag>noneOf(Goal.Flag.class);
        this.newGoalRate = 3;
        this.profiler = supplier;
    }
    
    public void addGoal(final int integer, final Goal aug) {
        this.availableGoals.add(new WrappedGoal(integer, aug));
    }
    
    public void removeGoal(final Goal aug) {
        this.availableGoals.stream().filter(avx -> avx.getGoal() == aug).filter(WrappedGoal::isRunning).forEach(WrappedGoal::stop);
        this.availableGoals.removeIf(avx -> avx.getGoal() == aug);
    }
    
    public void tick() {
        final ProfilerFiller ami2 = this.profiler.get();
        ami2.push("goalCleanup");
        this.getRunningGoals().filter(avx -> !avx.isRunning() || avx.getFlags().stream().anyMatch(this.disabledFlags::contains) || !avx.canContinueToUse()).forEach(Goal::stop);
        this.lockedFlags.forEach((a, avx) -> {
            if (!avx.isRunning()) {
                this.lockedFlags.remove(a);
            }
            return;
        });
        ami2.pop();
        ami2.push("goalUpdate");
        final WrappedGoal avx2;
        this.availableGoals.stream().filter(avx -> !avx.isRunning()).filter(avx -> avx.getFlags().stream().noneMatch(this.disabledFlags::contains)).filter(avx -> avx.getFlags().stream().allMatch(a -> this.lockedFlags.getOrDefault(a, GoalSelector.NO_GOAL).canBeReplacedBy(avx))).filter(WrappedGoal::canUse).forEach(avx -> {
            avx.getFlags().forEach(a -> {
                avx2 = this.lockedFlags.getOrDefault(a, GoalSelector.NO_GOAL);
                avx2.stop();
                this.lockedFlags.put(a, avx);
                return;
            });
            avx.start();
            return;
        });
        ami2.pop();
        ami2.push("goalTick");
        this.getRunningGoals().forEach(WrappedGoal::tick);
        ami2.pop();
    }
    
    public Stream<WrappedGoal> getRunningGoals() {
        return this.availableGoals.stream().filter(WrappedGoal::isRunning);
    }
    
    public void disableControlFlag(final Goal.Flag a) {
        this.disabledFlags.add(a);
    }
    
    public void enableControlFlag(final Goal.Flag a) {
        this.disabledFlags.remove(a);
    }
    
    public void setControlFlag(final Goal.Flag a, final boolean boolean2) {
        if (boolean2) {
            this.enableControlFlag(a);
        }
        else {
            this.disableControlFlag(a);
        }
    }
    
    static {
        LOGGER = LogManager.getLogger();
        NO_GOAL = new WrappedGoal(Integer.MAX_VALUE, new Goal() {
            @Override
            public boolean canUse() {
                return false;
            }
        }) {
            @Override
            public boolean isRunning() {
                return false;
            }
        };
    }
}
