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.google.common.collect.ImmutableList.Builder;
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.Map.Entry;
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.entity.LivingEntity;
import net.minecraft.world.entity.ai.behavior.BehaviorControl;
import net.minecraft.world.entity.ai.behavior.Behavior.Status;
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.entity.schedule.Schedule;
import org.apache.commons.lang3.mutable.MutableObject;
import org.jetbrains.annotations.Nullable;
import org.slf4j.Logger;

public class Brain<E extends LivingEntity> {
	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.<MemoryModuleType<?>, Optional<? extends ExpirableValue<?>>>newHashMap();
	private final Map<SensorType<? extends Sensor<? super E>>, Sensor<? super E>> sensors = Maps.<SensorType<? extends Sensor<? super E>>, Sensor<? super E>>newLinkedHashMap();
	private final Map<Integer, Map<Activity, Set<BehaviorControl<? super E>>>> availableBehaviorsByPriority = Maps.newTreeMap();
	private Schedule schedule = Schedule.EMPTY;
	private final Map<Activity, Set<Pair<MemoryModuleType<?>, MemoryStatus>>> activityRequirements = Maps.<Activity, Set<Pair<MemoryModuleType<?>, MemoryStatus>>>newHashMap();
	private final Map<Activity, Set<MemoryModuleType<?>>> activityMemoriesToEraseWhenStopped = Maps.<Activity, Set<MemoryModuleType<?>>>newHashMap();
	private Set<Activity> coreActivities = Sets.<Activity>newHashSet();
	private final Set<Activity> activeActivities = Sets.<Activity>newHashSet();
	private Activity defaultActivity = Activity.IDLE;
	private long lastScheduleUpdate = -9999L;

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

	public static <E extends LivingEntity> Codec<Brain<E>> codec(
		Collection<? extends MemoryModuleType<?>> memoryTypes, Collection<? extends SensorType<? extends Sensor<? super E>>> sensorTypes
	) {
		final MutableObject<Codec<Brain<E>>> mutableObject = new MutableObject<>();
		mutableObject.setValue(
			(new MapCodec<Brain<E>>() {
					@Override
					public <T> Stream<T> keys(DynamicOps<T> dynamicOps) {
						return memoryTypes.stream()
							.flatMap(memoryModuleType -> memoryModuleType.getCodec().map(codec -> BuiltInRegistries.MEMORY_MODULE_TYPE.getKey(memoryModuleType)).stream())
							.map(resourceLocation -> dynamicOps.createString(resourceLocation.toString()));
					}

					@Override
					public <T> DataResult<Brain<E>> decode(DynamicOps<T> dynamicOps, MapLike<T> mapLike) {
						MutableObject<DataResult<Builder<Brain.MemoryValue<?>>>> mutableObjectx = new MutableObject<>(DataResult.success(ImmutableList.builder()));
						mapLike.entries()
							.forEach(
								pair -> {
									DataResult<MemoryModuleType<?>> dataResult = BuiltInRegistries.MEMORY_MODULE_TYPE.byNameCodec().parse(dynamicOps, (T)pair.getFirst());
									DataResult<? extends Brain.MemoryValue<?>> dataResult2 = dataResult.flatMap(
										memoryModuleType -> this.captureRead(memoryModuleType, dynamicOps, (T)pair.getSecond())
									);
									mutableObject.setValue(mutableObject.getValue().apply2(Builder::add, dataResult2));
								}
							);
						ImmutableList<Brain.MemoryValue<?>> immutableList = (ImmutableList<Brain.MemoryValue<?>>)mutableObjectx.getValue()
							.resultOrPartial(Brain.LOGGER::error)
							.map(Builder::build)
							.orElseGet(ImmutableList::of);
						return DataResult.success(new Brain<>(memoryTypes, sensorTypes, immutableList, mutableObject::getValue));
					}

					private <T, U> DataResult<Brain.MemoryValue<U>> captureRead(MemoryModuleType<U> memoryModuleType, DynamicOps<T> dynamicOps, T object) {
						return ((DataResult)memoryModuleType.getCodec()
								.map(DataResult::success)
								.orElseGet(() -> DataResult.error(() -> "No codec for memory: " + memoryModuleType)))
							.flatMap(codec -> codec.parse(dynamicOps, object))
							.map(expirableValue -> new Brain.MemoryValue<>(memoryModuleType, Optional.of(expirableValue)));
					}

					public <T> RecordBuilder<T> encode(Brain<E> input, DynamicOps<T> ops, RecordBuilder<T> prefix) {
						input.memories().forEach(memoryValue -> memoryValue.serialize(ops, prefix));
						return prefix;
					}
				})
				.fieldOf("memories")
				.codec()
		);
		return mutableObject.getValue();
	}

