/*
    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.techhider;

import com.comphenix.tinyprotocol.Reflection;
import de.steamwar.core.Core;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;
import net.minecraft.core.IRegistry;
import net.minecraft.network.protocol.game.ClientboundLevelChunkPacketData;
import net.minecraft.network.protocol.game.ClientboundLevelChunkWithLightPacket;
import net.minecraft.resources.MinecraftKey;
import net.minecraft.util.SimpleBitStorage;
import net.minecraft.world.level.block.entity.TileEntityTypes;
import org.bukkit.entity.Player;

import java.util.Collections;
import java.util.List;
import java.util.Set;
import java.util.function.BiFunction;
import java.util.function.UnaryOperator;
import java.util.stream.Collectors;

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

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

    private static final Reflection.FieldAccessor<Integer> chunkXField = Reflection.getField(ClientboundLevelChunkWithLightPacket.class, int.class, 0);
    private static final Reflection.FieldAccessor<Integer> chunkZField = 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);

    @Override
    public BiFunction<Player, Object, Object> chunkHiderGenerator(TechHider.LocationEvaluator locationEvaluator, int obfuscationTarget, Set<Integer> obfuscate, Set<String> hiddenBlockEntities) {
        return (p, packet) -> {
            int chunkX = chunkXField.get(packet);
            int chunkZ = chunkZField.get(packet);
            if (locationEvaluator.skipChunk(p, chunkX, chunkZ))
                return packet;

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

            tileEntities.set(dataWrapper, ((List<?>)tileEntities.get(dataWrapper)).stream().filter(te -> tileEntityVisible(hiddenBlockEntities, te)).collect(Collectors.toList()));

            ByteBuf in = Unpooled.wrappedBuffer(dataField.get(dataWrapper));
            ByteBuf out = Unpooled.buffer(in.readableBytes() + 64);
            int xOffset = 16*chunkX;
            int zOffset = 16*chunkZ;
            for(int yOffset = p.getWorld().getMinHeight(); yOffset < p.getWorld().getMaxHeight(); yOffset += 16) {
                int finalYOffset = yOffset;
                int chunkY = yOffset/16;
                dataHider((x, y, z) -> locationEvaluator.check(p, x+xOffset, y+finalYOffset, z+zOffset), obfuscationTarget, obfuscate, in, out, locationEvaluator.skipChunkSection(p, chunkX, chunkY, chunkZ), locationEvaluator.blockPrecise(p, chunkX, chunkY, chunkZ));
            }
            out.writeBytes(in); // MC appends a 0 byte at the end if there is a full chunk, idk why

            byte[] data = new byte[out.readableBytes()];
            out.readBytes(data);
            dataField.set(dataWrapper, data);

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

    public static final Class<?> tileEntity = Reflection.getClass("net.minecraft.network.protocol.game.ClientboundLevelChunkPacketData$a");
    protected static final Reflection.FieldAccessor<TileEntityTypes> entityType = Reflection.getField(tileEntity, TileEntityTypes.class, 0);
    private static final IRegistry<?> tileEntityTypes = Reflection.getField(Core.getVersion() > 18 ? Reflection.getClass("net.minecraft.core.registries.BuiltInRegistries") : IRegistry.class, IRegistry.class, 0, TileEntityTypes.class).get(null);
    private static final Reflection.MethodInvoker getKey = Reflection.getTypedMethod(IRegistry.class, null, MinecraftKey.class, Object.class);
    private static final Reflection.MethodInvoker getName = Reflection.getTypedMethod(MinecraftKey.class, null, String.class);
    protected boolean tileEntityVisible(Set<String> hiddenBlockEntities, Object tile) {
        return !hiddenBlockEntities.contains((String) getName.invoke(getKey.invoke(tileEntityTypes, entityType.get(tile))));
    }

    private void dataHider(PosEvaluator locationEvaluator, int obfuscationTarget, Set<Integer> obfuscate, ByteBuf in, ByteBuf out, boolean skip, boolean blockPrecise) {
        out.writeShort(in.readShort()); // Block count

        containerWalker(in, out, 15, obfuscationTarget, skip ? Collections.emptySet() : obfuscate, (paletteTarget, dataArrayLength, bitsPerBlock) -> {
            Set<Integer> palettedObfuscate;
            if(bitsPerBlock < 15) {
                palettedObfuscate = Collections.emptySet();
            } else {
                paletteTarget = obfuscationTarget;
                palettedObfuscate = obfuscate;
            }

            if(skip || dataArrayLength == 0 || (!blockPrecise && bitsPerBlock < 15)) {
                out.writeBytes(in, dataArrayLength*8);
                return;
            }

            long[] array = new long[dataArrayLength];
            for(int i = 0; i < dataArrayLength; i++)
                array[i] = in.readLong();
            SimpleBitStorage values = new SimpleBitStorage(bitsPerBlock, 4096, array);

            for (int y = 0; y < 16; y++) {
                for (int z = 0; z < 16; z++) {
                    for (int x = 0; x < 16; x++) {
                        int pos = (((y * 16) + z) * 16) + x;

                        switch (locationEvaluator.test(x, y, z)) {
                            case SKIP:
                                break;
                            case CHECK:
                                if(!palettedObfuscate.contains(values.a(pos)))
                                    break;
                            default:
                                values.b(pos, paletteTarget);
                        }
                    }
                }
            }

            for (long l : values.a())
                out.writeLong(l);
        });
        containerWalker(in, out, 6, 0, Collections.emptySet(), (paletteTarget, dataArrayLength, bitsPerBlock) -> out.writeBytes(in, dataArrayLength * 8)); // Biomes
    }

    private void containerWalker(ByteBuf in, ByteBuf out, int globalPalette, int obfuscationTarget, Set<Integer> obfuscate, TriConsumer<Integer, Integer, Byte> dataArray) {
        byte bitsPerBlock = in.readByte();
        out.writeByte(bitsPerBlock);

        //blockId -> obfuscate.contains(blockId) ? obfuscationTarget : blockId
        if(bitsPerBlock == 0) {
            int value = ProtocolUtils.readVarInt(in);
            ProtocolUtils.writeVarInt(out, obfuscate.contains(value) ? obfuscationTarget : value);
        }else if(bitsPerBlock < globalPalette) {
            obfuscationTarget = ChunkHider.processPalette(obfuscationTarget, obfuscate, in, out);
        }

        int dataArrayLength = ProtocolUtils.readVarInt(in);
        ProtocolUtils.writeVarInt(out, dataArrayLength);
        dataArray.accept(obfuscationTarget, dataArrayLength, bitsPerBlock);
    }

    private interface TriConsumer<X, Y, Z> {
        void accept(X x, Y y, Z z);
    }
}
