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 alb, final float float2) {
        this(akz -> alb.equals(akz.getType().getCategory()), float2);
    }
    
    public SetEntityLookTarget(final EntityType<?> aku, final float float2) {
        this(akz -> aku.equals(akz.getType()), float2);
    }
    
    public SetEntityLookTarget(final Predicate<LivingEntity> predicate, final float float2) {
        super(ImmutableMap.<MemoryModuleType<PositionWrapper>, MemoryStatus>of(MemoryModuleType.LOOK_TARGET, MemoryStatus.VALUE_ABSENT, (MemoryModuleType<PositionWrapper>)MemoryModuleType.VISIBLE_LIVING_ENTITIES, MemoryStatus.VALUE_PRESENT));
        this.predicate = predicate;
        this.maxDistSqr = float2 * float2;
    }
    
    @Override
    protected boolean checkExtraStartConditions(final ServerLevel xd, final LivingEntity akz) {
        return akz.getBrain().<List<LivingEntity>>getMemory(MemoryModuleType.VISIBLE_LIVING_ENTITIES).get().stream().anyMatch(this.predicate);
    }
    
    @Override
    protected void start(final ServerLevel xd, final LivingEntity akz, final long long3) {
        final Brain<?> alp6 = akz.getBrain();
        alp6.<List<LivingEntity>>getMemory(MemoryModuleType.VISIBLE_LIVING_ENTITIES).ifPresent(list -> list.stream().filter(this.predicate).filter(akz2 -> akz2.distanceToSqr(akz) <= this.maxDistSqr).findFirst().ifPresent(akz -> alp6.<EntityPosWrapper>setMemory((MemoryModuleType<EntityPosWrapper>)MemoryModuleType.LOOK_TARGET, new EntityPosWrapper(akz))));
    }
}
