package net.minecraft.world.level.chunk.storage;

import java.util.AbstractList;
import org.apache.logging.log4j.LogManager;
import it.unimi.dsi.fastutil.shorts.ShortListIterator;
import net.minecraft.nbt.ShortTag;
import it.unimi.dsi.fastutil.shorts.ShortList;
import it.unimi.dsi.fastutil.longs.LongOpenHashSet;
import java.util.Locale;
import com.google.common.collect.Maps;
import it.unimi.dsi.fastutil.longs.LongSet;
import net.minecraft.world.level.levelgen.structure.StructureStart;
import net.minecraft.world.level.levelgen.feature.StructureFeature;
import net.minecraft.world.level.Level;
import net.minecraft.world.level.block.entity.BlockEntity;
import net.minecraft.world.entity.EntityType;
import javax.annotation.Nullable;
import net.minecraft.nbt.LongArrayTag;
import java.util.Map;
import java.util.Collection;
import net.minecraft.world.entity.Entity;
import java.util.Arrays;
import net.minecraft.nbt.Tag;
import net.minecraft.SharedConstants;
import java.util.Iterator;
import net.minecraft.world.level.chunk.ChunkAccess;
import net.minecraft.world.level.TickList;
import net.minecraft.world.level.lighting.LevelLightEngine;
import net.minecraft.world.level.chunk.ChunkSource;
import net.minecraft.nbt.ListTag;
import net.minecraft.world.level.biome.BiomeSource;
import net.minecraft.world.level.chunk.ChunkGenerator;
import java.util.BitSet;
import net.minecraft.world.level.levelgen.GenerationStep;
import net.minecraft.world.level.chunk.ImposterProtoChunk;
import java.util.Set;
import java.util.EnumSet;
import net.minecraft.world.level.levelgen.Heightmap;
import net.minecraft.core.BlockPos;
import net.minecraft.world.level.chunk.LevelChunk;
import net.minecraft.world.level.ChunkTickList;
import net.minecraft.core.Registry;
import net.minecraft.world.level.chunk.ChunkStatus;
import net.minecraft.world.level.chunk.DataLayer;
import net.minecraft.core.SectionPos;
import net.minecraft.world.level.LightLayer;
import net.minecraft.world.level.chunk.LevelChunkSection;
import net.minecraft.world.level.material.Fluid;
import net.minecraft.world.level.material.Fluids;
import net.minecraft.world.level.block.Block;
import net.minecraft.world.level.chunk.ProtoTickList;
import net.minecraft.world.level.chunk.UpgradeData;
import net.minecraft.world.level.chunk.ChunkBiomeContainer;
import java.util.Objects;
import net.minecraft.world.level.chunk.ProtoChunk;
import net.minecraft.nbt.CompoundTag;
import net.minecraft.world.level.ChunkPos;
import net.minecraft.world.entity.ai.village.poi.PoiManager;
import net.minecraft.world.level.levelgen.structure.templatesystem.StructureManager;
import net.minecraft.server.level.ServerLevel;
import org.apache.logging.log4j.Logger;

public class ChunkSerializer {
    private static final Logger LOGGER;
    
