package net.minecraft.world.entity.ai;

import com.google.common.collect.ImmutableMap;
import java.util.stream.Collectors;
import net.minecraft.server.level.ServerLevel;
import com.mojang.datafixers.types.DynamicOps;
import net.minecraft.nbt.CompoundTag;
import net.minecraft.nbt.NbtOps;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.ImmutableList;
import java.util.stream.Stream;
import javax.annotation.Nullable;
import java.util.function.Supplier;
import java.util.Iterator;
import net.minecraft.resources.ResourceLocation;
import net.minecraft.core.Registry;
import java.util.function.Function;
import com.google.common.collect.Sets;
import com.google.common.collect.Maps;
import com.mojang.datafixers.Dynamic;
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 java.util.Optional;
import net.minecraft.world.entity.ai.memory.MemoryModuleType;
import java.util.Map;
import net.minecraft.util.Serializable;
import net.minecraft.world.entity.LivingEntity;

public class Brain<E extends LivingEntity> implements Serializable {
    private final Map<MemoryModuleType<?>, Optional<?>> memories;
    private final Map<SensorType<? extends Sensor<? super E>>, Sensor<? super E>> sensors;
    private final Map<Integer, Map<Activity, Set<Behavior<? super E>>>> availableGoalsByPriority;
    private Schedule schedule;
    private final Map<Activity, Set<Pair<MemoryModuleType<?>, MemoryStatus>>> activityRequirements;
    private Set<Activity> coreActivities;
    private final Set<Activity> activeActivities;
    private Activity defaultActivity;
    private long lastScheduleUpdate;
    
    public <T> Brain(final Collection<MemoryModuleType<?>> collection1, final Collection<SensorType<? extends Sensor<? super E>>> collection2, final Dynamic<T> dynamic) {
        this.memories = Maps.newHashMap();
        this.sensors = Maps.newLinkedHashMap();
        this.availableGoalsByPriority = Maps.newTreeMap();
        this.schedule = Schedule.EMPTY;
        this.activityRequirements = Maps.newHashMap();
        this.coreActivities = Sets.newHashSet();
        this.activeActivities = Sets.newHashSet();
        this.defaultActivity = Activity.IDLE;
        this.lastScheduleUpdate = -9999L;
        final Optional<?> optional;
        collection1.forEach(arl -> optional = this.memories.put(arl, Optional.empty()));
        final Sensor<? super E> sensor;
        collection2.forEach(asg -> sensor = this.sensors.put(asg, asg.create()));
        final Iterator<MemoryModuleType<?>> iterator;
        MemoryModuleType<?> arl2;
        this.sensors.values().forEach(asf -> {
            asf.requires().iterator();
            while (iterator.hasNext()) {
                arl2 = iterator.next();
                this.memories.put(arl2, Optional.empty());
            }
            return;
        });
        for (final Map.Entry<Dynamic<T>, Dynamic<T>> entry6 : dynamic.get("memories").asMap((Function)Function.identity(), (Function)Function.identity()).entrySet()) {
            this.<T, Object>readMemory(Registry.MEMORY_MODULE_TYPE.get(new ResourceLocation(entry6.getKey().asString(""))), entry6.getValue());
        }
    }
    
    public boolean hasMemoryValue(final MemoryModuleType<?> arl) {
        return this.checkMemory(arl, MemoryStatus.VALUE_PRESENT);
    }
    
    private <T, U> void readMemory(final MemoryModuleType<U> arl, final Dynamic<T> dynamic) {
        this.<U>setMemory(arl, arl.getDeserializer().<Throwable>orElseThrow(RuntimeException::new).apply(dynamic));
    }
    
    public <U> void eraseMemory(final MemoryModuleType<U> arl) {
        this.<U>setMemory(arl, Optional.<U>empty());
    }
    
    public <U> void setMemory(final MemoryModuleType<U> arl, @Nullable final U object) {
        this.<U>setMemory(arl, Optional.ofNullable((U)object));
    }
    
    public <U> void setMemory(final MemoryModuleType<U> arl, final Optional<U> optional) {
        if (this.memories.containsKey(arl)) {
            if (optional.isPresent() && this.isEmptyCollection(optional.get())) {
                this.<U>eraseMemory(arl);
            }
            else {
                this.memories.put(arl, optional);
            }
        }
    }
    
    public <U> Optional<U> getMemory(final MemoryModuleType<U> arl) {
        return (Optional<U>)this.memories.get(arl);
    }
    
    public boolean checkMemory(final MemoryModuleType<?> arl, final MemoryStatus arm) {
        final Optional<?> optional4 = this.memories.get(arl);
        return optional4 != null && (arm == MemoryStatus.REGISTERED || (arm == MemoryStatus.VALUE_PRESENT && optional4.isPresent()) || (arm == MemoryStatus.VALUE_ABSENT && !optional4.isPresent()));
    }
    
    public Schedule getSchedule() {
        return this.schedule;
    }
    
    public void setSchedule(final Schedule azt) {
        this.schedule = azt;
    }
    
    public void setCoreActivities(final Set<Activity> set) {
        this.coreActivities = set;
    }
    
