/*
 *   This file is a part of the SteamWar software.
 *
 *   Copyright (C) 2025 SteamWar.de-Serverteam
 *
 *   This program is free software: you can redistribute it and/or modify
 *   it under the terms of the GNU Affero General Public License as published by
 *   the Free Software Foundation, either version 3 of the License, or
 *   (at your option) any later version.
 *
 *   This program is distributed in the hope that it will be useful,
 *   but WITHOUT ANY WARRANTY; without even the implied warranty of
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 *   GNU Affero General Public License for more details.
 *
 *   You should have received a copy of the GNU Affero General Public License
 *   along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

package de.steamwar.techhider;

import de.steamwar.Reflection;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;
import it.unimi.dsi.fastutil.ints.Int2IntMap;
import it.unimi.dsi.fastutil.ints.Int2IntOpenHashMap;
import it.unimi.dsi.fastutil.ints.IntArrayList;
import net.minecraft.core.SectionPos;
import net.minecraft.network.protocol.game.ClientboundLevelChunkPacketData;
import net.minecraft.network.protocol.game.ClientboundLevelChunkWithLightPacket;
import net.minecraft.util.SimpleBitStorage;
import net.minecraft.world.level.block.Block;
import net.minecraft.world.level.block.entity.BlockEntityType;
import net.minecraft.world.level.block.state.BlockState;
import org.bukkit.entity.Player;

import java.util.*;
import java.util.function.UnaryOperator;


public abstract class ChunkHider {
    private static final UnaryOperator<Object> chunkPacketShallowCloner = ProtocolUtils.shallowCloneGenerator(ClientboundLevelChunkWithLightPacket.class);
    private static final UnaryOperator<Object> chunkDataShallowCloner = ProtocolUtils.shallowCloneGenerator(ClientboundLevelChunkPacketData.class);

    private static final Reflection.Field<ClientboundLevelChunkPacketData> levelChunkPacketDataField = Reflection.getField(ClientboundLevelChunkWithLightPacket.class, ClientboundLevelChunkPacketData.class, 0);

    private static final Reflection.Field<byte[]> chunkBlockDataField = Reflection.getField(ClientboundLevelChunkPacketData.class, byte[].class, 0);
    private static final Reflection.Field<List> chunkBlockEntitiesDataField = Reflection.getField(ClientboundLevelChunkPacketData.class, List.class, 0);

    private final int SECTION_SPAN_SIZE = 16;
    private final byte BIT_PER_BLOCK_INDIRECTION_LIMIT = 8;
    private final byte BIT_PER_BIOME_INDIRECTION_LIMIT = 3;

    private final int BLOCKS_PER_SECTION = 4096;
    private final int BIOMES_PER_SECTION = 64;

    private final byte BITS_PER_LONG = 64;

    private final int blockIdUsedForHiding = 315;

    public ChunkHider(Block blockUsedForObfuscation) {
        // blockIdUsedForHiding =  Block.BLOCK_STATE_REGISTRY.getId(blockUsedForObfuscation.defaultBlockState());
    }

    private int getLongsRequiredToEncodeEntries(int bitsPerEntry, int entryCount) {
        int entriesPerLong =  BITS_PER_LONG / bitsPerEntry;
        int dataLengthAsLongCount = (entryCount + entriesPerLong - 1) / entriesPerLong;

        return dataLengthAsLongCount;
    }

    private long[] readSectionDataFromBuffer(ByteBuf dataSource, short bitsPerEntry, int entryCount) {
        int dataLengthAsLongCount = getLongsRequiredToEncodeEntries(bitsPerEntry, entryCount);
        long[] dataArray = new long[dataLengthAsLongCount];

        for(int i = 0; i < dataLengthAsLongCount; i++){
            dataArray[i] = dataSource.readLong();
        }

        return dataArray;
    }

    public ClientboundLevelChunkWithLightPacket processLevelChunkWithLightPacket(Player player, ClientboundLevelChunkWithLightPacket packet) {
        int chunkX = packet.getX();
        int chunkZ = packet.getZ();
        ClientboundLevelChunkPacketData chunkData = packet.getChunkData();

        ByteBuf in = Unpooled.wrappedBuffer(chunkData.getReadBuffer());
        ByteBuf out = Unpooled.buffer(in.readableBytes() + 64);

        int worldMinHeight = player.getWorld().getMinHeight();
        int worldMaxHeight = player.getWorld().getMaxHeight();

        for (int yOffset = worldMinHeight; yOffset < worldMaxHeight; yOffset += SECTION_SPAN_SIZE) {
            short blockCount = in.readShort();

            byte bitsPerBlock = in.readByte();

            if(bitsPerBlock == 0) {
                int sectionBlockId = ProtocolUtils.readVarInt(in);

                out.writeShort(blockCount);
                out.writeByte(bitsPerBlock);
                ProtocolUtils.writeVarInt(out, sectionBlockId);
            }
            else if (bitsPerBlock <= BIT_PER_BLOCK_INDIRECTION_LIMIT) {
                int palletLength = ProtocolUtils.readVarInt(in);
                int[] pallet = ProtocolUtils.readVarIntArray(in, palletLength);

                long[] rawData = readSectionDataFromBuffer(in, bitsPerBlock,  BLOCKS_PER_SECTION);
                SimpleBitStorage data = new SimpleBitStorage(bitsPerBlock, BLOCKS_PER_SECTION, rawData);

                int[] resolvedData = new int[BLOCKS_PER_SECTION];
                for(int i = 0; i < BLOCKS_PER_SECTION; i++){
                    int palletReference = data.get(i);
                    resolvedData[i] = pallet[palletReference];
                }

                int[] obfuscatedBlockData = obfuscateBlockDataArray(player, chunkX, chunkZ, yOffset, resolvedData);

                long[] reEncodedData = encodeDirectBlockDataArray(obfuscatedBlockData);


                out.writeShort(blockCount);
                out.writeByte(15);
                for(long rawDataSegment : reEncodedData) {
                    out.writeLong(rawDataSegment);
                }
                /*
                Int2IntMap blockIdToPalletIndex = new Int2IntOpenHashMap();
                IntArrayList newPallet = new IntArrayList();

                for(int blockId : obfuscatedData) {
                    if(!blockIdToPalletIndex.containsKey(blockId)) {
                        newPallet.add(blockId);
                        blockIdToPalletIndex.put(blockId, newPallet.size());
                    }
                }

                byte newBitsPerBlock = (byte) getUnsignedBitLength(newPallet.size());
                int[] newPalletRaw = newPallet.toArray(new int[newPallet.size()]);

                SimpleBitStorage reEncodedData = new SimpleBitStorage(bitsPerBlock, BLOCKS_PER_SECTION, new long[rawData.length]);
                for(int i = 0; i < obfuscatedData.length; i++) {
                    int blockId = obfuscatedData[i];
                    int palletReference = blockIdToPalletIndex.get(blockId);

                    reEncodedData.set(i, palletReference);
                }

                out.writeShort(blockCount);
                out.writeByte(newBitsPerBlock);
                ProtocolUtils.writeVarInt(out, palletLength);
                ProtocolUtils.writeVarIntArray(out, newPalletRaw);
                for(long rawDataSegment : reEncodedData.getRaw()) {
                    out.writeLong(rawDataSegment);
                }*/

            }
            else {
                long[] rawData = readSectionDataFromBuffer(in, bitsPerBlock,  BLOCKS_PER_SECTION);

                int[] blockData = resolveDirectBlockDataArray(rawData, bitsPerBlock);

                int[] obfuscatedBlockData = obfuscateBlockDataArray(player, chunkX, chunkZ, yOffset, blockData);

                long[] reEncodedData = encodeDirectBlockDataArray(obfuscatedBlockData);


                out.writeShort(blockCount);
                out.writeByte(15);
                for(long rawDataSegment : reEncodedData) {
                    out.writeLong(rawDataSegment);
                }
            }

            copyOverSectionBiomeData(in, out);
        }

        if (in.readableBytes() != 0) {
            return null;
        }

        byte[] data = new byte[out.readableBytes()];
        out.readBytes(data);

        List<Object> blockEntities = chunkBlockEntitiesDataField.get(chunkData);
        List<Object> filteredBlockEntities = filterBlockEntities(player, blockEntities);

        return buildNewChunkPacket(packet, data, filteredBlockEntities);

    }

    private int[] obfuscateBlockDataArray(Player player, int chunkX, int chunkZ, int yOffset, int[] blockDataArray) {
        int[] obfuscatedData = new int[BLOCKS_PER_SECTION];
        for (int sectionY = 0; sectionY < SECTION_SPAN_SIZE; sectionY++) {
            for (int sectionZ = 0; sectionZ < SECTION_SPAN_SIZE; sectionZ++) {
                for (int sectionX = 0; sectionX < SECTION_SPAN_SIZE; sectionX++) {
                    int blockDataIndex = (((sectionY * SECTION_SPAN_SIZE) + sectionZ) * SECTION_SPAN_SIZE) + sectionX;

                    int worldX = sectionX + (SECTION_SPAN_SIZE * chunkX);
                    int worldY = sectionY + yOffset;
                    int worldZ = sectionZ + (SECTION_SPAN_SIZE * chunkZ);

                    int blockId = blockDataArray[blockDataIndex];
                    BlockState blockState =  Block.BLOCK_STATE_REGISTRY.byId(blockId);
                    Block block = blockState.getBlock();



                    if(isPlayerPrivilegedToAccessPosition(player, worldX, worldY, worldZ) && isPlayerPrivilegedToAccessBlock(player, worldX, worldY, worldZ, block)) {
                        obfuscatedData[blockDataIndex] = blockId;
                    }
                    else {
                        obfuscatedData[blockDataIndex] = blockIdUsedForHiding;
                    }
                }
            }
        }

        return obfuscatedData;
    }

    private int[] resolveDirectBlockDataArray(long[] rawDataArray, byte bitsPerBlock) {
        SimpleBitStorage data = new SimpleBitStorage(bitsPerBlock, BLOCKS_PER_SECTION, rawDataArray);

        int[] resolvedData = new int[BLOCKS_PER_SECTION];
        for(int i = 0; i < BLOCKS_PER_SECTION; i++){
            resolvedData[i] = data.get(i);
        }

        return resolvedData;
    }

    private long[] encodeDirectBlockDataArray(int[] blockDataArray) {
        int longsRequiredToEncodeData = getLongsRequiredToEncodeEntries(15, blockDataArray.length);

        SimpleBitStorage reEncodedData = new SimpleBitStorage(15, BLOCKS_PER_SECTION, new long[longsRequiredToEncodeData]);
        for(int i = 0; i < blockDataArray.length; i++) {
            int blockId = blockDataArray[i];

            reEncodedData.set(i, blockId);
        }

        return reEncodedData.getRaw();
    }

    public abstract boolean isPlayerPrivilegedToAccessPosition(Player p, int blockX, int blockY, int blockZ);
    public abstract boolean isPlayerPrivilegedToAccessBlock(Player p, int blockX, int blockY, int blockZ, Block block);
    public abstract boolean isPlayerPrivilegedToAccessBlockEntity(Player p, int blockX, int blockY, int blockZ, BlockEntityType<?> type);

    private ClientboundLevelChunkWithLightPacket buildNewChunkPacket(ClientboundLevelChunkWithLightPacket originalPacket, byte[] newBlockDataBuffer, List<Object> newBlockEntities) {
        ClientboundLevelChunkWithLightPacket clonedPacket = (ClientboundLevelChunkWithLightPacket) chunkPacketShallowCloner.apply(originalPacket);
        ClientboundLevelChunkPacketData clonedPacketChunkData = (ClientboundLevelChunkPacketData) chunkDataShallowCloner.apply(originalPacket.getChunkData());

        chunkBlockDataField.set(clonedPacketChunkData, newBlockDataBuffer);
        chunkBlockEntitiesDataField.set(clonedPacketChunkData, newBlockEntities);

        levelChunkPacketDataField.set(clonedPacket, clonedPacketChunkData);

        return  clonedPacket;
    }


    private static final Class<?> blockEntitiyInfoClass = Reflection.getClass("net.minecraft.network.protocol.game.ClientboundLevelChunkPacketData$BlockEntityInfo");

    private static final Reflection.Field<BlockEntityType> blockEntityInfoTypeField = Reflection.getField(blockEntitiyInfoClass, BlockEntityType.class, 0);
    private static final Reflection.Field<Integer> packedXZField = Reflection.getField(blockEntitiyInfoClass, int.class, 0);
    private static final Reflection.Field<Integer> yField = Reflection.getField(blockEntitiyInfoClass, int.class, 1);
    private List<Object> filterBlockEntities(Player player, List<Object> blockEntities) {
        return blockEntities.stream()
                .filter((blockEntityInfo) -> {
                    BlockEntityType<?> type = blockEntityInfoTypeField.get(blockEntityInfo);

                    int packedXZ = packedXZField.get(blockEntityInfo);

                    int y = yField.get(blockEntityInfo);
                    int x = SectionPos.sectionRelativeX((short) packedXZ);
                    int z = SectionPos.sectionRelativeZ((short) packedXZ);

                    return isPlayerPrivilegedToAccessPosition(player, x, y, z) && isPlayerPrivilegedToAccessBlockEntity(player, x, y, z, type);
                }).toList();
    }

    private void copyOverSectionBiomeData(ByteBuf oldData, ByteBuf newData){
        short bitsPerBiome = oldData.readByte();
        newData.writeByte(bitsPerBiome);
        if(bitsPerBiome == 0) {
            int sectionBiomeId = ProtocolUtils.readVarInt(oldData);
            ProtocolUtils.writeVarInt(newData, sectionBiomeId);
        }
        else if(bitsPerBiome <= BIT_PER_BIOME_INDIRECTION_LIMIT) {
            int palletLength = ProtocolUtils.readVarInt(oldData);
            ProtocolUtils.writeVarInt(newData, palletLength);

            for(int i = 0; i < palletLength; i++) {
                int palletEntry = ProtocolUtils.readVarInt(oldData);
                ProtocolUtils.writeVarInt(newData, palletEntry);
            }

            long[] rawData = readSectionDataFromBuffer(oldData, bitsPerBiome,  BIOMES_PER_SECTION);
            for(long rawDataSegment : rawData ) {
                newData.writeLong(rawDataSegment);
            }
        }
        else {
            long[] rawData = readSectionDataFromBuffer(oldData, bitsPerBiome,  BIOMES_PER_SECTION);
            for(long rawDataSegment : rawData ) {
                newData.writeLong(rawDataSegment);
            }
        }

    }
}
