/*
 * Decompiled with CFR 0.152.
 * 
 * Could not load the following classes:
 *  com.google.common.annotations.VisibleForTesting
 *  com.google.common.collect.ImmutableList
 *  com.google.common.collect.ImmutableList$Builder
 *  com.google.common.collect.ImmutableSet
 *  com.google.common.collect.Maps
 *  com.google.common.collect.Sets
 *  com.mojang.datafixers.util.Pair
 *  com.mojang.logging.LogUtils
 *  com.mojang.serialization.Codec
 *  com.mojang.serialization.DataResult
 *  com.mojang.serialization.Dynamic
 *  com.mojang.serialization.DynamicOps
 *  com.mojang.serialization.MapCodec
 *  com.mojang.serialization.MapLike
 *  com.mojang.serialization.RecordBuilder
 *  it.unimi.dsi.fastutil.objects.ObjectArrayList
 *  org.apache.commons.lang3.mutable.MutableObject
 *  org.jspecify.annotations.Nullable
 *  org.slf4j.Logger
 */
package net.minecraft.world.entity.ai;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Maps;
import com.google.common.collect.Sets;
import com.mojang.datafixers.util.Pair;
import com.mojang.logging.LogUtils;
import com.mojang.serialization.Codec;
import com.mojang.serialization.DataResult;
import com.mojang.serialization.Dynamic;
import com.mojang.serialization.DynamicOps;
import com.mojang.serialization.MapCodec;
import com.mojang.serialization.MapLike;
import com.mojang.serialization.RecordBuilder;
import it.unimi.dsi.fastutil.objects.ObjectArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.function.Supplier;
import java.util.stream.Stream;
import net.minecraft.core.registries.BuiltInRegistries;
import net.minecraft.server.level.ServerLevel;
import net.minecraft.util.VisibleForDebug;
import net.minecraft.world.attribute.EnvironmentAttribute;
import net.minecraft.world.attribute.EnvironmentAttributeSystem;
import net.minecraft.world.entity.Entity;
import net.minecraft.world.entity.LivingEntity;
import net.minecraft.world.entity.ai.behavior.Behavior;
import net.minecraft.world.entity.ai.behavior.BehaviorControl;
import net.minecraft.world.entity.ai.memory.ExpirableValue;
import net.minecraft.world.entity.ai.memory.MemoryModuleType;
import net.minecraft.world.entity.ai.memory.MemoryStatus;
import net.minecraft.world.entity.ai.sensing.Sensor;
import net.minecraft.world.entity.ai.sensing.SensorType;
import net.minecraft.world.entity.schedule.Activity;
import net.minecraft.world.phys.Vec3;
import org.apache.commons.lang3.mutable.MutableObject;
import org.jspecify.annotations.Nullable;
import org.slf4j.Logger;

public class Brain<E extends LivingEntity> {
    private static final Logger LOGGER = LogUtils.getLogger();
    private final Supplier<Codec<Brain<E>>> codec;
    private static final int SCHEDULE_UPDATE_DELAY = 20;
    private final Map<MemoryModuleType<?>, Optional<? extends ExpirableValue<?>>> memories = Maps.newHashMap();
    private final Map<SensorType<? extends Sensor<? super E>>, Sensor<? super E>> sensors = Maps.newLinkedHashMap();
    private final Map<Integer, Map<Activity, Set<BehaviorControl<? super E>>>> availableBehaviorsByPriority = Maps.newTreeMap();
    private @Nullable EnvironmentAttribute<Activity> schedule;
    private final Map<Activity, Set<Pair<MemoryModuleType<?>, MemoryStatus>>> activityRequirements = Maps.newHashMap();
    private final Map<Activity, Set<MemoryModuleType<?>>> activityMemoriesToEraseWhenStopped = Maps.newHashMap();
    private Set<Activity> coreActivities = Sets.newHashSet();
    private final Set<Activity> activeActivities = Sets.newHashSet();
    private Activity defaultActivity = Activity.IDLE;
    private long lastScheduleUpdate = -9999L;

