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

import java.util.Optional;
import net.minecraft.world.entity.EntitySelector;
import net.minecraft.server.level.ServerLevel;
import java.util.Map;
import com.google.common.collect.ImmutableMap;
import net.minecraft.world.entity.ai.memory.MemoryStatus;
import net.minecraft.world.entity.ai.memory.MemoryModuleType;
import net.minecraft.world.entity.LivingEntity;
import java.util.function.Predicate;
import net.minecraft.world.entity.Mob;

public class StopAttackingIfTargetInvalid<E extends Mob> extends Behavior<E> {
    private final Predicate<LivingEntity> stopAttackingWhen;
    
    public StopAttackingIfTargetInvalid(final Predicate<LivingEntity> predicate) {
        super(ImmutableMap.<MemoryModuleType<LivingEntity>, MemoryStatus>of(MemoryModuleType.ATTACK_TARGET, MemoryStatus.VALUE_PRESENT, (MemoryModuleType<LivingEntity>)MemoryModuleType.CANT_REACH_WALK_TARGET_SINCE, MemoryStatus.REGISTERED));
        this.stopAttackingWhen = predicate;
    }
    
    public StopAttackingIfTargetInvalid() {
        this(aoy -> false);
    }
    
    @Override
    protected void start(final ServerLevel zd, final E aoz, final long long3) {
        if (StopAttackingIfTargetInvalid.<E>isTiredOfTryingToReachTarget(aoz)) {
            this.clearAttackTarget(aoz);
            return;
        }
        if (this.isCurrentTargetDeadOrRemoved(aoz)) {
            this.clearAttackTarget(aoz);
            return;
        }
        if (this.isCurrentTargetInDifferentLevel(aoz)) {
            this.clearAttackTarget(aoz);
            return;
        }
        if (!EntitySelector.ATTACK_ALLOWED.test(this.getAttackTarget(aoz))) {
            this.clearAttackTarget(aoz);
            return;
        }
        if (this.stopAttackingWhen.test(this.getAttackTarget(aoz))) {
            this.clearAttackTarget(aoz);
        }
    }
    
    private boolean isCurrentTargetInDifferentLevel(final E aoz) {
        return this.getAttackTarget(aoz).level != aoz.level;
    }
    
    private LivingEntity getAttackTarget(final E aoz) {
        return aoz.getBrain().<LivingEntity>getMemory(MemoryModuleType.ATTACK_TARGET).get();
    }
    
    private static <E extends LivingEntity> boolean isTiredOfTryingToReachTarget(final E aoy) {
        final Optional<Long> optional2 = aoy.getBrain().<Long>getMemory(MemoryModuleType.CANT_REACH_WALK_TARGET_SINCE);
        return optional2.isPresent() && aoy.level.getGameTime() - optional2.get() > 200L;
    }
    
    private boolean isCurrentTargetDeadOrRemoved(final E aoz) {
        final Optional<LivingEntity> optional3 = aoz.getBrain().<LivingEntity>getMemory(MemoryModuleType.ATTACK_TARGET);
        return optional3.isPresent() && !optional3.get().isAlive();
    }
    
    private void clearAttackTarget(final E aoz) {
        aoz.getBrain().<LivingEntity>eraseMemory(MemoryModuleType.ATTACK_TARGET);
    }
}
