package net.minecraft.world.entity.ai.village.poi;

import it.unimi.dsi.fastutil.longs.Long2ByteOpenHashMap;
import it.unimi.dsi.fastutil.longs.Long2ByteMap;
import net.minecraft.server.level.SectionTracker;
import net.minecraft.world.level.chunk.ChunkStatus;
import com.mojang.datafixers.util.Pair;
import net.minecraft.world.level.LevelReader;
import net.minecraft.world.level.block.state.BlockState;
import java.util.function.BiConsumer;
import net.minecraft.Util;
import net.minecraft.world.level.chunk.LevelChunkSection;
import java.util.function.BooleanSupplier;
import java.util.Collections;
import java.util.stream.Collector;
import java.util.stream.Collectors;
import java.util.List;
import java.util.Random;
import java.util.Comparator;
import java.util.function.Function;
import java.util.stream.IntStream;
import net.minecraft.core.Vec3i;
import net.minecraft.world.level.ChunkPos;
import java.util.stream.Stream;
import java.util.Optional;
import java.util.function.Predicate;
import net.minecraft.core.SectionPos;
import net.minecraft.core.BlockPos;
import it.unimi.dsi.fastutil.longs.LongOpenHashSet;
import net.minecraft.util.datafix.DataFixTypes;
import com.mojang.datafixers.DataFixer;
import java.io.File;
import it.unimi.dsi.fastutil.longs.LongSet;
import net.minecraft.world.level.chunk.storage.SectionStorage;

public class PoiManager extends SectionStorage<PoiSection> {
    private final DistanceTracker distanceTracker;
    private final LongSet loadedChunks;
    
    public PoiManager(final File file, final DataFixer dataFixer, final boolean boolean3) {
        super(file, PoiSection::codec, PoiSection::new, dataFixer, DataFixTypes.POI_CHUNK, boolean3);
        this.loadedChunks = (LongSet)new LongOpenHashSet();
        this.distanceTracker = new DistanceTracker();
    }
    
    public void add(final BlockPos fu, final PoiType ayc) {
        this.getOrCreate(SectionPos.of(fu).asLong()).add(fu, ayc);
    }
    
    public void remove(final BlockPos fu) {
        this.getOrCreate(SectionPos.of(fu).asLong()).remove(fu);
    }
    
    public long getCountInRange(final Predicate<PoiType> predicate, final BlockPos fu, final int integer, final Occupancy b) {
        return this.getInRange(predicate, fu, integer, b).count();
    }
    
    public boolean existsAtPosition(final PoiType ayc, final BlockPos fu) {
        final Optional<PoiType> optional4 = this.getOrCreate(SectionPos.of(fu).asLong()).getType(fu);
        return optional4.isPresent() && optional4.get().equals(ayc);
    }
    
    public Stream<PoiRecord> getInSquare(final Predicate<PoiType> predicate, final BlockPos fu, final int integer, final Occupancy b) {
        final int integer2 = Math.floorDiv(integer, 16) + 1;
        return ChunkPos.rangeClosed(new ChunkPos(fu), integer2).<PoiRecord>flatMap(bph -> this.getInChunk(predicate, bph, b));
    }
    
    public Stream<PoiRecord> getInRange(final Predicate<PoiType> predicate, final BlockPos fu, final int integer, final Occupancy b) {
        final int integer2 = integer * integer;
        return this.getInSquare(predicate, fu, integer, b).filter(aya -> aya.getPos().distSqr(fu) <= integer2);
    }
    
    public Stream<PoiRecord> getInChunk(final Predicate<PoiType> predicate, final ChunkPos bph, final Occupancy b) {
        return IntStream.range(0, 16).boxed().map(integer -> this.getOrLoad(SectionPos.of(bph, integer).asLong())).filter(Optional::isPresent).<PoiRecord>flatMap(optional -> optional.get().getRecords(predicate, b));
    }
    
    public Stream<BlockPos> findAll(final Predicate<PoiType> predicate1, final Predicate<BlockPos> predicate2, final BlockPos fu, final int integer, final Occupancy b) {
        return this.getInRange(predicate1, fu, integer, b).<BlockPos>map(PoiRecord::getPos).filter(predicate2);
    }
    
    public Optional<BlockPos> find(final Predicate<PoiType> predicate1, final Predicate<BlockPos> predicate2, final BlockPos fu, final int integer, final Occupancy b) {
        return this.findAll(predicate1, predicate2, fu, integer, b).findFirst();
    }
    
    public Optional<BlockPos> findClosest(final Predicate<PoiType> predicate, final BlockPos fu, final int integer, final Occupancy b) {
        return this.getInRange(predicate, fu, integer, b).<BlockPos>map(PoiRecord::getPos).min(Comparator.comparingDouble(fu2 -> fu2.distSqr(fu)));
    }
    
    public Optional<BlockPos> take(final Predicate<PoiType> predicate1, final Predicate<BlockPos> predicate2, final BlockPos fu, final int integer) {
        return this.getInRange(predicate1, fu, integer, Occupancy.HAS_SPACE).filter(aya -> predicate2.test(aya.getPos())).findFirst().<BlockPos>map(aya -> {
            aya.acquireTicket();
            return aya.getPos();
        });
    }
    