    public static <E extends LivingEntity> Provider<E> provider(Collection<? extends MemoryModuleType<?>> memoryTypes, Collection<? extends SensorType<? extends Sensor<? super E>>> sensorTypes) {
        return new Provider(memoryTypes, sensorTypes);
    }

    public static <E extends LivingEntity> Codec<Brain<E>> codec(final Collection<? extends MemoryModuleType<?>> memoryTypes, final Collection<? extends SensorType<? extends Sensor<? super E>>> sensorTypes) {
        final MutableObject codecReference = new MutableObject();
        codecReference.setValue((Object)new MapCodec<Brain<E>>(){

            public <T> Stream<T> keys(DynamicOps<T> ops) {
                return memoryTypes.stream().flatMap(t -> t.getCodec().map(c -> BuiltInRegistries.MEMORY_MODULE_TYPE.getKey((MemoryModuleType<?>)t)).stream()).map(l -> ops.createString(l.toString()));
            }

            public <T> DataResult<Brain<E>> decode(DynamicOps<T> ops, MapLike<T> input) {
                MutableObject result = new MutableObject((Object)DataResult.success((Object)ImmutableList.builder()));
                input.entries().forEach(pair -> {
                    DataResult typeResult = BuiltInRegistries.MEMORY_MODULE_TYPE.byNameCodec().parse(ops, pair.getFirst());
                    DataResult entryResult = typeResult.flatMap(type -> this.captureRead((MemoryModuleType)type, ops, (Object)pair.getSecond()));
                    result.setValue((Object)((DataResult)result.get()).apply2(ImmutableList.Builder::add, entryResult));
                });
                ImmutableList memories = ((DataResult)result.get()).resultOrPartial(arg_0 -> ((Logger)LOGGER).error(arg_0)).map(ImmutableList.Builder::build).orElseGet(ImmutableList::of);
                return DataResult.success(new Brain(memoryTypes, sensorTypes, (ImmutableList<MemoryValue<?>>)memories, codecReference));
            }

            private <T, U> DataResult<MemoryValue<U>> captureRead(MemoryModuleType<U> type, DynamicOps<T> ops, T input) {
                return type.getCodec().map(DataResult::success).orElseGet(() -> DataResult.error(() -> "No codec for memory: " + String.valueOf(type))).flatMap(c -> c.parse(ops, input)).map(v -> new MemoryValue(type, Optional.of(v)));
            }

            public <T> RecordBuilder<T> encode(Brain<E> input, DynamicOps<T> ops, RecordBuilder<T> prefix) {
                input.memories().forEach(m -> m.serialize(ops, prefix));
                return prefix;
            }
        }.fieldOf("memories").codec());
        return (Codec)codecReference.get();
    }

    public Brain(Collection<? extends MemoryModuleType<?>> memoryTypes, Collection<? extends SensorType<? extends Sensor<? super E>>> sensorTypes, ImmutableList<MemoryValue<?>> memories, Supplier<Codec<Brain<E>>> codec) {
        this.codec = codec;
        for (MemoryModuleType<?> memoryModuleType : memoryTypes) {
            this.memories.put(memoryModuleType, Optional.empty());
        }
        for (SensorType sensorType : sensorTypes) {
            this.sensors.put(sensorType, (Sensor<E>)sensorType.create());
        }
        for (Sensor sensor : this.sensors.values()) {
            for (MemoryModuleType<?> type : sensor.requires()) {
                this.memories.put(type, Optional.empty());
            }
        }
        for (MemoryValue memoryValue : memories) {
            memoryValue.setMemoryInternal(this);
        }
    }

    public <T> DataResult<T> serializeStart(DynamicOps<T> ops) {
        return this.codec.get().encodeStart(ops, (Object)this);
    }

    private Stream<MemoryValue<?>> memories() {
        return this.memories.entrySet().stream().map(e -> MemoryValue.createUnchecked((MemoryModuleType)e.getKey(), (Optional)e.getValue()));
    }

    public boolean hasMemoryValue(MemoryModuleType<?> type) {
        return this.checkMemory(type, MemoryStatus.VALUE_PRESENT);
    }

