/*
    This file is a part of the SteamWar software.

    Copyright (C) 2021  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.fightsystem.utils;

import com.comphenix.tinyprotocol.Reflection;
import de.steamwar.fightsystem.Config;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.UnpooledByteBufAllocator;
import net.minecraft.core.IRegistry;
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.entity.TileEntityTypes;
import org.bukkit.World;
import org.bukkit.entity.Player;

import java.nio.ByteBuffer;
import java.util.List;
import java.util.Set;
import java.util.function.IntFunction;
import java.util.function.UnaryOperator;
import java.util.stream.Collectors;

public class TechHider18 implements TechHider.ChunkHider {

    @Override
    public Class<?> mapChunkPacket() {
        return ClientboundLevelChunkWithLightPacket.class;
    }

    private static final UnaryOperator<Object> chunkPacketCloner = ProtocolAPI.shallowCloneGenerator(ClientboundLevelChunkWithLightPacket.class);
    private static final UnaryOperator<Object> chunkDataCloner = ProtocolAPI.shallowCloneGenerator(ClientboundLevelChunkPacketData.class);

    private static final Reflection.FieldAccessor<Integer> chunkX = Reflection.getField(ClientboundLevelChunkWithLightPacket.class, int.class, 0);
    private static final Reflection.FieldAccessor<Integer> chunkZ = Reflection.getField(ClientboundLevelChunkWithLightPacket.class, int.class, 1);
    private static final Reflection.FieldAccessor<ClientboundLevelChunkPacketData> chunkData = Reflection.getField(ClientboundLevelChunkWithLightPacket.class, ClientboundLevelChunkPacketData.class, 0);

    private static final Reflection.FieldAccessor<byte[]> dataField = Reflection.getField(ClientboundLevelChunkPacketData.class, byte[].class, 0);
    private static final Reflection.FieldAccessor<List> tileEntities = Reflection.getField(ClientboundLevelChunkPacketData.class, List.class, 0);
    public static final Class<?> tileEntity = Reflection.getClass("net.minecraft.network.protocol.game.ClientboundLevelChunkPacketData$a");
    private static final Reflection.FieldAccessor<TileEntityTypes> entityType = Reflection.getField(tileEntity, TileEntityTypes.class, 0);

    private final Set<Integer> hiddenBlockIds = BlockIdWrapper.impl.getHiddenBlockIds();
    private final int obfuscateWith = BlockIdWrapper.impl.getObfuscateWith();

    @Override
    public Object mapChunkHider(Player p, Object packet) {
        if(TechHider.bypass(p, chunkX.get(packet), chunkZ.get(packet)))
            return packet;

        packet = chunkPacketCloner.apply(packet);
        Object data = chunkDataCloner.apply(chunkData.get(packet));

        tileEntities.set(data, ((List<?>)tileEntities.get(data)).stream().filter(
                tile -> Config.HiddenBlockEntities.contains(IRegistry.aa.b(entityType.get(tile)).a())
        ).collect(Collectors.toList()));

        World world = p.getWorld();
        int sections = (world.getMaxHeight() - world.getMinHeight()) / 16;
        dataField.set(data, dataHider(dataField.get(data), sections));

        chunkData.set(packet, data);
        return packet;
    }

    private byte[] dataHider(byte[] data, int sections) {
        ByteBuf buffer = UnpooledByteBufAllocator.DEFAULT.directBuffer(data.length + 100);
        int i = 0;

        while(sections-- > 0) {
            buffer.writeBytes(data, i, 2); // Block count
            i += 2;

            i = containerWalker(data, buffer, i, 15, blockId -> hiddenBlockIds.contains(blockId) ? obfuscateWith : blockId, (curI, dataArrayLength, bitsPerBlock) -> {
                if(bitsPerBlock < 15) {
                    buffer.writeBytes(data, curI, dataArrayLength * 8);
                } else {
                    ByteBuffer source = ByteBuffer.wrap(data, curI, dataArrayLength * 8);
                    long[] array = new long[dataArrayLength];
                    source.asLongBuffer().get(array);
                    SimpleBitStorage values = new SimpleBitStorage(bitsPerBlock, 4096, array);

                    for (int pos = 0; pos < 4096; pos++) {
                        if (hiddenBlockIds.contains(values.a(pos))) {
                            values.b(pos, obfuscateWith);
                        }
                    }

                    for (long l : values.a())
                        buffer.writeLong(l);
                }
            });
            i = containerWalker(data, buffer, i, 6, value -> value, (curI, dataArrayLength, bitsPerBlock) -> buffer.writeBytes(data, curI, dataArrayLength * 8));
        }
        buffer.writeBytes(data, i, data.length - i); // MC appends a 0 byte at the end if there is a full chunk, idk why

        byte[] outdata = new byte[buffer.readableBytes()];
        buffer.readBytes(outdata);
        return outdata;
    }

    private int containerWalker(byte[] data, ByteBuf buffer, int i, int globalPalette, IntFunction<Integer> palette, Region.TriConsumer<Integer, Integer, Byte> dataArray) {
        byte bitsPerBlock = data[i++];
        buffer.writeByte(bitsPerBlock);

        if(bitsPerBlock == 0) {
            int paletteValue = TechHider.readVarInt(data, i);
            i += TechHider.readVarIntLength(data, i);
            buffer.writeBytes(TechHider.writeVarInt(palette.apply(paletteValue)));
        }else if(bitsPerBlock < globalPalette) {
            int paletteLength = TechHider.readVarInt(data, i);
            int paletteLengthLength = TechHider.readVarIntLength(data, i);
            buffer.writeBytes(data, i, paletteLengthLength);
            i += paletteLengthLength;

            for(int actPaletteId = 0; actPaletteId < paletteLength; actPaletteId++) {
                int paletteValue = TechHider.readVarInt(data, i);
                i += TechHider.readVarIntLength(data, i);
                buffer.writeBytes(TechHider.writeVarInt(palette.apply(paletteValue)));
            }
        }

        int dataArrayLength = TechHider.readVarInt(data, i);
        int dataArrayLengthLength = TechHider.readVarIntLength(data, i);
        buffer.writeBytes(data, i, dataArrayLengthLength);
        i += dataArrayLengthLength;

        dataArray.accept(i, dataArrayLength, bitsPerBlock);
        i += dataArrayLength * 8;

        return i;
    }
}