    @Deprecated
    public Stream<Behavior<? super E>> getRunningBehaviorsStream() {
        return this.availableGoalsByPriority.values().stream().flatMap(map -> map.values().stream()).<Behavior<? super E>>flatMap(Collection::stream).filter(amb -> amb.getStatus() == Behavior.Status.RUNNING);
    }
    
    public void setActivity(final Activity azr) {
        this.activeActivities.clear();
        this.activeActivities.addAll(this.coreActivities);
        final boolean boolean3 = this.activityRequirements.keySet().contains(azr) && this.activityRequirementsAreMet(azr);
        this.activeActivities.add(boolean3 ? azr : this.defaultActivity);
    }
    
    public void updateActivity(final long long1, final long long2) {
        if (long2 - this.lastScheduleUpdate > 20L) {
            this.lastScheduleUpdate = long2;
            final Activity azr6 = this.getSchedule().getActivityAt((int)(long1 % 24000L));
            if (!this.activeActivities.contains(azr6)) {
                this.setActivity(azr6);
            }
        }
    }
    
    public void setDefaultActivity(final Activity azr) {
        this.defaultActivity = azr;
    }
    
    public void addActivity(final Activity azr, final ImmutableList<Pair<Integer, ? extends Behavior<? super E>>> immutableList) {
        this.addActivity(azr, immutableList, ImmutableSet.of());
    }
    
    public void addActivity(final Activity azr, final ImmutableList<Pair<Integer, ? extends Behavior<? super E>>> immutableList, final Set<Pair<MemoryModuleType<?>, MemoryStatus>> set) {
        this.activityRequirements.put(azr, set);
        immutableList.forEach(pair -> this.availableGoalsByPriority.computeIfAbsent((Integer)pair.getFirst(), integer -> Maps.newHashMap()).computeIfAbsent(azr, azr -> Sets.newLinkedHashSet()).add(pair.getSecond()));
    }
    
    public boolean isActive(final Activity azr) {
        return this.activeActivities.contains(azr);
    }
    
    public Brain<E> copyWithoutGoals() {
        final Brain<E> alp2 = new Brain<E>(this.memories.keySet(), this.sensors.keySet(), (Dynamic<T>)new Dynamic((DynamicOps)NbtOps.INSTANCE, new CompoundTag()));
        final Optional<?> optional2;
        this.memories.forEach((arl, optional) -> optional.ifPresent(object -> optional2 = alp2.memories.put(arl, Optional.of(object))));
        return alp2;
    }
    
    public void tick(final ServerLevel xd, final E akz) {
        this.tickEachSensor(xd, akz);
        this.startEachNonRunningBehavior(xd, akz);
        this.tickEachRunningBehavior(xd, akz);
    }
    
    public void stopAll(final ServerLevel xd, final E akz) {
        final long long4 = akz.level.getGameTime();
        this.getRunningBehaviorsStream().forEach(amb -> amb.doStop(xd, akz, long4));
    }
    
    @Override
    public <T> T serialize(final DynamicOps<T> dynamicOps) {
        final T object3 = (T)dynamicOps.createMap((Map)this.memories.entrySet().stream().filter(entry -> entry.getKey().getDeserializer().isPresent() && ((Optional)entry.getValue()).isPresent()).map(entry -> Pair.of(dynamicOps.createString(Registry.MEMORY_MODULE_TYPE.getKey(entry.getKey()).toString()), ((Serializable)((Optional)entry.getValue()).get()).<T>serialize(dynamicOps))).collect(Collectors.toMap(Pair::getFirst, Pair::getSecond)));
        return (T)dynamicOps.createMap((Map)ImmutableMap.<Object, T>of(dynamicOps.createString("memories"), object3));
    }
    
    private void tickEachSensor(final ServerLevel xd, final E akz) {
        this.sensors.values().forEach(asf -> asf.tick(xd, akz));
    }
    
    private void startEachNonRunningBehavior(final ServerLevel xd, final E akz) {
        final long long4 = xd.getGameTime();
        this.availableGoalsByPriority.values().stream().flatMap(map -> map.entrySet().stream()).filter(entry -> this.activeActivities.contains(entry.getKey())).map(Map.Entry::getValue).flatMap(Collection::stream).filter(amb -> amb.getStatus() == Behavior.Status.STOPPED).forEach(amb -> amb.tryStart(xd, akz, long4));
    }
    
    private void tickEachRunningBehavior(final ServerLevel xd, final E akz) {
        final long long4 = xd.getGameTime();
        this.getRunningBehaviorsStream().forEach(amb -> amb.tickOrStop(xd, akz, long4));
    }
    
    private boolean activityRequirementsAreMet(final Activity azr) {
        final MemoryModuleType<?> arl3;
        final MemoryStatus arm4;
        return this.activityRequirements.get(azr).stream().allMatch(pair -> {
            arl3 = pair.getFirst();
            arm4 = (MemoryStatus)pair.getSecond();
            return this.checkMemory(arl3, arm4);
        });
    }
    
    private boolean isEmptyCollection(final Object object) {
        return object instanceof Collection && ((Collection)object).isEmpty();
    }
}
