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

import net.minecraft.world.entity.ai.Brain;
import net.minecraft.world.entity.Entity;
import java.util.List;
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.EntityType;
import net.minecraft.world.entity.MobCategory;
import java.util.function.Predicate;
import net.minecraft.world.entity.LivingEntity;

public class SetEntityLookTarget extends Behavior<LivingEntity> {
    private final Predicate<LivingEntity> predicate;
    private final float maxDistSqr;
    
    public SetEntityLookTarget(final MobCategory apa, final float float2) {
        this(aoy -> apa.equals(aoy.getType().getCategory()), float2);
    }
    
    public SetEntityLookTarget(final EntityType<?> aoq, final float float2) {
        this(aoy -> aoq.equals(aoy.getType()), float2);
    }
    
    public SetEntityLookTarget(final float float1) {
        this(aoy -> true, float1);
    }
    
    public SetEntityLookTarget(final Predicate<LivingEntity> predicate, final float float2) {
        super(ImmutableMap.<MemoryModuleType<PositionTracker>, MemoryStatus>of(MemoryModuleType.LOOK_TARGET, MemoryStatus.VALUE_ABSENT, (MemoryModuleType<PositionTracker>)MemoryModuleType.VISIBLE_LIVING_ENTITIES, MemoryStatus.VALUE_PRESENT));
        this.predicate = predicate;
        this.maxDistSqr = float2 * float2;
    }
    
    @Override
    protected boolean checkExtraStartConditions(final ServerLevel zd, final LivingEntity aoy) {
        return aoy.getBrain().<List<LivingEntity>>getMemory(MemoryModuleType.VISIBLE_LIVING_ENTITIES).get().stream().anyMatch(this.predicate);
    }
    
    @Override
    protected void start(final ServerLevel zd, final LivingEntity aoy, final long long3) {
        final Brain<?> apr6 = aoy.getBrain();
        apr6.<List<LivingEntity>>getMemory(MemoryModuleType.VISIBLE_LIVING_ENTITIES).ifPresent(list -> list.stream().filter(this.predicate).filter(aoy2 -> aoy2.distanceToSqr(aoy) <= this.maxDistSqr).findFirst().ifPresent(aoy -> apr6.<EntityTracker>setMemory((MemoryModuleType<EntityTracker>)MemoryModuleType.LOOK_TARGET, new EntityTracker(aoy, true))));
    }
}