    public void clearMemories() {
        this.memories.keySet().forEach(key -> this.memories.put((MemoryModuleType<?>)key, Optional.empty()));
    }

    public <U> void eraseMemory(MemoryModuleType<U> type) {
        this.setMemory(type, Optional.empty());
    }

    public <U> void setMemory(MemoryModuleType<U> type, @Nullable U value) {
        this.setMemory(type, Optional.ofNullable(value));
    }

    public <U> void setMemoryWithExpiry(MemoryModuleType<U> type, U value, long timeToLive) {
        this.setMemoryInternal(type, Optional.of(ExpirableValue.of(value, timeToLive)));
    }

    public <U> void setMemory(MemoryModuleType<U> type, Optional<? extends U> optionalValue) {
        this.setMemoryInternal(type, optionalValue.map(ExpirableValue::of));
    }

    private <U> void setMemoryInternal(MemoryModuleType<U> type, Optional<? extends ExpirableValue<?>> optionalExpirableValue) {
        if (this.memories.containsKey(type)) {
            if (optionalExpirableValue.isPresent() && this.isEmptyCollection(optionalExpirableValue.get().getValue())) {
                this.eraseMemory(type);
            } else {
                this.memories.put(type, optionalExpirableValue);
            }
        }
    }

    public <U> Optional<U> getMemory(MemoryModuleType<U> type) {
        Optional<ExpirableValue<?>> expirableValue = this.memories.get(type);
        if (expirableValue == null) {
            throw new IllegalStateException("Unregistered memory fetched: " + String.valueOf(type));
        }
        return expirableValue.map(ExpirableValue::getValue);
    }

    public <U> @Nullable Optional<U> getMemoryInternal(MemoryModuleType<U> type) {
        Optional<ExpirableValue<?>> expirableValue = this.memories.get(type);
        if (expirableValue == null) {
            return null;
        }
        return expirableValue.map(ExpirableValue::getValue);
    }

    public <U> long getTimeUntilExpiry(MemoryModuleType<U> type) {
        Optional<ExpirableValue<?>> memory = this.memories.get(type);
        return memory.map(ExpirableValue::getTimeToLive).orElse(0L);
    }

    @Deprecated
    @VisibleForDebug
    public Map<MemoryModuleType<?>, Optional<? extends ExpirableValue<?>>> getMemories() {
        return this.memories;
    }

    public <U> boolean isMemoryValue(MemoryModuleType<U> memoryType, U value) {
        if (!this.hasMemoryValue(memoryType)) {
            return false;
        }
        return this.getMemory(memoryType).filter(memory -> memory.equals(value)).isPresent();
    }

    public boolean checkMemory(MemoryModuleType<?> type, MemoryStatus status) {
        Optional<ExpirableValue<?>> optionalExpirableValue = this.memories.get(type);
        if (optionalExpirableValue == null) {
            return false;
        }
        return status == MemoryStatus.REGISTERED || status == MemoryStatus.VALUE_PRESENT && optionalExpirableValue.isPresent() || status == MemoryStatus.VALUE_ABSENT && optionalExpirableValue.isEmpty();
    }

    public void setSchedule(EnvironmentAttribute<Activity> schedule) {
        this.schedule = schedule;
    }

    public void setCoreActivities(Set<Activity> activities) {
        this.coreActivities = activities;
    }

    @Deprecated
    @VisibleForDebug
    public Set<Activity> getActiveActivities() {
        return this.activeActivities;
    }

    @Deprecated
    @VisibleForDebug
    public List<BehaviorControl<? super E>> getRunningBehaviors() {
        ObjectArrayList runningBehaviours = new ObjectArrayList();
        for (Map<Activity, Set<BehaviorControl<E>>> behavioursByActivities : this.availableBehaviorsByPriority.values()) {
            for (Set<BehaviorControl<E>> behaviors : behavioursByActivities.values()) {
                for (BehaviorControl<E> behavior : behaviors) {
                    if (behavior.getStatus() != Behavior.Status.RUNNING) continue;
                    runningBehaviours.add(behavior);
                }
            }
        }
        return runningBehaviours;
    }

