package net.minecraft.world.entity.ai;

import com.mojang.serialization.Dynamic;
import org.apache.logging.log4j.LogManager;
import net.minecraft.server.level.ServerLevel;
import com.google.common.collect.ImmutableSet;
import it.unimi.dsi.fastutil.objects.ObjectArrayList;
import java.util.List;
import javax.annotation.Nullable;
import com.google.common.collect.UnmodifiableIterator;
import java.util.Iterator;
import com.google.common.collect.Sets;
import com.google.common.collect.Maps;
import net.minecraft.resources.ResourceLocation;
import com.mojang.serialization.RecordBuilder;
import java.util.function.Function;
import java.util.function.Consumer;
import java.util.function.BiFunction;
import com.google.common.collect.ImmutableList;
import com.mojang.serialization.DataResult;
import com.mojang.serialization.MapLike;
import net.minecraft.Util;
import net.minecraft.core.Registry;
import java.util.stream.Stream;
import com.mojang.serialization.DynamicOps;
import com.mojang.serialization.MapCodec;
import org.apache.commons.lang3.mutable.MutableObject;
import java.util.Collection;
import net.minecraft.world.entity.ai.memory.MemoryStatus;
import com.mojang.datafixers.util.Pair;
import net.minecraft.world.entity.schedule.Schedule;
import net.minecraft.world.entity.ai.behavior.Behavior;
import java.util.Set;
import net.minecraft.world.entity.schedule.Activity;
import net.minecraft.world.entity.ai.sensing.Sensor;
import net.minecraft.world.entity.ai.sensing.SensorType;
import net.minecraft.world.entity.ai.memory.ExpirableValue;
import java.util.Optional;
import net.minecraft.world.entity.ai.memory.MemoryModuleType;
import java.util.Map;
import com.mojang.serialization.Codec;
import java.util.function.Supplier;
import org.apache.logging.log4j.Logger;
import net.minecraft.world.entity.LivingEntity;

public class Brain<E extends LivingEntity> {
    private static final Logger LOGGER;
    private final Supplier<Codec<Brain<E>>> codec;
    private final Map<MemoryModuleType<?>, Optional<? extends ExpirableValue<?>>> memories;
    private final Map<SensorType<? extends Sensor<? super E>>, Sensor<? super E>> sensors;
    private final Map<Integer, Map<Activity, Set<Behavior<? super E>>>> availableBehaviorsByPriority;
    private Schedule schedule;
    private final Map<Activity, Set<Pair<MemoryModuleType<?>, MemoryStatus>>> activityRequirements;
    private final Map<Activity, Set<MemoryModuleType<?>>> activityMemoriesToEraseWhenStopped;
    private Set<Activity> coreActivities;
    private final Set<Activity> activeActivities;
    private Activity defaultActivity;
    private long lastScheduleUpdate;
    
    public static <E extends LivingEntity> Provider<E> provider(final Collection<? extends MemoryModuleType<?>> collection1, final Collection<? extends SensorType<? extends Sensor<? super E>>> collection2) {
        return new Provider<E>((Collection)collection1, (Collection)collection2);
    }
    