	public Brain(
		Collection<? extends MemoryModuleType<?>> memoryModuleTypes,
		Collection<? extends SensorType<? extends Sensor<? super E>>> sensorTypes,
		ImmutableList<Brain.MemoryValue<?>> memoryValues,
		Supplier<Codec<Brain<E>>> codec
	) {
		this.codec = codec;

		for (MemoryModuleType<?> memoryModuleType : memoryModuleTypes) {
			this.memories.put(memoryModuleType, Optional.empty());
		}

		for (SensorType<? extends Sensor<? super E>> sensorType : sensorTypes) {
			this.sensors.put(sensorType, sensorType.create());
		}

		for (Sensor<? super E> sensor : this.sensors.values()) {
			for (MemoryModuleType<?> memoryModuleType2 : sensor.requires()) {
				this.memories.put(memoryModuleType2, Optional.empty());
			}
		}

		for (Brain.MemoryValue<?> memoryValue : memoryValues) {
			memoryValue.setMemoryInternal(this);
		}
	}

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

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

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

	public void clearMemories() {
		this.memories.keySet().forEach(memoryModuleType -> this.memories.put(memoryModuleType, Optional.empty()));
	}

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

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

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

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

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

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

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

	public <U> long getTimeUntilExpiry(MemoryModuleType<U> memoryType) {
		Optional<? extends ExpirableValue<?>> optional = (Optional<? extends ExpirableValue<?>>)this.memories.get(memoryType);
		return (Long)optional.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 memory) {
		return !this.hasMemoryValue(memoryType) ? false : this.getMemory(memoryType).filter(object2 -> object2.equals(memory)).isPresent();
	}

	public boolean checkMemory(MemoryModuleType<?> memoryType, MemoryStatus memoryStatus) {
		Optional<? extends ExpirableValue<?>> optional = (Optional<? extends ExpirableValue<?>>)this.memories.get(memoryType);
		return optional == null
			? false
			: memoryStatus == MemoryStatus.REGISTERED
				|| memoryStatus == MemoryStatus.VALUE_PRESENT && optional.isPresent()
				|| memoryStatus == MemoryStatus.VALUE_ABSENT && optional.isEmpty();
	}

	public Schedule getSchedule() {
		return this.schedule;
	}

	public void setSchedule(Schedule newSchedule) {
		this.schedule = newSchedule;
	}

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

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

	@Deprecated
	@VisibleForDebug
	public List<BehaviorControl<? super E>> getRunningBehaviors() {
		List<BehaviorControl<? super E>> list = new ObjectArrayList<>();

		for (Map<Activity, Set<BehaviorControl<? super E>>> map : this.availableBehaviorsByPriority.values()) {
			for (Set<BehaviorControl<? super E>> set : map.values()) {
				for (BehaviorControl<? super E> behaviorControl : set) {
					if (behaviorControl.getStatus() == Status.RUNNING) {
						list.add(behaviorControl);
					}
				}
			}
		}

		return list;
	}

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

	public Optional<Activity> getActiveNonCoreActivity() {
		for (Activity activity : this.activeActivities) {
			if (!this.coreActivities.contains(activity)) {
				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)) {
			this.eraseMemoriesForOtherActivitesThan(activity);
			this.activeActivities.clear();
			this.activeActivities.addAll(this.coreActivities);
			this.activeActivities.add(activity);
		}
	}

	private void eraseMemoriesForOtherActivitesThan(Activity activity) {
		for (Activity activity2 : this.activeActivities) {
			if (activity2 != activity) {
				Set<MemoryModuleType<?>> set = (Set<MemoryModuleType<?>>)this.activityMemoriesToEraseWhenStopped.get(activity2);
				if (set != null) {
					for (MemoryModuleType<?> memoryModuleType : set) {
						this.eraseMemory(memoryModuleType);
					}
				}
			}
		}
	}

	public void updateActivityFromSchedule(long dayTime, long gameTime) {
		if (gameTime - this.lastScheduleUpdate > 20L) {
			this.lastScheduleUpdate = gameTime;
			Activity activity = this.getSchedule().getActivityAt((int)(dayTime % 24000L));
			if (!this.activeActivities.contains(activity)) {
				this.setActiveActivityIfPossible(activity);
			}
		}
	}

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

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

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

	public void addActivityAndRemoveMemoryWhenStopped(
		Activity activity, int priorityStart, ImmutableList<? extends BehaviorControl<? super E>> tasks, MemoryModuleType<?> memoryType
	) {
		Set<Pair<MemoryModuleType<?>, MemoryStatus>> set = ImmutableSet.of(Pair.of(memoryType, MemoryStatus.VALUE_PRESENT));
		Set<MemoryModuleType<?>> set2 = ImmutableSet.of(memoryType);
		this.addActivityAndRemoveMemoriesWhenStopped(activity, this.createPriorityPairs(priorityStart, tasks), set, set2);
	}

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

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

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