    public void useDefaultActivity() {
        this.setActiveActivity(this.defaultActivity);
    }

    public Optional<Activity> getActiveNonCoreActivity() {
        for (Activity activity : this.activeActivities) {
            if (this.coreActivities.contains(activity)) continue;
            return Optional.of(activity);
        }
        return Optional.empty();
    }

    public void setActiveActivityIfPossible(Activity activity) {
        if (this.activityRequirementsAreMet(activity)) {
            this.setActiveActivity(activity);
        } else {
            this.useDefaultActivity();
        }
    }

    private void setActiveActivity(Activity activity) {
        if (this.isActive(activity)) {
            return;
        }
        this.eraseMemoriesForOtherActivitesThan(activity);
        this.activeActivities.clear();
        this.activeActivities.addAll(this.coreActivities);
        this.activeActivities.add(activity);
    }

    private void eraseMemoriesForOtherActivitesThan(Activity activity) {
        for (Activity oldActivity : this.activeActivities) {
            Set<MemoryModuleType<?>> memoryModuleTypes;
            if (oldActivity == activity || (memoryModuleTypes = this.activityMemoriesToEraseWhenStopped.get(oldActivity)) == null) continue;
            for (MemoryModuleType<?> memoryModuleType : memoryModuleTypes) {
                this.eraseMemory(memoryModuleType);
            }
        }
    }

    public void updateActivityFromSchedule(EnvironmentAttributeSystem environmentAttributes, long gameTime, Vec3 pos) {
        if (gameTime - this.lastScheduleUpdate > 20L) {
            Activity scheduledActivity;
            this.lastScheduleUpdate = gameTime;
            Activity activity = scheduledActivity = this.schedule != null ? environmentAttributes.getValue(this.schedule, pos) : Activity.IDLE;
            if (!this.activeActivities.contains(scheduledActivity)) {
                this.setActiveActivityIfPossible(scheduledActivity);
            }
        }
    }

    public void setActiveActivityToFirstValid(List<Activity> activities) {
        for (Activity activity : activities) {
            if (!this.activityRequirementsAreMet(activity)) continue;
            this.setActiveActivity(activity);
            break;
        }
    }

    public void setDefaultActivity(Activity activity) {
        this.defaultActivity = activity;
    }

    public void addActivity(Activity activity, int priorityOfFirstBehavior, ImmutableList<? extends BehaviorControl<? super E>> behaviorList) {
        this.addActivity(activity, this.createPriorityPairs(priorityOfFirstBehavior, behaviorList));
    }

    public void addActivityAndRemoveMemoryWhenStopped(Activity activity, int priorityOfFirstBehavior, ImmutableList<? extends BehaviorControl<? super E>> behaviorList, MemoryModuleType<?> memoryThatMustHaveValueAndWillBeErasedAfter) {
        ImmutableSet conditions = ImmutableSet.of((Object)Pair.of(memoryThatMustHaveValueAndWillBeErasedAfter, (Object)((Object)MemoryStatus.VALUE_PRESENT)));
        ImmutableSet memoriesToEraseWhenStopped = ImmutableSet.of(memoryThatMustHaveValueAndWillBeErasedAfter);
        this.addActivityAndRemoveMemoriesWhenStopped(activity, (ImmutableList<? extends Pair<Integer, ? extends BehaviorControl<? super E>>>)this.createPriorityPairs(priorityOfFirstBehavior, behaviorList), (Set<Pair<MemoryModuleType<?>, MemoryStatus>>)conditions, (Set<MemoryModuleType<?>>)memoriesToEraseWhenStopped);
    }

    public void addActivity(Activity activity, ImmutableList<? extends Pair<Integer, ? extends BehaviorControl<? super E>>> behaviorPriorityPairs) {
        this.addActivityAndRemoveMemoriesWhenStopped(activity, behaviorPriorityPairs, (Set<Pair<MemoryModuleType<?>, MemoryStatus>>)ImmutableSet.of(), Sets.newHashSet());
    }