    public static <E extends LivingEntity> Codec<Brain<E>> codec(final Collection<? extends MemoryModuleType<?>> collection1, final Collection<? extends SensorType<? extends Sensor<? super E>>> collection2) {
        final MutableObject<Codec<Brain<E>>> mutableObject3 = (MutableObject<Codec<Brain<E>>>)new MutableObject();
        mutableObject3.setValue(new MapCodec<Brain<E>>() {
            public <T> Stream<T> keys(final DynamicOps<T> dynamicOps) {
                return collection1.stream().flatMap(awp -> Util.toStream(awp.getCodec().map(codec -> Registry.MEMORY_MODULE_TYPE.getKey(awp)))).<T>map(uh -> dynamicOps.createString(uh.toString()));
            }
            
            public <T> DataResult<Brain<E>> decode(final DynamicOps<T> dynamicOps, final MapLike<T> mapLike) {
                final MutableObject<DataResult<ImmutableList.Builder<MemoryValue<?>>>> mutableObject4 = (MutableObject<DataResult<ImmutableList.Builder<MemoryValue<?>>>>)new MutableObject((Object)DataResult.success((Object)ImmutableList.builder()));
                final DataResult<MemoryModuleType<?>> dataResult5;
                final DataResult<? extends MemoryValue<?>> dataResult6;
                final MutableObject mutableObject5;
                mapLike.entries().forEach(pair -> {
                    dataResult5 = (DataResult<MemoryModuleType<?>>)Registry.MEMORY_MODULE_TYPE.parse((DynamicOps)dynamicOps, pair.getFirst());
                    dataResult6 = dataResult5.flatMap(awp -> this.captureRead(awp, (com.mojang.serialization.DynamicOps<Object>)dynamicOps, pair.getSecond()));
                    mutableObject5.setValue((Object)((DataResult)mutableObject5.getValue()).apply2((BiFunction)ImmutableList.Builder::add, (DataResult)dataResult6));
                    return;
                });
                final ImmutableList<MemoryValue<?>> immutableList5 = ((DataResult)mutableObject4.getValue()).resultOrPartial((Consumer)Brain.LOGGER::error).<ImmutableList<MemoryValue<?>>>map(ImmutableList.Builder::build).orElseGet(ImmutableList::of);
                return (DataResult<Brain<E>>)DataResult.success((Object)new Brain(collection1, collection2, immutableList5, mutableObject3::getValue));
            }
            
            private <T, U> DataResult<MemoryValue<U>> captureRead(final MemoryModuleType<U> awp, final DynamicOps<T> dynamicOps, final T object) {
                return (DataResult<MemoryValue<U>>)awp.getCodec().<DataResult>map(DataResult::success).orElseGet(() -> DataResult.error("No codec for memory: " + awp)).flatMap(codec -> codec.parse((DynamicOps)dynamicOps, (Object)object)).map(awo -> new MemoryValue((MemoryModuleType)awp, (Optional)Optional.<ExpirableValue>of(awo)));
            }
            
            public <T> RecordBuilder<T> encode(final Brain<E> apr, final DynamicOps<T> dynamicOps, final RecordBuilder<T> recordBuilder) {
                ((Brain<LivingEntity>)apr).memories().forEach(a -> a.<T>serialize(dynamicOps, recordBuilder));
                return recordBuilder;
            }
        }.fieldOf("memories").codec());
        return (Codec<Brain<E>>)mutableObject3.getValue();
    }
    
    public Brain(final Collection<? extends MemoryModuleType<?>> collection1, final Collection<? extends SensorType<? extends Sensor<? super E>>> collection2, final ImmutableList<MemoryValue<?>> immutableList, final Supplier<Codec<Brain<E>>> supplier) {
        this.memories = Maps.newHashMap();
        this.sensors = Maps.newLinkedHashMap();
        this.availableBehaviorsByPriority = Maps.newTreeMap();
        this.schedule = Schedule.EMPTY;
        this.activityRequirements = Maps.newHashMap();
        this.activityMemoriesToEraseWhenStopped = Maps.newHashMap();
        this.coreActivities = Sets.newHashSet();
        this.activeActivities = Sets.newHashSet();
        this.defaultActivity = Activity.IDLE;
        this.lastScheduleUpdate = -9999L;
        this.codec = supplier;
        for (final MemoryModuleType<?> awp7 : collection1) {
            this.memories.put(awp7, Optional.empty());
        }
        for (final SensorType<? extends Sensor<? super E>> axo7 : collection2) {
            this.sensors.put(axo7, axo7.create());
        }
        for (final Sensor<? super E> axn7 : this.sensors.values()) {
            for (final MemoryModuleType<?> awp8 : axn7.requires()) {
                this.memories.put(awp8, Optional.empty());
            }
        }
        for (final MemoryValue<?> a7 : immutableList) {
            ((MemoryValue<Object>)a7).setMemoryInternal(this);
        }
    }
    
    public <T> DataResult<T> serializeStart(final DynamicOps<T> dynamicOps) {
        return (DataResult<T>)this.codec.get().encodeStart((DynamicOps)dynamicOps, this);
    }
    
    private Stream<MemoryValue<?>> memories() {
        return this.memories.entrySet().stream().<MemoryValue<?>>map(entry -> MemoryValue.createUnchecked((MemoryModuleType<Object>)entry.getKey(), (Optional)entry.getValue()));
    }
    
    public boolean hasMemoryValue(final MemoryModuleType<?> awp) {
        return this.checkMemory(awp, MemoryStatus.VALUE_PRESENT);
    }
    
    public <U> void eraseMemory(final MemoryModuleType<U> awp) {
        this.<U>setMemory(awp, Optional.empty());
    }
    
    public <U> void setMemory(final MemoryModuleType<U> awp, @Nullable final U object) {
        this.<U>setMemory(awp, Optional.ofNullable((U)object));
    }
    
