package net.minecraft.world.entity.ai.behavior;

import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Lists;
import java.util.Comparator;
import java.util.List;
import java.util.Optional;
import java.util.function.Function;
import java.util.function.ToIntFunction;
import net.minecraft.core.BlockPos;
import net.minecraft.core.Direction;
import net.minecraft.server.level.ServerLevel;
import net.minecraft.sounds.SoundEvent;
import net.minecraft.sounds.SoundSource;
import net.minecraft.util.Mth;
import net.minecraft.world.entity.LivingEntity;
import net.minecraft.world.entity.PathfinderMob;
import net.minecraft.world.entity.ai.Brain;
import net.minecraft.world.entity.ai.memory.MemoryModuleType;
import net.minecraft.world.entity.ai.memory.MemoryStatus;
import net.minecraft.world.entity.ai.memory.WalkTarget;
import net.minecraft.world.entity.ai.navigation.PathNavigation;
import net.minecraft.world.entity.ai.targeting.TargetingConditions;
import net.minecraft.world.level.pathfinder.Path;
import net.minecraft.world.level.pathfinder.WalkNodeEvaluator;
import net.minecraft.world.phys.Vec3;

public class PrepareRamNearestTarget<E extends PathfinderMob> extends Behavior<E> {
	public static final int TIME_OUT_DURATION = 160;
	private final ToIntFunction<E> getCooldownOnFail;
	private final int minRamDistance;
	private final int maxRamDistance;
	private final float walkSpeed;
	private final TargetingConditions ramTargeting;
	private final int ramPrepareTime;
	private final Function<E, SoundEvent> getPrepareRamSound;
	private Optional<Long> reachedRamPositionTimestamp = Optional.empty();
	private Optional<PrepareRamNearestTarget.RamCandidate> ramCandidate = Optional.empty();

	public PrepareRamNearestTarget(
		ToIntFunction<E> getCooldownOnFall,
		int minRamDistance,
		int maxRamDistance,
		float walkSpeed,
		TargetingConditions ramTargeting,
		int ramPrepareTime,
		Function<E, SoundEvent> getPrepareRamSound
	) {
		super(
			ImmutableMap.of(
				MemoryModuleType.LOOK_TARGET,
				MemoryStatus.REGISTERED,
				MemoryModuleType.RAM_COOLDOWN_TICKS,
				MemoryStatus.VALUE_ABSENT,
				MemoryModuleType.NEAREST_VISIBLE_LIVING_ENTITIES,
				MemoryStatus.VALUE_PRESENT,
				MemoryModuleType.RAM_TARGET,
				MemoryStatus.VALUE_ABSENT
			),
			160
		);
		this.getCooldownOnFail = getCooldownOnFall;
		this.minRamDistance = minRamDistance;
		this.maxRamDistance = maxRamDistance;
		this.walkSpeed = walkSpeed;
		this.ramTargeting = ramTargeting;
		this.ramPrepareTime = ramPrepareTime;
		this.getPrepareRamSound = getPrepareRamSound;
	}

	protected void start(ServerLevel serverLevel, PathfinderMob pathfinderMob, long l) {
		Brain<?> brain = pathfinderMob.getBrain();
		brain.getMemory(MemoryModuleType.NEAREST_VISIBLE_LIVING_ENTITIES)
			.flatMap(
				nearestVisibleLivingEntities -> nearestVisibleLivingEntities.findClosest(livingEntity -> this.ramTargeting.test(serverLevel, pathfinderMob, livingEntity))
			)
			.ifPresent(livingEntity -> this.chooseRamPosition(pathfinderMob, livingEntity));
	}

	protected void stop(ServerLevel serverLevel, E pathfinderMob, long l) {
		Brain<?> brain = pathfinderMob.getBrain();
		if (!brain.hasMemoryValue(MemoryModuleType.RAM_TARGET)) {
			serverLevel.broadcastEntityEvent(pathfinderMob, (byte)59);
			brain.setMemory(MemoryModuleType.RAM_COOLDOWN_TICKS, this.getCooldownOnFail.applyAsInt(pathfinderMob));
		}
	}

	protected boolean canStillUse(ServerLevel serverLevel, PathfinderMob pathfinderMob, long l) {
		return this.ramCandidate.isPresent() && ((PrepareRamNearestTarget.RamCandidate)this.ramCandidate.get()).getTarget().isAlive();
	}