    public Optional<BlockPos> getRandom(final Predicate<PoiType> predicate1, final Predicate<BlockPos> predicate2, final Occupancy b, final BlockPos fu, final int integer, final Random random) {
        final List<PoiRecord> list8 = this.getInRange(predicate1, fu, integer, b).collect(Collectors.toList());
        Collections.shuffle(list8, random);
        return list8.stream().filter(aya -> predicate2.test(aya.getPos())).findFirst().<BlockPos>map(PoiRecord::getPos);
    }
    
    public boolean release(final BlockPos fu) {
        return this.getOrCreate(SectionPos.of(fu).asLong()).release(fu);
    }
    
    public boolean exists(final BlockPos fu, final Predicate<PoiType> predicate) {
        return this.getOrLoad(SectionPos.of(fu).asLong()).<Boolean>map(ayb -> ayb.exists(fu, predicate)).orElse(false);
    }
    
    public Optional<PoiType> getType(final BlockPos fu) {
        final PoiSection ayb3 = this.getOrCreate(SectionPos.of(fu).asLong());
        return ayb3.getType(fu);
    }
    
    public int sectionsToVillage(final SectionPos go) {
        this.distanceTracker.runAllUpdates();
        return this.distanceTracker.getLevel(go.asLong());
    }
    
    private boolean isVillageCenter(final long long1) {
        final Optional<PoiSection> optional4 = this.get(long1);
        return optional4 != null && optional4.<Boolean>map(ayb -> ayb.getRecords(PoiType.ALL, Occupancy.IS_OCCUPIED).count() > 0L).orElse(false);
    }
    
    public void tick(final BooleanSupplier booleanSupplier) {
        super.tick(booleanSupplier);
        this.distanceTracker.runAllUpdates();
    }
    
    @Override
    protected void setDirty(final long long1) {
        super.setDirty(long1);
        this.distanceTracker.update(long1, this.distanceTracker.getLevelFromSource(long1), false);
    }
    
    @Override
    protected void onSectionLoad(final long long1) {
        this.distanceTracker.update(long1, this.distanceTracker.getLevelFromSource(long1), false);
    }
    
    public void checkConsistencyWithBlocks(final ChunkPos bph, final LevelChunkSection chk) {
        final SectionPos go4 = SectionPos.of(bph, chk.bottomBlockY() >> 4);
        final SectionPos go5;
        final SectionPos go6;
        PoiSection ayb2;
        Util.<PoiSection>ifElse(this.getOrLoad(go4.asLong()), ayb -> ayb.refresh(biConsumer -> {
            if (mayHavePoi(chk)) {
                this.updateFromSection(chk, go5, biConsumer);
            }
        }), () -> {
            if (mayHavePoi(chk)) {
                ayb2 = this.getOrCreate(go6.asLong());
                this.updateFromSection(chk, go6, ayb2::add);
            }
        });
    }
    
    private static boolean mayHavePoi(final LevelChunkSection chk) {
        return chk.maybeHas(PoiType.ALL_STATES::contains);
    }
    
    private void updateFromSection(final LevelChunkSection chk, final SectionPos go, final BiConsumer<BlockPos, PoiType> biConsumer) {
        final BlockState cfj4;
        go.blocksInside().forEach(fu -> {
            cfj4 = chk.getBlockState(SectionPos.sectionRelative(fu.getX()), SectionPos.sectionRelative(fu.getY()), SectionPos.sectionRelative(fu.getZ()));
            PoiType.forState(cfj4).ifPresent(ayc -> biConsumer.accept(fu, ayc));
        });
    }
    
    public void ensureLoadedAndValid(final LevelReader bqd, final BlockPos fu, final int integer) {
        SectionPos.aroundChunk(new ChunkPos(fu), Math.floorDiv(integer, 16)).map(go -> Pair.of(go, this.getOrLoad(go.asLong()))).filter(pair -> !((Optional)pair.getSecond()).map(PoiSection::isValid).orElse(false)).map(pair -> ((SectionPos)pair.getFirst()).chunk()).filter(bph -> this.loadedChunks.add(bph.toLong())).forEach(bph -> bqd.getChunk(bph.x, bph.z, ChunkStatus.EMPTY));
    }
    
    public enum Occupancy {
        HAS_SPACE(PoiRecord::hasSpace), 
        IS_OCCUPIED(PoiRecord::isOccupied), 
        ANY(aya -> true);
        
        private final Predicate<? super PoiRecord> test;
        
        private Occupancy(final Predicate<? super PoiRecord> predicate) {
            this.test = predicate;
        }
        
        public Predicate<? super PoiRecord> getTest() {
            return this.test;
        }
    }
    
    final class DistanceTracker extends SectionTracker {
        private final Long2ByteMap levels;
        
        protected DistanceTracker() {
            super(7, 16, 256);
            (this.levels = (Long2ByteMap)new Long2ByteOpenHashMap()).defaultReturnValue((byte)7);
        }
        
        @Override
        protected int getLevelFromSource(final long long1) {
            return PoiManager.this.isVillageCenter(long1) ? 0 : 7;
        }
        
        @Override
        protected int getLevel(final long long1) {
            return this.levels.get(long1);
        }
        
        @Override
        protected void setLevel(final long long1, final int integer) {
            if (integer > 6) {
                this.levels.remove(long1);
            }
            else {
                this.levels.put(long1, (byte)integer);
            }
        }
        
        public void runAllUpdates() {
            super.runUpdates(Integer.MAX_VALUE);
        }
    }
}