    public <U> void setMemoryWithExpiry(final MemoryModuleType<U> awp, final U object, final long long3) {
        this.<U>setMemoryInternal(awp, Optional.of(ExpirableValue.<U>of(object, long3)));
    }
    
    public <U> void setMemory(final MemoryModuleType<U> awp, final Optional<? extends U> optional) {
        this.setMemoryInternal((MemoryModuleType<Object>)awp, optional.map(ExpirableValue::of));
    }
    
    private <U> void setMemoryInternal(final MemoryModuleType<U> awp, final Optional<? extends ExpirableValue<?>> optional) {
        if (this.memories.containsKey(awp)) {
            if (optional.isPresent() && this.isEmptyCollection(((ExpirableValue)optional.get()).getValue())) {
                this.<U>eraseMemory(awp);
            }
            else {
                this.memories.put(awp, optional);
            }
        }
    }
    
    public <U> Optional<U> getMemory(final MemoryModuleType<U> awp) {
        return this.memories.get(awp).<U>map(ExpirableValue::getValue);
    }
    
    public <U> boolean isMemoryValue(final MemoryModuleType<U> awp, final U object) {
        return this.hasMemoryValue(awp) && this.<U>getMemory(awp).filter(object2 -> object2.equals(object)).isPresent();
    }
    
    public boolean checkMemory(final MemoryModuleType<?> awp, final MemoryStatus awq) {
        final Optional<? extends ExpirableValue<?>> optional4 = this.memories.get(awp);
        return optional4 != null && (awq == MemoryStatus.REGISTERED || (awq == MemoryStatus.VALUE_PRESENT && optional4.isPresent()) || (awq == MemoryStatus.VALUE_ABSENT && !optional4.isPresent()));
    }
    
    public Schedule getSchedule() {
        return this.schedule;
    }
    
    public void setSchedule(final Schedule bfn) {
        this.schedule = bfn;
    }
    
    public void setCoreActivities(final Set<Activity> set) {
        this.coreActivities = set;
    }
    
    @Deprecated
    public List<Behavior<? super E>> getRunningBehaviors() {
        final List<Behavior<? super E>> list2 = (List<Behavior<? super E>>)new ObjectArrayList();
        for (final Map<Activity, Set<Behavior<? super E>>> map4 : this.availableBehaviorsByPriority.values()) {
            for (final Set<Behavior<? super E>> set6 : map4.values()) {
                for (final Behavior<? super E> aqh8 : set6) {
                    if (aqh8.getStatus() == Behavior.Status.RUNNING) {
                        list2.add(aqh8);
                    }
                }
            }
        }
        return list2;
    }
    
    public void useDefaultActivity() {
        this.setActiveActivity(this.defaultActivity);
    }
    
    public Optional<Activity> getActiveNonCoreActivity() {
        for (final Activity bfl3 : this.activeActivities) {
            if (!this.coreActivities.contains(bfl3)) {
                return Optional.<Activity>of(bfl3);
            }
        }
        return Optional.<Activity>empty();
    }
    
    public void setActiveActivityIfPossible(final Activity bfl) {
        if (this.activityRequirementsAreMet(bfl)) {
            this.setActiveActivity(bfl);
        }
        else {
            this.useDefaultActivity();
        }
    }
    
    private void setActiveActivity(final Activity bfl) {
        if (this.isActive(bfl)) {
            return;
        }
        this.eraseMemoriesForOtherActivitesThan(bfl);
        this.activeActivities.clear();
        this.activeActivities.addAll(this.coreActivities);
        this.activeActivities.add(bfl);
    }
    
    private void eraseMemoriesForOtherActivitesThan(final Activity bfl) {
        for (final Activity bfl2 : this.activeActivities) {
            if (bfl2 != bfl) {
                final Set<MemoryModuleType<?>> set5 = this.activityMemoriesToEraseWhenStopped.get(bfl2);
                if (set5 == null) {
                    continue;
                }
                for (final MemoryModuleType<?> awp7 : set5) {
                    this.eraseMemory(awp7);
                }
            }
        }
    }
    
    public void updateActivityFromSchedule(final long long1, final long long2) {
        if (long2 - this.lastScheduleUpdate > 20L) {
            this.lastScheduleUpdate = long2;
            final Activity bfl6 = this.getSchedule().getActivityAt((int)(long1 % 24000L));
            if (!this.activeActivities.contains(bfl6)) {
                this.setActiveActivityIfPossible(bfl6);
            }
        }
    }
    