    public static ProtoChunk read(final ServerLevel zd, final StructureManager cva, final PoiManager axz, final ChunkPos bph, final CompoundTag le) {
        final ChunkGenerator cha6 = zd.getChunkSource().getGenerator();
        final BiomeSource brh7 = cha6.getBiomeSource();
        final CompoundTag le2 = le.getCompound("Level");
        final ChunkPos bph2 = new ChunkPos(le2.getInt("xPos"), le2.getInt("zPos"));
        if (!Objects.equals(bph, bph2)) {
            ChunkSerializer.LOGGER.error("Chunk file at {} is in the wrong location; relocating. (Expected {}, got {})", bph, bph, bph2);
        }
        final ChunkBiomeContainer cgz10 = new ChunkBiomeContainer(bph, brh7, (int[])(le2.contains("Biomes", 11) ? le2.getIntArray("Biomes") : null));
        final UpgradeData cht11 = le2.contains("UpgradeData", 10) ? new UpgradeData(le2.getCompound("UpgradeData")) : UpgradeData.EMPTY;
        final ProtoTickList<Block> chs12 = new ProtoTickList<Block>(bvr -> bvr == null || bvr.defaultBlockState().isAir(), bph, le2.getList("ToBeTicked", 9));
        final ProtoTickList<Fluid> chs13 = new ProtoTickList<Fluid>(cwz -> cwz == null || cwz == Fluids.EMPTY, bph, le2.getList("LiquidsToBeTicked", 9));
        final boolean boolean14 = le2.getBoolean("isLightOn");
        final ListTag lk15 = le2.getList("Sections", 10);
        final int integer16 = 16;
        final LevelChunkSection[] arr17 = new LevelChunkSection[16];
        final boolean boolean15 = zd.dimensionType().hasSkyLight();
        final ChunkSource chb19 = zd.getChunkSource();
        final LevelLightEngine cwr20 = chb19.getLightEngine();
        if (boolean14) {
            cwr20.retainData(bph, true);
        }
        for (int integer17 = 0; integer17 < lk15.size(); ++integer17) {
            final CompoundTag le3 = lk15.getCompound(integer17);
            final int integer18 = le3.getByte("Y");
            if (le3.contains("Palette", 9) && le3.contains("BlockStates", 12)) {
                final LevelChunkSection chk24 = new LevelChunkSection(integer18 << 4);
                chk24.getStates().read(le3.getList("Palette", 10), le3.getLongArray("BlockStates"));
                chk24.recalcBlockCounts();
                if (!chk24.isEmpty()) {
                    arr17[integer18] = chk24;
                }
                axz.checkConsistencyWithBlocks(bph, chk24);
            }
            if (boolean14) {
                if (le3.contains("BlockLight", 7)) {
                    cwr20.queueSectionData(LightLayer.BLOCK, SectionPos.of(bph, integer18), new DataLayer(le3.getByteArray("BlockLight")), true);
                }
                if (boolean15 && le3.contains("SkyLight", 7)) {
                    cwr20.queueSectionData(LightLayer.SKY, SectionPos.of(bph, integer18), new DataLayer(le3.getByteArray("SkyLight")), true);
                }
            }
        }
        final long long21 = le2.getLong("InhabitedTime");
        final ChunkStatus.ChunkType a23 = getChunkTypeFromTag(le);
        ChunkAccess cgy24;
        if (a23 == ChunkStatus.ChunkType.LEVELCHUNK) {
            TickList<Block> bqr25;
            if (le2.contains("TileTicks", 9)) {
                bqr25 = ChunkTickList.create(le2.getList("TileTicks", 10), Registry.BLOCK::getKey, Registry.BLOCK::get);
            }
            else {
                bqr25 = chs12;
            }
            TickList<Fluid> bqr26;
            if (le2.contains("LiquidTicks", 9)) {
                bqr26 = ChunkTickList.create(le2.getList("LiquidTicks", 10), Registry.FLUID::getKey, Registry.FLUID::get);
            }
            else {
                bqr26 = chs13;
            }
            cgy24 = new LevelChunk(zd.getLevel(), bph, cgz10, cht11, bqr25, bqr26, long21, arr17, chj -> postLoadChunk(le2, chj));
        }
        else {
            final ProtoChunk chr25 = new ProtoChunk(bph, cht11, arr17, chs12, chs13);
            chr25.setBiomes(cgz10);
            cgy24 = chr25;
            cgy24.setInhabitedTime(long21);
            chr25.setStatus(ChunkStatus.byName(le2.getString("Status")));
            if (cgy24.getStatus().isOrAfter(ChunkStatus.FEATURES)) {
                chr25.setLightEngine(cwr20);
            }
            if (!boolean14 && cgy24.getStatus().isOrAfter(ChunkStatus.LIGHT)) {
                for (final BlockPos fu27 : BlockPos.betweenClosed(bph.getMinBlockX(), 0, bph.getMinBlockZ(), bph.getMaxBlockX(), 255, bph.getMaxBlockZ())) {
                    if (cgy24.getBlockState(fu27).getLightEmission() != 0) {
                        chr25.addLight(fu27);
                    }
                }
            }
        }
        cgy24.setLightCorrect(boolean14);
        final CompoundTag le4 = le2.getCompound("Heightmaps");
        final EnumSet<Heightmap.Types> enumSet26 = EnumSet.<Heightmap.Types>noneOf(Heightmap.Types.class);
        for (final Heightmap.Types a24 : cgy24.getStatus().heightmapsAfter()) {
            final String string29 = a24.getSerializationKey();
            if (le4.contains(string29, 12)) {
                cgy24.setHeightmap(a24, le4.getLongArray(string29));
            }
            else {
                enumSet26.add(a24);
            }
        }
        Heightmap.primeHeightmaps(cgy24, enumSet26);
        final CompoundTag le5 = le2.getCompound("Structures");
        cgy24.setAllStarts(unpackStructureStart(cva, le5, zd.getSeed()));
        cgy24.setAllReferences(unpackStructureReferences(bph, le5));
        if (le2.getBoolean("shouldSave")) {
            cgy24.setUnsaved(true);
        }
        final ListTag lk16 = le2.getList("PostProcessing", 9);
        for (int integer19 = 0; integer19 < lk16.size(); ++integer19) {
            final ListTag lk17 = lk16.getList(integer19);
            for (int integer20 = 0; integer20 < lk17.size(); ++integer20) {
                cgy24.addPackedPostProcess(lk17.getShort(integer20), integer19);
            }
        }
        if (a23 == ChunkStatus.ChunkType.LEVELCHUNK) {
            return new ImposterProtoChunk((LevelChunk)cgy24);
        }
        final ProtoChunk chr26 = (ProtoChunk)cgy24;
        final ListTag lk17 = le2.getList("Entities", 10);
        for (int integer20 = 0; integer20 < lk17.size(); ++integer20) {
            chr26.addEntity(lk17.getCompound(integer20));
        }
        final ListTag lk18 = le2.getList("TileEntities", 10);
        for (int integer21 = 0; integer21 < lk18.size(); ++integer21) {
            final CompoundTag le6 = lk18.getCompound(integer21);
            cgy24.setBlockEntityNbt(le6);
        }
        final ListTag lk19 = le2.getList("Lights", 9);
        for (int integer22 = 0; integer22 < lk19.size(); ++integer22) {
            final ListTag lk20 = lk19.getList(integer22);
            for (int integer23 = 0; integer23 < lk20.size(); ++integer23) {
                chr26.addLight(lk20.getShort(integer23), integer22);
            }
        }
        final CompoundTag le6 = le2.getCompound("CarvingMasks");
        for (final String string30 : le6.getAllKeys()) {
            final GenerationStep.Carving a25 = GenerationStep.Carving.valueOf(string30);
            chr26.setCarvingMask(a25, BitSet.valueOf(le6.getByteArray(string30)));
        }
        return chr26;
    }
    
