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.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 ProfilerFiller profiler;
    private final EnumSet<Goal.Flag> disabledFlags;
    private int newGoalRate;
    
    public GoalSelector(final ProfilerFiller aio) {
        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 = aio;
    }
    
    public void addGoal(final int integer, final Goal apg) {
        this.availableGoals.add(new WrappedGoal(integer, apg));
    }
    
    public void removeGoal(final Goal apg) {
        this.availableGoals.stream().filter(aqv -> aqv.getGoal() == apg).filter(WrappedGoal::isRunning).forEach(WrappedGoal::stop);
        this.availableGoals.removeIf(aqv -> aqv.getGoal() == apg);
    }
    
    public void tick() {
        this.profiler.push("goalCleanup");
        this.getRunningGoals().filter(aqv -> !aqv.isRunning() || aqv.getFlags().stream().anyMatch(this.disabledFlags::contains) || !aqv.canContinueToUse()).forEach(Goal::stop);
        this.lockedFlags.forEach((a, aqv) -> {
            if (!aqv.isRunning()) {
                this.lockedFlags.remove(a);
            }
            return;
        });
        this.profiler.pop();
        this.profiler.push("goalUpdate");
        final WrappedGoal aqv2;
        this.availableGoals.stream().filter(aqv -> !aqv.isRunning()).filter(aqv -> aqv.getFlags().stream().noneMatch(this.disabledFlags::contains)).filter(aqv -> aqv.getFlags().stream().allMatch(a -> this.lockedFlags.getOrDefault(a, GoalSelector.NO_GOAL).canBeReplacedBy(aqv))).filter(WrappedGoal::canUse).forEach(aqv -> {
            aqv.getFlags().forEach(a -> {
                aqv2 = this.lockedFlags.getOrDefault(a, GoalSelector.NO_GOAL);
                aqv2.stop();
                this.lockedFlags.put(a, aqv);
                return;
            });
            aqv.start();
            return;
        });
        this.profiler.pop();
        this.profiler.push("goalTick");
        this.getRunningGoals().forEach(WrappedGoal::tick);
        this.profiler.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;
            }
        };
    }
}