    public void setActiveActivityToFirstValid(final List<Activity> list) {
        for (final Activity bfl4 : list) {
            if (this.activityRequirementsAreMet(bfl4)) {
                this.setActiveActivity(bfl4);
                break;
            }
        }
    }
    
    public void setDefaultActivity(final Activity bfl) {
        this.defaultActivity = bfl;
    }
    
    public void addActivity(final Activity bfl, final int integer, final ImmutableList<? extends Behavior<? super E>> immutableList) {
        this.addActivity(bfl, this.createPriorityPairs(integer, immutableList));
    }
    
    public void addActivityAndRemoveMemoryWhenStopped(final Activity bfl, final int integer, final ImmutableList<? extends Behavior<? super E>> immutableList, final MemoryModuleType<?> awp) {
        final Set<Pair<MemoryModuleType<?>, MemoryStatus>> set6 = ImmutableSet.of(Pair.of(awp, MemoryStatus.VALUE_PRESENT));
        final Set<MemoryModuleType<?>> set7 = ImmutableSet.<MemoryModuleType<?>>of(awp);
        this.addActivityAndRemoveMemoriesWhenStopped(bfl, this.createPriorityPairs(integer, immutableList), set6, set7);
    }
    
    public void addActivity(final Activity bfl, final ImmutableList<? extends Pair<Integer, ? extends Behavior<? super E>>> immutableList) {
        this.addActivityAndRemoveMemoriesWhenStopped(bfl, immutableList, ImmutableSet.of(), Sets.newHashSet());
    }
    
    public void addActivityWithConditions(final Activity bfl, final ImmutableList<? extends Pair<Integer, ? extends Behavior<? super E>>> immutableList, final Set<Pair<MemoryModuleType<?>, MemoryStatus>> set) {
        this.addActivityAndRemoveMemoriesWhenStopped(bfl, immutableList, set, Sets.newHashSet());
    }
    
    private void addActivityAndRemoveMemoriesWhenStopped(final Activity bfl, final ImmutableList<? extends Pair<Integer, ? extends Behavior<? super E>>> immutableList, final Set<Pair<MemoryModuleType<?>, MemoryStatus>> set3, final Set<MemoryModuleType<?>> set4) {
        this.activityRequirements.put(bfl, set3);
        if (!set4.isEmpty()) {
            this.activityMemoriesToEraseWhenStopped.put(bfl, set4);
        }
        for (final Pair<Integer, ? extends Behavior<? super E>> pair7 : immutableList) {
            this.availableBehaviorsByPriority.computeIfAbsent((Integer)pair7.getFirst(), integer -> Maps.newHashMap()).computeIfAbsent(bfl, bfl -> Sets.newLinkedHashSet()).add(pair7.getSecond());
        }
    }
    
    public boolean isActive(final Activity bfl) {
        return this.activeActivities.contains(bfl);
    }
    
    public Brain<E> copyWithoutBehaviors() {
        final Brain<E> apr2 = new Brain<E>(this.memories.keySet(), this.sensors.keySet(), ImmutableList.<MemoryValue<?>>of(), this.codec);
        for (final Map.Entry<MemoryModuleType<?>, Optional<? extends ExpirableValue<?>>> entry4 : this.memories.entrySet()) {
            final MemoryModuleType<?> awp5 = entry4.getKey();
            if (entry4.getValue().isPresent()) {
                apr2.memories.put(awp5, entry4.getValue());
            }
        }
        return apr2;
    }
    
    public void tick(final ServerLevel zd, final E aoy) {
        this.forgetOutdatedMemories();
        this.tickSensors(zd, aoy);
        this.startEachNonRunningBehavior(zd, aoy);
        this.tickEachRunningBehavior(zd, aoy);
    }
    
    private void tickSensors(final ServerLevel zd, final E aoy) {
        for (final Sensor<? super E> axn5 : this.sensors.values()) {
            axn5.tick(zd, aoy);
        }
    }
    
    private void forgetOutdatedMemories() {
        for (final Map.Entry<MemoryModuleType<?>, Optional<? extends ExpirableValue<?>>> entry3 : this.memories.entrySet()) {
            if (entry3.getValue().isPresent()) {
                final ExpirableValue<?> awo4 = entry3.getValue().get();
                awo4.tick();
                if (!awo4.hasExpired()) {
                    continue;
                }
                this.eraseMemory(entry3.getKey());
            }
        }
    }
    