    public static CompoundTag write(final ServerLevel zd, final ChunkAccess cgy) {
        final ChunkPos bph3 = cgy.getPos();
        final CompoundTag le4 = new CompoundTag();
        final CompoundTag le5 = new CompoundTag();
        le4.putInt("DataVersion", SharedConstants.getCurrentVersion().getWorldVersion());
        le4.put("Level", le5);
        le5.putInt("xPos", bph3.x);
        le5.putInt("zPos", bph3.z);
        le5.putLong("LastUpdate", zd.getGameTime());
        le5.putLong("InhabitedTime", cgy.getInhabitedTime());
        le5.putString("Status", cgy.getStatus().getName());
        final UpgradeData cht6 = cgy.getUpgradeData();
        if (!cht6.isEmpty()) {
            le5.put("UpgradeData", cht6.write());
        }
        final LevelChunkSection[] arr7 = cgy.getSections();
        final ListTag lk8 = new ListTag();
        final LevelLightEngine cwr9 = zd.getChunkSource().getLightEngine();
        final boolean boolean10 = cgy.isLightCorrect();
        for (int integer11 = -1; integer11 < 17; ++integer11) {
            final int integer12 = integer11;
            final int n;
            final LevelChunkSection chk2 = Arrays.<LevelChunkSection>stream(arr7).filter(chk -> chk != null && chk.bottomBlockY() >> 4 == n).findFirst().orElse(LevelChunk.EMPTY_SECTION);
            final DataLayer chd14 = cwr9.getLayerListener(LightLayer.BLOCK).getDataLayerData(SectionPos.of(bph3, integer12));
            final DataLayer chd15 = cwr9.getLayerListener(LightLayer.SKY).getDataLayerData(SectionPos.of(bph3, integer12));
            if (chk2 != LevelChunk.EMPTY_SECTION || chd14 != null || chd15 != null) {
                final CompoundTag le6 = new CompoundTag();
                le6.putByte("Y", (byte)(integer12 & 0xFF));
                if (chk2 != LevelChunk.EMPTY_SECTION) {
                    chk2.getStates().write(le6, "Palette", "BlockStates");
                }
                if (chd14 != null && !chd14.isEmpty()) {
                    le6.putByteArray("BlockLight", chd14.getData());
                }
                if (chd15 != null && !chd15.isEmpty()) {
                    le6.putByteArray("SkyLight", chd15.getData());
                }
                ((AbstractList<CompoundTag>)lk8).add(le6);
            }
        }
        le5.put("Sections", lk8);
        if (boolean10) {
            le5.putBoolean("isLightOn", true);
        }
        final ChunkBiomeContainer cgz11 = cgy.getBiomes();
        if (cgz11 != null) {
            le5.putIntArray("Biomes", cgz11.writeBiomes());
        }
        final ListTag lk9 = new ListTag();
        for (final BlockPos fu14 : cgy.getBlockEntitiesPos()) {
            final CompoundTag le7 = cgy.getBlockEntityNbtForSaving(fu14);
            if (le7 != null) {
                ((AbstractList<CompoundTag>)lk9).add(le7);
            }
        }
        le5.put("TileEntities", lk9);
        final ListTag lk10 = new ListTag();
        if (cgy.getStatus().getChunkType() == ChunkStatus.ChunkType.LEVELCHUNK) {
            final LevelChunk chj14 = (LevelChunk)cgy;
            chj14.setLastSaveHadEntities(false);
            for (int integer13 = 0; integer13 < chj14.getEntitySections().length; ++integer13) {
                for (final Entity aom17 : chj14.getEntitySections()[integer13]) {
                    final CompoundTag le8 = new CompoundTag();
                    if (aom17.save(le8)) {
                        chj14.setLastSaveHadEntities(true);
                        ((AbstractList<CompoundTag>)lk10).add(le8);
                    }
                }
            }
        }
        else {
            final ProtoChunk chr14 = (ProtoChunk)cgy;
            lk10.addAll(chr14.getEntities());
            le5.put("Lights", packOffsets(chr14.getPackedLights()));
            final CompoundTag le7 = new CompoundTag();
            for (final GenerationStep.Carving a19 : GenerationStep.Carving.values()) {
                final BitSet bitSet20 = chr14.getCarvingMask(a19);
                if (bitSet20 != null) {
                    le7.putByteArray(a19.toString(), bitSet20.toByteArray());
                }
            }
            le5.put("CarvingMasks", le7);
        }
        le5.put("Entities", lk10);
        final TickList<Block> bqr14 = cgy.getBlockTicks();
        if (bqr14 instanceof ProtoTickList) {
            le5.put("ToBeTicked", ((ProtoTickList)bqr14).save());
        }
        else if (bqr14 instanceof ChunkTickList) {
            le5.put("TileTicks", ((ChunkTickList)bqr14).save());
        }
        else {
            le5.put("TileTicks", zd.getBlockTicks().save(bph3));
        }
        final TickList<Fluid> bqr15 = cgy.getLiquidTicks();
        if (bqr15 instanceof ProtoTickList) {
            le5.put("LiquidsToBeTicked", ((ProtoTickList)bqr15).save());
        }
        else if (bqr15 instanceof ChunkTickList) {
            le5.put("LiquidTicks", ((ChunkTickList)bqr15).save());
        }
        else {
            le5.put("LiquidTicks", zd.getLiquidTicks().save(bph3));
        }
        le5.put("PostProcessing", packOffsets(cgy.getPostProcessing()));
        final CompoundTag le6 = new CompoundTag();
        for (final Map.Entry<Heightmap.Types, Heightmap> entry18 : cgy.getHeightmaps()) {
            if (cgy.getStatus().heightmapsAfter().contains(entry18.getKey())) {
                le6.put(entry18.getKey().getSerializationKey(), new LongArrayTag(entry18.getValue().getRawData()));
            }
        }
        le5.put("Heightmaps", le6);
        le5.put("Structures", packStructureData(bph3, cgy.getAllStarts(), cgy.getAllReferences()));
        return le4;
    }
    
