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

import java.util.UUID;
import net.minecraft.world.level.GameRules;
import net.minecraft.world.entity.EntityType;
import net.minecraft.core.BlockPos;
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 java.util.function.BiPredicate;
import net.minecraft.world.entity.LivingEntity;

public class StartCelebratingIfTargetDead extends Behavior<LivingEntity> {
    private final int celebrateDuration;
    private final BiPredicate<LivingEntity, LivingEntity> dancePredicate;
    
    public StartCelebratingIfTargetDead(final int integer, final BiPredicate<LivingEntity, LivingEntity> biPredicate) {
        super(ImmutableMap.<MemoryModuleType<LivingEntity>, MemoryStatus>of(MemoryModuleType.ATTACK_TARGET, MemoryStatus.VALUE_PRESENT, (MemoryModuleType<LivingEntity>)MemoryModuleType.ANGRY_AT, MemoryStatus.REGISTERED, (MemoryModuleType<LivingEntity>)MemoryModuleType.CELEBRATE_LOCATION, MemoryStatus.VALUE_ABSENT, (MemoryModuleType<LivingEntity>)MemoryModuleType.DANCING, MemoryStatus.REGISTERED));
        this.celebrateDuration = integer;
        this.dancePredicate = biPredicate;
    }
    
    @Override
    protected boolean checkExtraStartConditions(final ServerLevel zd, final LivingEntity aoy) {
        return this.getAttackTarget(aoy).isDeadOrDying();
    }
    
    @Override
    protected void start(final ServerLevel zd, final LivingEntity aoy, final long long3) {
        final LivingEntity aoy2 = this.getAttackTarget(aoy);
        if (this.dancePredicate.test(aoy, aoy2)) {
            aoy.getBrain().<Boolean>setMemoryWithExpiry(MemoryModuleType.DANCING, true, this.celebrateDuration);
        }
        aoy.getBrain().<BlockPos>setMemoryWithExpiry(MemoryModuleType.CELEBRATE_LOCATION, aoy2.blockPosition(), this.celebrateDuration);
        if (aoy2.getType() != EntityType.PLAYER || zd.getGameRules().getBoolean(GameRules.RULE_FORGIVE_DEAD_PLAYERS)) {
            aoy.getBrain().<LivingEntity>eraseMemory(MemoryModuleType.ATTACK_TARGET);
            aoy.getBrain().<UUID>eraseMemory(MemoryModuleType.ANGRY_AT);
        }
    }
    
    private LivingEntity getAttackTarget(final LivingEntity aoy) {
        return aoy.getBrain().<LivingEntity>getMemory(MemoryModuleType.ATTACK_TARGET).get();
    }
}
