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

import net.minecraft.world.entity.LivingEntity;
import java.util.Optional;
import net.minecraft.world.level.pathfinder.Path;
import java.util.stream.Stream;
import java.util.function.Predicate;
import net.minecraft.world.entity.ai.village.poi.PoiManager;
import net.minecraft.world.entity.ai.village.poi.PoiType;
import net.minecraft.server.level.ServerLevel;
import net.minecraft.core.BlockPos;
import com.google.common.collect.ImmutableSet;
import net.minecraft.world.entity.ai.memory.MemoryModuleType;
import java.util.Set;
import it.unimi.dsi.fastutil.longs.Long2LongOpenHashMap;
import it.unimi.dsi.fastutil.longs.Long2LongMap;
import net.minecraft.world.entity.Mob;

public class NearestBedSensor extends Sensor<Mob> {
    private final Long2LongMap batchCache;
    private int triedCount;
    private long lastUpdate;
    
    public NearestBedSensor() {
        super(20);
        this.batchCache = (Long2LongMap)new Long2LongOpenHashMap();
    }
    
    @Override
    public Set<MemoryModuleType<?>> requires() {
        return ImmutableSet.of(MemoryModuleType.NEAREST_BED);
    }
    
    @Override
    protected void doTick(final ServerLevel zd, final Mob aoz) {
        if (!aoz.isBaby()) {
            return;
        }
        this.triedCount = 0;
        this.lastUpdate = zd.getGameTime() + zd.getRandom().nextInt(20);
        final PoiManager axz4 = zd.getPoiManager();
        final long long3;
        final Predicate<BlockPos> predicate5 = fu -> {
            long3 = fu.asLong();
            if (this.batchCache.containsKey(long3)) {
                return false;
            }
            else if (++this.triedCount >= 5) {
                return false;
            }
            else {
                this.batchCache.put(long3, this.lastUpdate + 40L);
                return true;
            }
        };
        final Stream<BlockPos> stream6 = axz4.findAll(PoiType.HOME.getPredicate(), predicate5, aoz.blockPosition(), 48, PoiManager.Occupancy.ANY);
        final Path czf7 = aoz.getNavigation().createPath(stream6, PoiType.HOME.getValidRange());
        if (czf7 != null && czf7.canReach()) {
            final BlockPos fu2 = czf7.getTarget();
            final Optional<PoiType> optional9 = axz4.getType(fu2);
            if (optional9.isPresent()) {
                aoz.getBrain().<BlockPos>setMemory(MemoryModuleType.NEAREST_BED, fu2);
            }
        }
        else if (this.triedCount < 5) {
            this.batchCache.long2LongEntrySet().removeIf(entry -> entry.getLongValue() < this.lastUpdate);
        }
    }
}