    public static ChunkStatus.ChunkType getChunkTypeFromTag(@Nullable final CompoundTag le) {
        if (le != null) {
            final ChunkStatus chc2 = ChunkStatus.byName(le.getCompound("Level").getString("Status"));
            if (chc2 != null) {
                return chc2.getChunkType();
            }
        }
        return ChunkStatus.ChunkType.PROTOCHUNK;
    }
    
    private static void postLoadChunk(final CompoundTag le, final LevelChunk chj) {
        final ListTag lk3 = le.getList("Entities", 10);
        final Level bqb4 = chj.getLevel();
        for (int integer5 = 0; integer5 < lk3.size(); ++integer5) {
            final CompoundTag le2 = lk3.getCompound(integer5);
            EntityType.loadEntityRecursive(le2, bqb4, aom -> {
                chj.addEntity(aom);
                return aom;
            });
            chj.setLastSaveHadEntities(true);
        }
        final ListTag lk4 = le.getList("TileEntities", 10);
        for (int integer6 = 0; integer6 < lk4.size(); ++integer6) {
            final CompoundTag le3 = lk4.getCompound(integer6);
            final boolean boolean8 = le3.getBoolean("keepPacked");
            if (boolean8) {
                chj.setBlockEntityNbt(le3);
            }
            else {
                final BlockPos fu9 = new BlockPos(le3.getInt("x"), le3.getInt("y"), le3.getInt("z"));
                final BlockEntity cdl10 = BlockEntity.loadStatic(chj.getBlockState(fu9), le3);
                if (cdl10 != null) {
                    chj.addBlockEntity(cdl10);
                }
            }
        }
    }
    
