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 java.util.function.BiConsumer;
import net.minecraft.world.level.block.state.BlockState;
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.Optional;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.stream.IntStream;
import net.minecraft.core.Vec3i;
import net.minecraft.world.level.ChunkPos;
import java.util.stream.Stream;
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) {
        super(file, PoiSection::new, PoiSection::new, dataFixer, DataFixTypes.POI_CHUNK);
        this.loadedChunks = (LongSet)new LongOpenHashSet();
        this.distanceTracker = new DistanceTracker();
    }
    
    public void add(final BlockPos fk, final PoiType asu) {
        this.getOrCreate(SectionPos.of(fk).asLong()).add(fk, asu);
    }
    
    public void remove(final BlockPos fk) {
        this.getOrCreate(SectionPos.of(fk).asLong()).remove(fk);
    }
    
    public long getCountInRange(final Predicate<PoiType> predicate, final BlockPos fk, final int integer, final Occupancy b) {
        return this.getInRange(predicate, fk, integer, b).count();
    }
    
    public Stream<PoiRecord> getInSquare(final Predicate<PoiType> predicate, final BlockPos fk, final int integer, final Occupancy b) {
        return ChunkPos.rangeClosed(new ChunkPos(fk), Math.floorDiv(integer, 16)).<PoiRecord>flatMap(bjh -> this.getInChunk(predicate, bjh, b));
    }
    
    public Stream<PoiRecord> getInRange(final Predicate<PoiType> predicate, final BlockPos fk, final int integer, final Occupancy b) {
        final int integer2 = integer * integer;
        return this.getInSquare(predicate, fk, integer, b).filter(ass -> ass.getPos().distSqr(fk) <= integer2);
    }
    
    public Stream<PoiRecord> getInChunk(final Predicate<PoiType> predicate, final ChunkPos bjh, final Occupancy b) {
        return IntStream.range(0, 16).boxed().<PoiRecord>flatMap(integer -> this.getInSection(predicate, SectionPos.of(bjh, integer).asLong(), b));
    }
    
    private Stream<PoiRecord> getInSection(final Predicate<PoiType> predicate, final long long2, final Occupancy b) {
        return this.getOrLoad(long2).<Stream<PoiRecord>>map(ast -> ast.getRecords(predicate, b)).orElseGet(Stream::empty);
    }
    
    public Stream<BlockPos> findAll(final Predicate<PoiType> predicate1, final Predicate<BlockPos> predicate2, final BlockPos fk, final int integer, final Occupancy b) {
        return this.getInRange(predicate1, fk, integer, b).<BlockPos>map(PoiRecord::getPos).filter(predicate2);
    }
    
    public Optional<BlockPos> find(final Predicate<PoiType> predicate1, final Predicate<BlockPos> predicate2, final BlockPos fk, final int integer, final Occupancy b) {
        return this.findAll(predicate1, predicate2, fk, integer, b).findFirst();
    }
    
    public Optional<BlockPos> findClosest(final Predicate<PoiType> predicate1, final Predicate<BlockPos> predicate2, final BlockPos fk, final int integer, final Occupancy b) {
        return this.getInRange(predicate1, fk, integer, b).<BlockPos>map(PoiRecord::getPos).sorted(Comparator.comparingDouble(fk2 -> fk2.distSqr(fk))).filter(predicate2).findFirst();
    }
    
    public Optional<BlockPos> take(final Predicate<PoiType> predicate1, final Predicate<BlockPos> predicate2, final BlockPos fk, final int integer) {
        return this.getInRange(predicate1, fk, integer, Occupancy.HAS_SPACE).filter(ass -> predicate2.test(ass.getPos())).findFirst().<BlockPos>map(ass -> {
            ass.acquireTicket();
            return ass.getPos();
        });
    }
    
    public Optional<BlockPos> getRandom(final Predicate<PoiType> predicate1, final Predicate<BlockPos> predicate2, final Occupancy b, final BlockPos fk, final int integer, final Random random) {
        final List<PoiRecord> list8 = this.getInRange(predicate1, fk, integer, b).collect(Collectors.toList());
        Collections.shuffle(list8, random);
        return list8.stream().filter(ass -> predicate2.test(ass.getPos())).findFirst().<BlockPos>map(PoiRecord::getPos);
    }
    
    public boolean release(final BlockPos fk) {
        return this.getOrCreate(SectionPos.of(fk).asLong()).release(fk);
    }
    
    public boolean exists(final BlockPos fk, final Predicate<PoiType> predicate) {
        return this.getOrLoad(SectionPos.of(fk).asLong()).<Boolean>map(ast -> ast.exists(fk, predicate)).orElse(false);
    }
    
    public Optional<PoiType> getType(final BlockPos fk) {
        final PoiSection ast3 = this.getOrCreate(SectionPos.of(fk).asLong());
        return ast3.getType(fk);
    }
    
    public int sectionsToVillage(final SectionPos gd) {
        this.distanceTracker.runAllUpdates();
        return this.distanceTracker.getLevel(gd.asLong());
    }
    
    private boolean isVillageCenter(final long long1) {
        final Optional<PoiSection> optional4 = this.get(long1);
        return optional4 != null && optional4.<Boolean>map(ast -> ast.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 bjh, final LevelChunkSection cam) {
        final SectionPos gd4 = SectionPos.of(bjh, cam.bottomBlockY() >> 4);
        final SectionPos gd5;
        final SectionPos gd6;
        PoiSection ast2;
        Util.<PoiSection>ifElse(this.getOrLoad(gd4.asLong()), ast -> ast.refresh(biConsumer -> {
            if (mayHavePoi(cam)) {
                this.updateFromSection(cam, gd5, biConsumer);
            }
        }), () -> {
            if (mayHavePoi(cam)) {
                ast2 = this.getOrCreate(gd6.asLong());
                this.updateFromSection(cam, gd6, ast2::add);
            }
        });
    }
    
    private static boolean mayHavePoi(final LevelChunkSection cam) {
        return PoiType.allPoiStates().anyMatch(cam::maybeHas);
    }
    
    private void updateFromSection(final LevelChunkSection cam, final SectionPos gd, final BiConsumer<BlockPos, PoiType> biConsumer) {
        final BlockState byj4;
        gd.blocksInside().forEach(fk -> {
            byj4 = cam.getBlockState(SectionPos.sectionRelative(fk.getX()), SectionPos.sectionRelative(fk.getY()), SectionPos.sectionRelative(fk.getZ()));
            PoiType.forState(byj4).ifPresent(asu -> biConsumer.accept(fk, asu));
        });
    }
    
    public void ensureLoadedAndValid(final LevelReader bjz, final BlockPos fk, final int integer) {
        SectionPos.aroundChunk(new ChunkPos(fk), Math.floorDiv(integer, 16)).map(gd -> Pair.of(gd, this.getOrLoad(gd.asLong()))).filter(pair -> !((Optional)pair.getSecond()).map(PoiSection::isValid).orElse(false)).map(pair -> ((SectionPos)pair.getFirst()).chunk()).filter(bjh -> this.loadedChunks.add(bjh.toLong())).forEach(bjh -> bjz.getChunk(bjh.x, bjh.z, ChunkStatus.EMPTY));
    }
    
    public enum Occupancy {
        HAS_SPACE(PoiRecord::hasSpace), 
        IS_OCCUPIED(PoiRecord::isOccupied), 
        ANY(ass -> 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);
        }
    }
}