	protected void tick(ServerLevel serverLevel, E pathfinderMob, long l) {
		if (!this.ramCandidate.isEmpty()) {
			pathfinderMob.getBrain()
				.setMemory(
					MemoryModuleType.WALK_TARGET, new WalkTarget(((PrepareRamNearestTarget.RamCandidate)this.ramCandidate.get()).getStartPosition(), this.walkSpeed, 0)
				);
			pathfinderMob.getBrain()
				.setMemory(MemoryModuleType.LOOK_TARGET, new EntityTracker(((PrepareRamNearestTarget.RamCandidate)this.ramCandidate.get()).getTarget(), true));
			boolean bl = !((PrepareRamNearestTarget.RamCandidate)this.ramCandidate.get())
				.getTarget()
				.blockPosition()
				.equals(((PrepareRamNearestTarget.RamCandidate)this.ramCandidate.get()).getTargetPosition());
			if (bl) {
				serverLevel.broadcastEntityEvent(pathfinderMob, (byte)59);
				pathfinderMob.getNavigation().stop();
				this.chooseRamPosition(pathfinderMob, ((PrepareRamNearestTarget.RamCandidate)this.ramCandidate.get()).target);
			} else {
				BlockPos blockPos = pathfinderMob.blockPosition();
				if (blockPos.equals(((PrepareRamNearestTarget.RamCandidate)this.ramCandidate.get()).getStartPosition())) {
					serverLevel.broadcastEntityEvent(pathfinderMob, (byte)58);
					if (this.reachedRamPositionTimestamp.isEmpty()) {
						this.reachedRamPositionTimestamp = Optional.of(l);
					}

					if (l - (Long)this.reachedRamPositionTimestamp.get() >= this.ramPrepareTime) {
						pathfinderMob.getBrain()
							.setMemory(
								MemoryModuleType.RAM_TARGET, this.getEdgeOfBlock(blockPos, ((PrepareRamNearestTarget.RamCandidate)this.ramCandidate.get()).getTargetPosition())
							);
						serverLevel.playSound(
							null, pathfinderMob, (SoundEvent)this.getPrepareRamSound.apply(pathfinderMob), SoundSource.NEUTRAL, 1.0F, pathfinderMob.getVoicePitch()
						);
						this.ramCandidate = Optional.empty();
					}
				}
			}
		}
	}

	private Vec3 getEdgeOfBlock(BlockPos pos, BlockPos other) {
		double d = 0.5;
		double e = 0.5 * Mth.sign(other.getX() - pos.getX());
		double f = 0.5 * Mth.sign(other.getZ() - pos.getZ());
		return Vec3.atBottomCenterOf(other).add(e, 0.0, f);
	}

	private Optional<BlockPos> calculateRammingStartPosition(PathfinderMob pathfinder, LivingEntity entity) {
		BlockPos blockPos = entity.blockPosition();
		if (!this.isWalkableBlock(pathfinder, blockPos)) {
			return Optional.empty();
		} else {
			List<BlockPos> list = Lists.<BlockPos>newArrayList();
			BlockPos.MutableBlockPos mutableBlockPos = blockPos.mutable();

			for (Direction direction : Direction.Plane.HORIZONTAL) {
				mutableBlockPos.set(blockPos);

				for (int i = 0; i < this.maxRamDistance; i++) {
					if (!this.isWalkableBlock(pathfinder, mutableBlockPos.move(direction))) {
						mutableBlockPos.move(direction.getOpposite());
						break;
					}
				}

				if (mutableBlockPos.distManhattan(blockPos) >= this.minRamDistance) {
					list.add(mutableBlockPos.immutable());
				}
			}

			PathNavigation pathNavigation = pathfinder.getNavigation();
			return list.stream().sorted(Comparator.comparingDouble(pathfinder.blockPosition()::distSqr)).filter(blockPosx -> {
				Path path = pathNavigation.createPath(blockPosx, 0);
				return path != null && path.canReach();
			}).findFirst();
		}
	}

	private boolean isWalkableBlock(PathfinderMob pathfinder, BlockPos pos) {
		return pathfinder.getNavigation().isStableDestination(pos) && pathfinder.getPathfindingMalus(WalkNodeEvaluator.getPathTypeStatic(pathfinder, pos)) == 0.0F;
	}

	private void chooseRamPosition(PathfinderMob pathfinder, LivingEntity entity) {
		this.reachedRamPositionTimestamp = Optional.empty();
		this.ramCandidate = this.calculateRammingStartPosition(pathfinder, entity)
			.map(blockPos -> new PrepareRamNearestTarget.RamCandidate(blockPos, entity.blockPosition(), entity));
	}

	public static class RamCandidate {
		private final BlockPos startPosition;
		private final BlockPos targetPosition;
		final LivingEntity target;

		public RamCandidate(BlockPos startPosition, BlockPos targetPosition, LivingEntity target) {
			this.startPosition = startPosition;
			this.targetPosition = targetPosition;
			this.target = target;
		}

		public BlockPos getStartPosition() {
			return this.startPosition;
		}

		public BlockPos getTargetPosition() {
			return this.targetPosition;
		}

		public LivingEntity getTarget() {
			return this.target;
		}
	}
}