    private static CompoundTag packStructureData(final ChunkPos bph, final Map<StructureFeature<?>, StructureStart<?>> map2, final Map<StructureFeature<?>, LongSet> map3) {
        final CompoundTag le4 = new CompoundTag();
        final CompoundTag le5 = new CompoundTag();
        for (final Map.Entry<StructureFeature<?>, StructureStart<?>> entry7 : map2.entrySet()) {
            le5.put(entry7.getKey().getFeatureName(), entry7.getValue().createTag(bph.x, bph.z));
        }
        le4.put("Starts", le5);
        final CompoundTag le6 = new CompoundTag();
        for (final Map.Entry<StructureFeature<?>, LongSet> entry8 : map3.entrySet()) {
            le6.put(entry8.getKey().getFeatureName(), new LongArrayTag(entry8.getValue()));
        }
        le4.put("References", le6);
        return le4;
    }
    
    private static Map<StructureFeature<?>, StructureStart<?>> unpackStructureStart(final StructureManager cva, final CompoundTag le, final long long3) {
        final Map<StructureFeature<?>, StructureStart<?>> map5 = Maps.newHashMap();
        final CompoundTag le2 = le.getCompound("Starts");
        for (final String string8 : le2.getAllKeys()) {
            final String string9 = string8.toLowerCase(Locale.ROOT);
            final StructureFeature<?> cml10 = StructureFeature.STRUCTURES_REGISTRY.get(string9);
            if (cml10 == null) {
                ChunkSerializer.LOGGER.error("Unknown structure start: {}", string9);
            }
            else {
                final StructureStart<?> ctz11 = StructureFeature.loadStaticStart(cva, le2.getCompound(string8), long3);
                if (ctz11 == null) {
                    continue;
                }
                map5.put(cml10, ctz11);
            }
        }
        return map5;
    }
    
    private static Map<StructureFeature<?>, LongSet> unpackStructureReferences(final ChunkPos bph, final CompoundTag le) {
        final Map<StructureFeature<?>, LongSet> map3 = Maps.newHashMap();
        final CompoundTag le2 = le.getCompound("References");
        for (final String string6 : le2.getAllKeys()) {
            final ChunkPos bph2;
            final Object o;
            map3.put(StructureFeature.STRUCTURES_REGISTRY.get(string6.toLowerCase(Locale.ROOT)), (LongSet)new LongOpenHashSet(Arrays.stream(le2.getLongArray(string6)).filter(long3 -> {
                bph2 = new ChunkPos(long3);
                if (bph2.getChessboardDistance(bph) > 8) {
                    ChunkSerializer.LOGGER.warn("Found invalid structure reference [ {} @ {} ] for chunk {}.", o, bph2, bph);
                    return false;
                }
                else {
                    return true;
                }
            }).toArray()));
        }
        return map3;
    }
    
    public static ListTag packOffsets(final ShortList[] arr) {
        final ListTag lk2 = new ListTag();
        for (final ShortList shortList6 : arr) {
            final ListTag lk3 = new ListTag();
            if (shortList6 != null) {
                for (final Short short9 : shortList6) {
                    ((AbstractList<ShortTag>)lk3).add(ShortTag.valueOf(short9));
                }
            }
            ((AbstractList<ListTag>)lk2).add(lk3);
        }
        return lk2;
    }
    
    static {
        LOGGER = LogManager.getLogger();
    }
}