    public void addActivityWithConditions(Activity activity, int priorityOfFirstBehavior, ImmutableList<? extends BehaviorControl<? super E>> behaviorList, Set<Pair<MemoryModuleType<?>, MemoryStatus>> conditions) {
        this.addActivityWithConditions(activity, this.createPriorityPairs(priorityOfFirstBehavior, behaviorList), conditions);
    }

    public void addActivityWithConditions(Activity activity, ImmutableList<? extends Pair<Integer, ? extends BehaviorControl<? super E>>> behaviorPriorityPairs, Set<Pair<MemoryModuleType<?>, MemoryStatus>> conditions) {
        this.addActivityAndRemoveMemoriesWhenStopped(activity, behaviorPriorityPairs, conditions, Sets.newHashSet());
    }

    public void addActivityAndRemoveMemoriesWhenStopped(Activity activity, ImmutableList<? extends Pair<Integer, ? extends BehaviorControl<? super E>>> behaviorPriorityPairs, Set<Pair<MemoryModuleType<?>, MemoryStatus>> conditions, Set<MemoryModuleType<?>> memoriesToEraseWhenStopped) {
        this.activityRequirements.put(activity, conditions);
        if (!memoriesToEraseWhenStopped.isEmpty()) {
            this.activityMemoriesToEraseWhenStopped.put(activity, memoriesToEraseWhenStopped);
        }
        for (Pair pair : behaviorPriorityPairs) {
            this.availableBehaviorsByPriority.computeIfAbsent((Integer)pair.getFirst(), key -> Maps.newHashMap()).computeIfAbsent(activity, key -> Sets.newLinkedHashSet()).add((BehaviorControl)pair.getSecond());
        }
    }

    @VisibleForTesting
    public void removeAllBehaviors() {
        this.availableBehaviorsByPriority.clear();
    }

    public boolean isActive(Activity activity) {
        return this.activeActivities.contains(activity);
    }

    public Brain<E> copyWithoutBehaviors() {
        Brain<E> brain = new Brain<E>(this.memories.keySet(), this.sensors.keySet(), ImmutableList.of(), this.codec);
        for (Map.Entry<MemoryModuleType<?>, Optional<ExpirableValue<?>>> memoryEntry : this.memories.entrySet()) {
            MemoryModuleType<?> memoryModuleType = memoryEntry.getKey();
            if (!memoryEntry.getValue().isPresent()) continue;
            brain.memories.put(memoryModuleType, memoryEntry.getValue());
        }
        return brain;
    }

    public void tick(ServerLevel level, E body) {
        this.forgetOutdatedMemories();
        this.tickSensors(level, body);
        this.startEachNonRunningBehavior(level, body);
        this.tickEachRunningBehavior(level, body);
    }

    private void tickSensors(ServerLevel level, E body) {
        for (Sensor<E> sensor : this.sensors.values()) {
            sensor.tick(level, body);
        }
    }

    private void forgetOutdatedMemories() {
        for (Map.Entry<MemoryModuleType<?>, Optional<ExpirableValue<?>>> entry : this.memories.entrySet()) {
            if (!entry.getValue().isPresent()) continue;
            ExpirableValue<?> memory = entry.getValue().get();
            if (memory.hasExpired()) {
                this.eraseMemory(entry.getKey());
            }
            memory.tick();
        }
    }

    public void stopAll(ServerLevel level, E body) {
        long timestamp = ((Entity)body).level().getGameTime();
        for (BehaviorControl<E> behavior : this.getRunningBehaviors()) {
            behavior.doStop(level, body, timestamp);
        }
    }

