package net.minecraft.network.protocol.game;

import com.google.common.collect.Lists;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;
import java.util.Iterator;
import java.util.List;
import java.util.Map.Entry;
import java.util.function.Consumer;
import javax.annotation.Nullable;
import net.minecraft.core.BlockPos;
import net.minecraft.core.Registry;
import net.minecraft.core.SectionPos;
import net.minecraft.core.BlockPos.MutableBlockPos;
import net.minecraft.nbt.CompoundTag;
import net.minecraft.nbt.LongArrayTag;
import net.minecraft.network.FriendlyByteBuf;
import net.minecraft.world.level.block.entity.BlockEntity;
import net.minecraft.world.level.block.entity.BlockEntityType;
import net.minecraft.world.level.chunk.LevelChunk;
import net.minecraft.world.level.chunk.LevelChunkSection;
import net.minecraft.world.level.levelgen.Heightmap;
import net.minecraft.world.level.levelgen.Heightmap.Types;

public class ClientboundLevelChunkPacketData {
    private static final int TWO_MEGABYTES = 2097152;
    private final CompoundTag heightmaps;
    private final byte[] buffer;
    private final List<ClientboundLevelChunkPacketData.BlockEntityInfo> blockEntitiesData;

    public ClientboundLevelChunkPacketData(LevelChunk var0) {
        this.heightmaps = new CompoundTag();
        Iterator var2 = var0.getHeightmaps().iterator();

        Entry var2;
        while(var2.hasNext()) {
            var2 = (Entry)var2.next();
            if (((Types)var2.getKey()).sendToClient()) {
                this.heightmaps.put(((Types)var2.getKey()).getSerializationKey(), new LongArrayTag(((Heightmap)var2.getValue()).getRawData()));
            }
        }

        this.buffer = new byte[calculateChunkSize(var0)];
        extractChunkData(new FriendlyByteBuf(this.getWriteBuffer()), var0);
        this.blockEntitiesData = Lists.newArrayList();
        var2 = var0.getBlockEntities().entrySet().iterator();

        while(var2.hasNext()) {
            var2 = (Entry)var2.next();
            this.blockEntitiesData.add(ClientboundLevelChunkPacketData.BlockEntityInfo.create((BlockEntity)var2.getValue()));
        }

    }

    public ClientboundLevelChunkPacketData(FriendlyByteBuf var0, int var1, int var2) {
        this.heightmaps = var0.readNbt();
        if (this.heightmaps == null) {
            throw new RuntimeException("Can't read heightmap in packet for [" + var1 + ", " + var2 + "]");
        } else {
            int var3 = var0.readVarInt();
            if (var3 > 2097152) {
                throw new RuntimeException("Chunk Packet trying to allocate too much memory on read.");
            } else {
                this.buffer = new byte[var3];
                var0.readBytes(this.buffer);
                this.blockEntitiesData = var0.readList(ClientboundLevelChunkPacketData.BlockEntityInfo::new);
            }
        }
    }

    public void write(FriendlyByteBuf var0) {
        var0.writeNbt(this.heightmaps);
        var0.writeVarInt(this.buffer.length);
        var0.writeBytes(this.buffer);
        var0.writeCollection(this.blockEntitiesData, (var0x, var1) -> {
            var1.write(var0x);
        });
    }

    private static int calculateChunkSize(LevelChunk var0) {
        int var1 = 0;
        LevelChunkSection[] var2 = var0.getSections();
        int var3 = var2.length;

        for(int var4 = 0; var4 < var3; ++var4) {
            LevelChunkSection var5 = var2[var4];
            var1 += var5.getSerializedSize();
        }

        return var1;
    }

    private ByteBuf getWriteBuffer() {
        ByteBuf var0 = Unpooled.wrappedBuffer(this.buffer);
        var0.writerIndex(0);
        return var0;
    }

    public static void extractChunkData(FriendlyByteBuf var0, LevelChunk var1) {
        LevelChunkSection[] var2 = var1.getSections();
        int var3 = var2.length;

        for(int var4 = 0; var4 < var3; ++var4) {
            LevelChunkSection var5 = var2[var4];
            var5.write(var0);
        }

    }

    public Consumer<ClientboundLevelChunkPacketData.BlockEntityTagOutput> getBlockEntitiesTagsConsumer(int var0, int var1) {
        return (var2) -> {
            this.getBlockEntitiesTags(var2, var0, var1);
        };
    }

    private void getBlockEntitiesTags(ClientboundLevelChunkPacketData.BlockEntityTagOutput var0, int var1, int var2) {
        int var3 = 16 * var1;
        int var4 = 16 * var2;
        MutableBlockPos var5 = new MutableBlockPos();
        Iterator var7 = this.blockEntitiesData.iterator();

        while(var7.hasNext()) {
            ClientboundLevelChunkPacketData.BlockEntityInfo var7 = (ClientboundLevelChunkPacketData.BlockEntityInfo)var7.next();
            int var8 = var3 + SectionPos.sectionRelative(var7.packedXZ >> 4);
            int var9 = var4 + SectionPos.sectionRelative(var7.packedXZ);
            var5.set(var8, var7.y, var9);
            var0.accept(var5, var7.type, var7.tag);
        }

    }

    public FriendlyByteBuf getReadBuffer() {
        return new FriendlyByteBuf(Unpooled.wrappedBuffer(this.buffer));
    }

    public CompoundTag getHeightmaps() {
        return this.heightmaps;
    }

    private static class BlockEntityInfo {
        final int packedXZ;
        final int y;
        final BlockEntityType<?> type;
        @Nullable
        final CompoundTag tag;

        private BlockEntityInfo(int var0, int var1, BlockEntityType<?> var2, @Nullable CompoundTag var3) {
            this.packedXZ = var0;
            this.y = var1;
            this.type = var2;
            this.tag = var3;
        }

        private BlockEntityInfo(FriendlyByteBuf var0) {
            this.packedXZ = var0.readByte();
            this.y = var0.readShort();
            int var1 = var0.readVarInt();
            this.type = (BlockEntityType)Registry.BLOCK_ENTITY_TYPE.byId(var1);
            this.tag = var0.readNbt();
        }

        void write(FriendlyByteBuf var0) {
            var0.writeByte(this.packedXZ);
            var0.writeShort(this.y);
            var0.writeVarInt(Registry.BLOCK_ENTITY_TYPE.getId(this.type));
            var0.writeNbt(this.tag);
        }

        static ClientboundLevelChunkPacketData.BlockEntityInfo create(BlockEntity var0) {
            CompoundTag var1 = var0.getUpdateTag();
            BlockPos var2 = var0.getBlockPos();
            int var3 = SectionPos.sectionRelative(var2.getX()) << 4 | SectionPos.sectionRelative(var2.getZ());
            return new ClientboundLevelChunkPacketData.BlockEntityInfo(var3, var2.getY(), var0.getType(), var1.isEmpty() ? null : var1);
        }
    }

    @FunctionalInterface
    public interface BlockEntityTagOutput {
        void accept(BlockPos var1, BlockEntityType<?> var2, @Nullable CompoundTag var3);
    }
}