    public void stopAll(final ServerLevel zd, final E aoy) {
        final long long4 = aoy.level.getGameTime();
        for (final Behavior<? super E> aqh7 : this.getRunningBehaviors()) {
            aqh7.doStop(zd, aoy, long4);
        }
    }
    
    private void startEachNonRunningBehavior(final ServerLevel zd, final E aoy) {
        final long long4 = zd.getGameTime();
        for (final Map<Activity, Set<Behavior<? super E>>> map7 : this.availableBehaviorsByPriority.values()) {
            for (final Map.Entry<Activity, Set<Behavior<? super E>>> entry9 : map7.entrySet()) {
                final Activity bfl10 = entry9.getKey();
                if (this.activeActivities.contains(bfl10)) {
                    final Set<Behavior<? super E>> set11 = entry9.getValue();
                    for (final Behavior<? super E> aqh13 : set11) {
                        if (aqh13.getStatus() == Behavior.Status.STOPPED) {
                            aqh13.tryStart(zd, aoy, long4);
                        }
                    }
                }
            }
        }
    }
    
    private void tickEachRunningBehavior(final ServerLevel zd, final E aoy) {
        final long long4 = zd.getGameTime();
        for (final Behavior<? super E> aqh7 : this.getRunningBehaviors()) {
            aqh7.tickOrStop(zd, aoy, long4);
        }
    }
    
    private boolean activityRequirementsAreMet(final Activity bfl) {
        if (!this.activityRequirements.containsKey(bfl)) {
            return false;
        }
        for (final Pair<MemoryModuleType<?>, MemoryStatus> pair4 : this.activityRequirements.get(bfl)) {
            final MemoryModuleType<?> awp5 = pair4.getFirst();
            final MemoryStatus awq6 = (MemoryStatus)pair4.getSecond();
            if (!this.checkMemory(awp5, awq6)) {
                return false;
            }
        }
        return true;
    }
    
    private boolean isEmptyCollection(final Object object) {
        return object instanceof Collection && ((Collection)object).isEmpty();
    }
    
    ImmutableList<? extends Pair<Integer, ? extends Behavior<? super E>>> createPriorityPairs(final int integer, final ImmutableList<? extends Behavior<? super E>> immutableList) {
        int integer2 = integer;
        final ImmutableList.Builder<Pair<Integer, ? extends Behavior<? super E>>> builder5 = ImmutableList.<Pair<Integer, ? extends Behavior<? super E>>>builder();
        for (final Behavior<? super E> aqh7 : immutableList) {
            builder5.add(Pair.of((integer2++), aqh7));
        }
        return builder5.build();
    }
    
    static {
        LOGGER = LogManager.getLogger();
    }
    
    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(final Collection<? extends MemoryModuleType<?>> collection1, final Collection<? extends SensorType<? extends Sensor<? super E>>> collection2) {
            this.memoryTypes = collection1;
            this.sensorTypes = collection2;
            this.codec = Brain.<E>codec(collection1, collection2);
        }
        
        public Brain<E> makeBrain(final Dynamic<?> dynamic) {
            return this.codec.parse((Dynamic)dynamic).resultOrPartial((Consumer)Brain.LOGGER::error).orElseGet(() -> new Brain(this.memoryTypes, this.sensorTypes, ImmutableList.<MemoryValue<?>>of(), () -> this.codec));
        }
    }
    
    static final class MemoryValue<U> {
        private final MemoryModuleType<U> type;
        private final Optional<? extends ExpirableValue<U>> value;
        
        private static <U> MemoryValue<U> createUnchecked(final MemoryModuleType<U> awp, final Optional<? extends ExpirableValue<?>> optional) {
            return new MemoryValue<U>(awp, optional);
        }
        
        private MemoryValue(final MemoryModuleType<U> awp, final Optional<? extends ExpirableValue<U>> optional) {
            this.type = awp;
            this.value = optional;
        }
        
        private void setMemoryInternal(final Brain<?> apr) {
            ((Brain<LivingEntity>)apr).setMemoryInternal(this.type, this.value);
        }
        
        public <T> void serialize(final DynamicOps<T> dynamicOps, final RecordBuilder<T> recordBuilder) {
            this.type.getCodec().ifPresent(codec -> this.value.ifPresent(awo -> recordBuilder.add(Registry.MEMORY_MODULE_TYPE.encodeStart((DynamicOps)dynamicOps, this.type), codec.encodeStart((DynamicOps)dynamicOps, awo))));
        }
    }
}
