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

import net.minecraft.world.entity.Entity;
import net.minecraft.world.entity.ai.Brain;
import net.minecraft.server.level.ServerLevel;
import net.minecraft.world.damagesource.DamageSource;
import com.google.common.collect.ImmutableSet;
import net.minecraft.world.entity.ai.memory.MemoryModuleType;
import java.util.Set;
import net.minecraft.world.entity.LivingEntity;

public class HurtBySensor extends Sensor<LivingEntity> {
    @Override
    public Set<MemoryModuleType<?>> requires() {
        return ImmutableSet.of(MemoryModuleType.HURT_BY, (MemoryModuleType<DamageSource>)MemoryModuleType.HURT_BY_ENTITY);
    }
    
    @Override
    protected void doTick(final ServerLevel zd, final LivingEntity aoy) {
        final Brain<?> apr4 = aoy.getBrain();
        final DamageSource anw5 = aoy.getLastDamageSource();
        if (anw5 != null) {
            apr4.<DamageSource>setMemory(MemoryModuleType.HURT_BY, aoy.getLastDamageSource());
            final Entity aom6 = anw5.getEntity();
            if (aom6 instanceof LivingEntity) {
                apr4.<LivingEntity>setMemory(MemoryModuleType.HURT_BY_ENTITY, (LivingEntity)aom6);
            }
        }
        else {
            apr4.<DamageSource>eraseMemory(MemoryModuleType.HURT_BY);
        }
        final Brain brain;
        apr4.<LivingEntity>getMemory(MemoryModuleType.HURT_BY_ENTITY).ifPresent(aoy -> {
            if (!aoy.isAlive() || aoy.level != zd) {
                brain.eraseMemory(MemoryModuleType.HURT_BY_ENTITY);
            }
        });
    }
}