    private void startEachNonRunningBehavior(ServerLevel level, E body) {
        long time = level.getGameTime();
        for (Map<Activity, Set<BehaviorControl<E>>> behavioursByActivities : this.availableBehaviorsByPriority.values()) {
            for (Map.Entry<Activity, Set<BehaviorControl<E>>> behavioursForActivity : behavioursByActivities.entrySet()) {
                Activity activity = behavioursForActivity.getKey();
                if (!this.activeActivities.contains(activity)) continue;
                Set<BehaviorControl<E>> behaviors = behavioursForActivity.getValue();
                for (BehaviorControl<E> behavior : behaviors) {
                    if (behavior.getStatus() != Behavior.Status.STOPPED) continue;
                    behavior.tryStart(level, body, time);
                }
            }
        }
    }

    private void tickEachRunningBehavior(ServerLevel level, E body) {
        long timestamp = level.getGameTime();
        for (BehaviorControl<E> behavior : this.getRunningBehaviors()) {
            behavior.tickOrStop(level, body, timestamp);
        }
    }

    private boolean activityRequirementsAreMet(Activity activity) {
        if (!this.activityRequirements.containsKey(activity)) {
            return false;
        }
        for (Pair<MemoryModuleType<?>, MemoryStatus> memoryRequirement : this.activityRequirements.get(activity)) {
            MemoryStatus memoryStatus;
            MemoryModuleType memoryType = (MemoryModuleType)memoryRequirement.getFirst();
            if (this.checkMemory(memoryType, memoryStatus = (MemoryStatus)((Object)memoryRequirement.getSecond()))) continue;
            return false;
        }
        return true;
    }

    private boolean isEmptyCollection(Object object) {
        return object instanceof Collection && ((Collection)object).isEmpty();
    }

    ImmutableList<? extends Pair<Integer, ? extends BehaviorControl<? super E>>> createPriorityPairs(int priorityOfFirstBehavior, ImmutableList<? extends BehaviorControl<? super E>> behaviorList) {
        int nextPrio = priorityOfFirstBehavior;
        ImmutableList.Builder listBuilder = ImmutableList.builder();
        for (BehaviorControl behavior : behaviorList) {
            listBuilder.add((Object)Pair.of((Object)nextPrio++, (Object)behavior));
        }
        return listBuilder.build();
    }

    public boolean isBrainDead() {
        return this.memories.isEmpty() && this.sensors.isEmpty() && this.availableBehaviorsByPriority.isEmpty();
    }

    public static final class Provider<E extends LivingEntity> {
        private final Collection<? extends MemoryModuleType<?>> memoryTypes;
        private final Collection<? extends SensorType<? extends Sensor<? super E>>> sensorTypes;
        private final Codec<Brain<E>> codec;

        private Provider(Collection<? extends MemoryModuleType<?>> memoryTypes, Collection<? extends SensorType<? extends Sensor<? super E>>> sensorTypes) {
            this.memoryTypes = memoryTypes;
            this.sensorTypes = sensorTypes;
            this.codec = Brain.codec(memoryTypes, sensorTypes);
        }

        public Brain<E> makeBrain(Dynamic<?> input) {
            return this.codec.parse(input).resultOrPartial(arg_0 -> ((Logger)LOGGER).error(arg_0)).orElseGet(() -> new Brain(this.memoryTypes, this.sensorTypes, ImmutableList.of(), () -> this.codec));
        }
    }

    private static final class MemoryValue<U> {
        private final MemoryModuleType<U> type;
        private final Optional<? extends ExpirableValue<U>> value;

        private static <U> MemoryValue<U> createUnchecked(MemoryModuleType<U> type, Optional<? extends ExpirableValue<?>> value) {
            return new MemoryValue<U>(type, value);
        }

        private MemoryValue(MemoryModuleType<U> type, Optional<? extends ExpirableValue<U>> value) {
            this.type = type;
            this.value = value;
        }

        private void setMemoryInternal(Brain<?> brain) {
            brain.setMemoryInternal(this.type, this.value);
        }

        public <T> void serialize(DynamicOps<T> ops, RecordBuilder<T> builder) {
            this.type.getCodec().ifPresent(codec -> this.value.ifPresent(v -> builder.add(BuiltInRegistries.MEMORY_MODULE_TYPE.byNameCodec().encodeStart(ops, this.type), codec.encodeStart(ops, v))));
        }
    }
}