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 net.minecraft.world.level.levelgen.structure.StructureFeatureIO;
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.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 xd, final StructureManager cmo, final PoiManager asr, final ChunkPos bjh, final CompoundTag jt) {
        final ChunkGenerator<?> caa6 = xd.getChunkSource().getGenerator();
        final BiomeSource bkw7 = caa6.getBiomeSource();
        final CompoundTag jt2 = jt.getCompound("Level");
        final ChunkPos bjh2 = new ChunkPos(jt2.getInt("xPos"), jt2.getInt("zPos"));
        if (!Objects.equals(bjh, bjh2)) {
            ChunkSerializer.LOGGER.error("Chunk file at {} is in the wrong location; relocating. (Expected {}, got {})", bjh, bjh, bjh2);
        }
        final ChunkBiomeContainer bzz10 = new ChunkBiomeContainer(bjh, bkw7, (int[])(jt2.contains("Biomes", 11) ? jt2.getIntArray("Biomes") : null));
        final UpgradeData cav11 = jt2.contains("UpgradeData", 10) ? new UpgradeData(jt2.getCompound("UpgradeData")) : UpgradeData.EMPTY;
        final ProtoTickList<Block> cau12 = new ProtoTickList<Block>(bph -> bph == null || bph.defaultBlockState().isAir(), bjh, jt2.getList("ToBeTicked", 9));
        final ProtoTickList<Fluid> cau13 = new ProtoTickList<Fluid>(coi -> coi == null || coi == Fluids.EMPTY, bjh, jt2.getList("LiquidsToBeTicked", 9));
        final boolean boolean14 = jt2.getBoolean("isLightOn");
        final ListTag jz15 = jt2.getList("Sections", 10);
        final int integer16 = 16;
        final LevelChunkSection[] arr17 = new LevelChunkSection[16];
        final boolean boolean15 = xd.getDimension().isHasSkyLight();
        final ChunkSource cad19 = xd.getChunkSource();
        final LevelLightEngine coa20 = cad19.getLightEngine();
        if (boolean14) {
            coa20.retainData(bjh, true);
        }
        for (int integer17 = 0; integer17 < jz15.size(); ++integer17) {
            final CompoundTag jt3 = jz15.getCompound(integer17);
            final int integer18 = jt3.getByte("Y");
            if (jt3.contains("Palette", 9) && jt3.contains("BlockStates", 12)) {
                final LevelChunkSection cam24 = new LevelChunkSection(integer18 << 4);
                cam24.getStates().read(jt3.getList("Palette", 10), jt3.getLongArray("BlockStates"));
                cam24.recalcBlockCounts();
                if (!cam24.isEmpty()) {
                    arr17[integer18] = cam24;
                }
                asr.checkConsistencyWithBlocks(bjh, cam24);
            }
            if (boolean14) {
                if (jt3.contains("BlockLight", 7)) {
                    coa20.queueSectionData(LightLayer.BLOCK, SectionPos.of(bjh, integer18), new DataLayer(jt3.getByteArray("BlockLight")));
                }
                if (boolean15 && jt3.contains("SkyLight", 7)) {
                    coa20.queueSectionData(LightLayer.SKY, SectionPos.of(bjh, integer18), new DataLayer(jt3.getByteArray("SkyLight")));
                }
            }
        }
        final long long21 = jt2.getLong("InhabitedTime");
        final ChunkStatus.ChunkType a23 = getChunkTypeFromTag(jt);
        ChunkAccess bzy24;
        if (a23 == ChunkStatus.ChunkType.LEVELCHUNK) {
            TickList<Block> bkl25;
            if (jt2.contains("TileTicks", 9)) {
                bkl25 = ChunkTickList.create(jt2.getList("TileTicks", 10), Registry.BLOCK::getKey, Registry.BLOCK::get);
            }
            else {
                bkl25 = cau12;
            }
            TickList<Fluid> bkl26;
            if (jt2.contains("LiquidTicks", 9)) {
                bkl26 = ChunkTickList.create(jt2.getList("LiquidTicks", 10), Registry.FLUID::getKey, Registry.FLUID::get);
            }
            else {
                bkl26 = cau13;
            }
            bzy24 = new LevelChunk(xd.getLevel(), bjh, bzz10, cav11, bkl25, bkl26, long21, arr17, cal -> postLoadChunk(jt2, cal));
        }
        else {
            final ProtoChunk cat25 = new ProtoChunk(bjh, cav11, arr17, cau12, cau13);
            cat25.setBiomes(bzz10);
            bzy24 = cat25;
            bzy24.setInhabitedTime(long21);
            cat25.setStatus(ChunkStatus.byName(jt2.getString("Status")));
            if (bzy24.getStatus().isOrAfter(ChunkStatus.FEATURES)) {
                cat25.setLightEngine(coa20);
            }
            if (!boolean14 && bzy24.getStatus().isOrAfter(ChunkStatus.LIGHT)) {
                for (final BlockPos fk27 : BlockPos.betweenClosed(bjh.getMinBlockX(), 0, bjh.getMinBlockZ(), bjh.getMaxBlockX(), 255, bjh.getMaxBlockZ())) {
                    if (bzy24.getBlockState(fk27).getLightEmission() != 0) {
                        cat25.addLight(fk27);
                    }
                }
            }
        }
        bzy24.setLightCorrect(boolean14);
        final CompoundTag jt4 = jt2.getCompound("Heightmaps");
        final EnumSet<Heightmap.Types> enumSet26 = EnumSet.<Heightmap.Types>noneOf(Heightmap.Types.class);
        for (final Heightmap.Types a24 : bzy24.getStatus().heightmapsAfter()) {
            final String string29 = a24.getSerializationKey();
            if (jt4.contains(string29, 12)) {
                bzy24.setHeightmap(a24, jt4.getLongArray(string29));
            }
            else {
                enumSet26.add(a24);
            }
        }
        Heightmap.primeHeightmaps(bzy24, enumSet26);
        final CompoundTag jt5 = jt2.getCompound("Structures");
        bzy24.setAllStarts(unpackStructureStart(caa6, cmo, jt5));
        bzy24.setAllReferences(unpackStructureReferences(bjh, jt5));
        if (jt2.getBoolean("shouldSave")) {
            bzy24.setUnsaved(true);
        }
        final ListTag jz16 = jt2.getList("PostProcessing", 9);
        for (int integer19 = 0; integer19 < jz16.size(); ++integer19) {
            final ListTag jz17 = jz16.getList(integer19);
            for (int integer20 = 0; integer20 < jz17.size(); ++integer20) {
                bzy24.addPackedPostProcess(jz17.getShort(integer20), integer19);
            }
        }
        if (a23 == ChunkStatus.ChunkType.LEVELCHUNK) {
            return new ImposterProtoChunk((LevelChunk)bzy24);
        }
        final ProtoChunk cat26 = (ProtoChunk)bzy24;
        final ListTag jz17 = jt2.getList("Entities", 10);
        for (int integer20 = 0; integer20 < jz17.size(); ++integer20) {
            cat26.addEntity(jz17.getCompound(integer20));
        }
        final ListTag jz18 = jt2.getList("TileEntities", 10);
        for (int integer21 = 0; integer21 < jz18.size(); ++integer21) {
            final CompoundTag jt6 = jz18.getCompound(integer21);
            bzy24.setBlockEntityNbt(jt6);
        }
        final ListTag jz19 = jt2.getList("Lights", 9);
        for (int integer22 = 0; integer22 < jz19.size(); ++integer22) {
            final ListTag jz20 = jz19.getList(integer22);
            for (int integer23 = 0; integer23 < jz20.size(); ++integer23) {
                cat26.addLight(jz20.getShort(integer23), integer22);
            }
        }
        final CompoundTag jt6 = jt2.getCompound("CarvingMasks");
        for (final String string30 : jt6.getAllKeys()) {
            final GenerationStep.Carving a25 = GenerationStep.Carving.valueOf(string30);
            cat26.setCarvingMask(a25, BitSet.valueOf(jt6.getByteArray(string30)));
        }
        return cat26;
    }
    
    public static CompoundTag write(final ServerLevel xd, final ChunkAccess bzy) {
        final ChunkPos bjh3 = bzy.getPos();
        final CompoundTag jt4 = new CompoundTag();
        final CompoundTag jt5 = new CompoundTag();
        jt4.putInt("DataVersion", SharedConstants.getCurrentVersion().getWorldVersion());
        jt4.put("Level", jt5);
        jt5.putInt("xPos", bjh3.x);
        jt5.putInt("zPos", bjh3.z);
        jt5.putLong("LastUpdate", xd.getGameTime());
        jt5.putLong("InhabitedTime", bzy.getInhabitedTime());
        jt5.putString("Status", bzy.getStatus().getName());
        final UpgradeData cav6 = bzy.getUpgradeData();
        if (!cav6.isEmpty()) {
            jt5.put("UpgradeData", cav6.write());
        }
        final LevelChunkSection[] arr7 = bzy.getSections();
        final ListTag jz8 = new ListTag();
        final LevelLightEngine coa9 = xd.getChunkSource().getLightEngine();
        final boolean boolean10 = bzy.isLightCorrect();
        for (int integer11 = -1; integer11 < 17; ++integer11) {
            final int integer12 = integer11;
            final int n;
            final LevelChunkSection cam2 = Arrays.<LevelChunkSection>stream(arr7).filter(cam -> cam != null && cam.bottomBlockY() >> 4 == n).findFirst().orElse(LevelChunk.EMPTY_SECTION);
            final DataLayer caf14 = coa9.getLayerListener(LightLayer.BLOCK).getDataLayerData(SectionPos.of(bjh3, integer12));
            final DataLayer caf15 = coa9.getLayerListener(LightLayer.SKY).getDataLayerData(SectionPos.of(bjh3, integer12));
            if (cam2 != LevelChunk.EMPTY_SECTION || caf14 != null || caf15 != null) {
                final CompoundTag jt6 = new CompoundTag();
                jt6.putByte("Y", (byte)(integer12 & 0xFF));
                if (cam2 != LevelChunk.EMPTY_SECTION) {
                    cam2.getStates().write(jt6, "Palette", "BlockStates");
                }
                if (caf14 != null && !caf14.isEmpty()) {
                    jt6.putByteArray("BlockLight", caf14.getData());
                }
                if (caf15 != null && !caf15.isEmpty()) {
                    jt6.putByteArray("SkyLight", caf15.getData());
                }
                ((AbstractList<CompoundTag>)jz8).add(jt6);
            }
        }
        jt5.put("Sections", jz8);
        if (boolean10) {
            jt5.putBoolean("isLightOn", true);
        }
        final ChunkBiomeContainer bzz11 = bzy.getBiomes();
        if (bzz11 != null) {
            jt5.putIntArray("Biomes", bzz11.writeBiomes());
        }
        final ListTag jz9 = new ListTag();
        for (final BlockPos fk14 : bzy.getBlockEntitiesPos()) {
            final CompoundTag jt7 = bzy.getBlockEntityNbtForSaving(fk14);
            if (jt7 != null) {
                ((AbstractList<CompoundTag>)jz9).add(jt7);
            }
        }
        jt5.put("TileEntities", jz9);
        final ListTag jz10 = new ListTag();
        if (bzy.getStatus().getChunkType() == ChunkStatus.ChunkType.LEVELCHUNK) {
            final LevelChunk cal14 = (LevelChunk)bzy;
            cal14.setLastSaveHadEntities(false);
            for (int integer13 = 0; integer13 < cal14.getEntitySections().length; ++integer13) {
                for (final Entity akq17 : cal14.getEntitySections()[integer13]) {
                    final CompoundTag jt8 = new CompoundTag();
                    if (akq17.save(jt8)) {
                        cal14.setLastSaveHadEntities(true);
                        ((AbstractList<CompoundTag>)jz10).add(jt8);
                    }
                }
            }
        }
        else {
            final ProtoChunk cat14 = (ProtoChunk)bzy;
            jz10.addAll(cat14.getEntities());
            jt5.put("Lights", packOffsets(cat14.getPackedLights()));
            final CompoundTag jt7 = new CompoundTag();
            for (final GenerationStep.Carving a19 : GenerationStep.Carving.values()) {
                jt7.putByteArray(a19.toString(), bzy.getCarvingMask(a19).toByteArray());
            }
            jt5.put("CarvingMasks", jt7);
        }
        jt5.put("Entities", jz10);
        final TickList<Block> bkl14 = bzy.getBlockTicks();
        if (bkl14 instanceof ProtoTickList) {
            jt5.put("ToBeTicked", ((ProtoTickList)bkl14).save());
        }
        else if (bkl14 instanceof ChunkTickList) {
            jt5.put("TileTicks", ((ChunkTickList)bkl14).save(xd.getGameTime()));
        }
        else {
            jt5.put("TileTicks", xd.getBlockTicks().save(bjh3));
        }
        final TickList<Fluid> bkl15 = bzy.getLiquidTicks();
        if (bkl15 instanceof ProtoTickList) {
            jt5.put("LiquidsToBeTicked", ((ProtoTickList)bkl15).save());
        }
        else if (bkl15 instanceof ChunkTickList) {
            jt5.put("LiquidTicks", ((ChunkTickList)bkl15).save(xd.getGameTime()));
        }
        else {
            jt5.put("LiquidTicks", xd.getLiquidTicks().save(bjh3));
        }
        jt5.put("PostProcessing", packOffsets(bzy.getPostProcessing()));
        final CompoundTag jt6 = new CompoundTag();
        for (final Map.Entry<Heightmap.Types, Heightmap> entry18 : bzy.getHeightmaps()) {
            if (bzy.getStatus().heightmapsAfter().contains(entry18.getKey())) {
                jt6.put(entry18.getKey().getSerializationKey(), new LongArrayTag(entry18.getValue().getRawData()));
            }
        }
        jt5.put("Heightmaps", jt6);
        jt5.put("Structures", packStructureData(bjh3, bzy.getAllStarts(), bzy.getAllReferences()));
        return jt4;
    }
    
    public static ChunkStatus.ChunkType getChunkTypeFromTag(@Nullable final CompoundTag jt) {
        if (jt != null) {
            final ChunkStatus cae2 = ChunkStatus.byName(jt.getCompound("Level").getString("Status"));
            if (cae2 != null) {
                return cae2.getChunkType();
            }
        }
        return ChunkStatus.ChunkType.PROTOCHUNK;
    }
    
    private static void postLoadChunk(final CompoundTag jt, final LevelChunk cal) {
        final ListTag jz3 = jt.getList("Entities", 10);
        final Level bjw4 = cal.getLevel();
        for (int integer5 = 0; integer5 < jz3.size(); ++integer5) {
            final CompoundTag jt2 = jz3.getCompound(integer5);
            EntityType.loadEntityRecursive(jt2, bjw4, akq -> {
                cal.addEntity(akq);
                return akq;
            });
            cal.setLastSaveHadEntities(true);
        }
        final ListTag jz4 = jt.getList("TileEntities", 10);
        for (int integer6 = 0; integer6 < jz4.size(); ++integer6) {
            final CompoundTag jt3 = jz4.getCompound(integer6);
            final boolean boolean8 = jt3.getBoolean("keepPacked");
            if (boolean8) {
                cal.setBlockEntityNbt(jt3);
            }
            else {
                final BlockEntity bwl9 = BlockEntity.loadStatic(jt3);
                if (bwl9 != null) {
                    cal.addBlockEntity(bwl9);
                }
            }
        }
    }
    
    private static CompoundTag packStructureData(final ChunkPos bjh, final Map<String, StructureStart> map2, final Map<String, LongSet> map3) {
        final CompoundTag jt4 = new CompoundTag();
        final CompoundTag jt5 = new CompoundTag();
        for (final Map.Entry<String, StructureStart> entry7 : map2.entrySet()) {
            jt5.put(entry7.getKey(), entry7.getValue().createTag(bjh.x, bjh.z));
        }
        jt4.put("Starts", jt5);
        final CompoundTag jt6 = new CompoundTag();
        for (final Map.Entry<String, LongSet> entry8 : map3.entrySet()) {
            jt6.put(entry8.getKey(), new LongArrayTag(entry8.getValue()));
        }
        jt4.put("References", jt6);
        return jt4;
    }
    
    private static Map<String, StructureStart> unpackStructureStart(final ChunkGenerator<?> caa, final StructureManager cmo, final CompoundTag jt) {
        final Map<String, StructureStart> map4 = Maps.newHashMap();
        final CompoundTag jt2 = jt.getCompound("Starts");
        for (final String string7 : jt2.getAllKeys()) {
            map4.put(string7, StructureFeatureIO.loadStaticStart(caa, cmo, jt2.getCompound(string7)));
        }
        return map4;
    }
    
    private static Map<String, LongSet> unpackStructureReferences(final ChunkPos bjh, final CompoundTag jt) {
        final Map<String, LongSet> map3 = Maps.newHashMap();
        final CompoundTag jt2 = jt.getCompound("References");
        for (final String string6 : jt2.getAllKeys()) {
            final ChunkPos bjh2;
            final Object o;
            map3.put(string6, (LongSet)new LongOpenHashSet(Arrays.stream(jt2.getLongArray(string6)).filter(long3 -> {
                bjh2 = new ChunkPos(long3);
                if (bjh2.getChessboardDistance(bjh) > 8) {
                    ChunkSerializer.LOGGER.warn("Found invalid structure reference [ {} @ {} ] for chunk {}.", o, bjh2, bjh);
                    return false;
                }
                else {
                    return true;
                }
            }).toArray()));
        }
        return map3;
    }
    
    public static ListTag packOffsets(final ShortList[] arr) {
        final ListTag jz2 = new ListTag();
        for (final ShortList shortList6 : arr) {
            final ListTag jz3 = new ListTag();
            if (shortList6 != null) {
                for (final Short short9 : shortList6) {
                    ((AbstractList<ShortTag>)jz3).add(ShortTag.valueOf(short9));
                }
            }
            ((AbstractList<ListTag>)jz2).add(jz3);
        }
        return jz2;
    }
    
    static {
        LOGGER = LogManager.getLogger();
    }
}