	public void addActivityAndRemoveMemoriesWhenStopped(
		Activity activity,
		ImmutableList<? extends Pair<Integer, ? extends BehaviorControl<? super E>>> tasks,
		Set<Pair<MemoryModuleType<?>, MemoryStatus>> memorieStatuses,
		Set<MemoryModuleType<?>> memoryTypes
	) {
		this.activityRequirements.put(activity, memorieStatuses);
		if (!memoryTypes.isEmpty()) {
			this.activityMemoriesToEraseWhenStopped.put(activity, memoryTypes);
		}

		for (Pair<Integer, ? extends BehaviorControl<? super E>> pair : tasks) {
			((Set)((Map)this.availableBehaviorsByPriority.computeIfAbsent(pair.getFirst(), integer -> Maps.newHashMap()))
					.computeIfAbsent(activity, activityx -> Sets.newLinkedHashSet()))
				.add(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<>(this.memories.keySet(), this.sensors.keySet(), ImmutableList.of(), this.codec);

		for (Entry<MemoryModuleType<?>, Optional<? extends ExpirableValue<?>>> entry : this.memories.entrySet()) {
			MemoryModuleType<?> memoryModuleType = (MemoryModuleType<?>)entry.getKey();
			if (((Optional)entry.getValue()).isPresent()) {
				brain.memories.put(memoryModuleType, (Optional)entry.getValue());
			}
		}

		return brain;
	}

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

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

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

				expirableValue.tick();
			}
		}
	}

	public void stopAll(ServerLevel level, E owner) {
		long l = owner.level().getGameTime();

		for (BehaviorControl<? super E> behaviorControl : this.getRunningBehaviors()) {
			behaviorControl.doStop(level, owner, l);
		}
	}

	private void startEachNonRunningBehavior(ServerLevel level, E entity) {
		long l = level.getGameTime();

		for (Map<Activity, Set<BehaviorControl<? super E>>> map : this.availableBehaviorsByPriority.values()) {
			for (Entry<Activity, Set<BehaviorControl<? super E>>> entry : map.entrySet()) {
				Activity activity = (Activity)entry.getKey();
				if (this.activeActivities.contains(activity)) {
					for (BehaviorControl<? super E> behaviorControl : (Set)entry.getValue()) {
						if (behaviorControl.getStatus() == Status.STOPPED) {
							behaviorControl.tryStart(level, entity, l);
						}
					}
				}
			}
		}
	}

	private void tickEachRunningBehavior(ServerLevel level, E entity) {
		long l = level.getGameTime();

		for (BehaviorControl<? super E> behaviorControl : this.getRunningBehaviors()) {
			behaviorControl.tickOrStop(level, entity, l);
		}
	}

	private boolean activityRequirementsAreMet(Activity activity) {
		if (!this.activityRequirements.containsKey(activity)) {
			return false;
		} else {
			for (Pair<MemoryModuleType<?>, MemoryStatus> pair : (Set)this.activityRequirements.get(activity)) {
				MemoryModuleType<?> memoryModuleType = pair.getFirst();
				MemoryStatus memoryStatus = pair.getSecond();
				if (!this.checkMemory(memoryModuleType, memoryStatus)) {
					return false;
				}
			}

			return true;
		}
	}

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

	ImmutableList<? extends Pair<Integer, ? extends BehaviorControl<? super E>>> createPriorityPairs(
		int priorityStart, ImmutableList<? extends BehaviorControl<? super E>> tasks
	) {
		int i = priorityStart;
		Builder<Pair<Integer, ? extends BehaviorControl<? super E>>> builder = ImmutableList.builder();

		for (BehaviorControl<? super E> behaviorControl : tasks) {
			builder.add(Pair.of(i++, behaviorControl));
		}

		return builder.build();
	}

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

		static <U> Brain.MemoryValue<U> createUnchecked(MemoryModuleType<U> memoryType, Optional<? extends ExpirableValue<?>> memory) {
			return new Brain.MemoryValue<>(memoryType, (Optional<? extends ExpirableValue<U>>)memory);
		}

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

		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(
							expirableValue -> builder.add(BuiltInRegistries.MEMORY_MODULE_TYPE.byNameCodec().encodeStart(ops, this.type), codec.encodeStart(ops, expirableValue))
						)
				);
		}
	}

	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;

		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<?> ops) {
			return (Brain<E>)this.codec
				.parse(ops)
				.resultOrPartial(Brain.LOGGER::error)
				.orElseGet(() -> new Brain(this.memoryTypes, this.sensorTypes, ImmutableList.of(), () -> this.codec));
		}
	}
}