package org.allaymc.server.network.processor.common;

import io.netty.buffer.ByteBuf;
import io.netty.buffer.ByteBufAllocator;
import io.netty.buffer.ByteBufOutputStream;
import io.netty.buffer.Unpooled;
import lombok.extern.slf4j.Slf4j;
import org.allaymc.api.player.Player;
import org.allaymc.api.registry.Registries;
import org.allaymc.api.world.chunk.OperationType;
import org.allaymc.api.world.dimension.DimensionType;
import org.allaymc.server.network.processor.PacketProcessor;
import org.allaymc.server.player.AllayPlayer;
import org.allaymc.server.player.ChunkCache;
import org.allaymc.server.world.chunk.AllayChunkSection;
import org.allaymc.server.world.chunk.ChunkEncoder;
import org.cloudburstmc.math.vector.Vector3i;
import org.cloudburstmc.nbt.NbtUtils;
import org.cloudburstmc.protocol.bedrock.data.HeightMapDataType;
import org.cloudburstmc.protocol.bedrock.data.SubChunkData;
import org.cloudburstmc.protocol.bedrock.data.SubChunkRequestResult;
import org.cloudburstmc.protocol.bedrock.packet.BedrockPacketType;
import org.cloudburstmc.protocol.bedrock.packet.PacketSignal;
import org.cloudburstmc.protocol.bedrock.packet.SubChunkRequestPacket;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.UUID;

/**
 * @author Cool_Loong | daoge_cmd
 */
@Slf4j
public class SubChunkRequestPacketProcessor extends PacketProcessor<SubChunkRequestPacket> {

    private record SubChunkInfo(
        Vector3i offset,
        SubChunkRequestResult result,
        HeightMapDataType hMapType,
        byte[] hMapData,
        byte[] sectionBlob,
        byte[] blockEntityBytes
    ) {}

    @Override
    public PacketSignal handleAsync(Player player, SubChunkRequestPacket packet, long receiveTime) {
        var allayPlayer = (AllayPlayer) player;
        var encoder = allayPlayer.getProtocol().getEncoder();
        var entity = player.getControlledEntity();
        var dimensionType = Registries.DIMENSIONS.getByK1(packet.getDimension());
        if (dimensionType == null || dimensionType != entity.getDimension().getDimensionType()) {
            // Outdated sub chunk request from a previous dimension
            allayPlayer.sendPacket(encoder.encodeSubChunk(
                    packet.getDimension(),
                    packet.getSubChunkPosition(),
                    false,
                    List.of()
            ));
            return PacketSignal.HANDLED;
        }

        ChunkCache cache = null;
        UUID playerId = null;
        int cacheGen = 0;

        if (allayPlayer.isCacheEffectivelyEnabled()) {
            cache = ChunkCache.getInstance();
            playerId = allayPlayer.getLoginData().getUuid();
            cacheGen = cache.getGeneration(playerId);
        }

        var centerPosition = packet.getSubChunkPosition();

        // Phase 1: Collect raw data
        var infos = new ArrayList<SubChunkInfo>(packet.getPositionOffsets().size());
        for (var offset : packet.getPositionOffsets()) {
            infos.add(collectSubChunkInfo(player, dimensionType, centerPosition, offset));
        }

        // Collect non-null section blobs for caching
        List<byte[]> nonAirBlobs = new ArrayList<>();
        for (var info : infos) {
            if (info.sectionBlob != null) {
                nonAirBlobs.add(info.sectionBlob);
            }
        }

        long[] hashes = null;
        if (cache != null && !nonAirBlobs.isEmpty()) {
            hashes = cache.tryStoreBlobsAndOpenTransaction(playerId, cacheGen, nonAirBlobs.toArray(new byte[0][]));
        }

        boolean cacheEnabled = hashes != null;
        var responseData = new ArrayList<SubChunkData>(infos.size());
        if (cacheEnabled) {
            // Cached path
            int hashIndex = 0;

            for (var info : infos) {
                var subChunkData = new SubChunkData();
                subChunkData.setPosition(info.offset);
                subChunkData.setResult(info.result);

                setHeightMaps(subChunkData, info.hMapType, info.hMapData);

                // Data
                if (info.result == SubChunkRequestResult.SUCCESS) {
                    subChunkData.setBlobId(hashes[hashIndex++]);
                    if (info.blockEntityBytes != null) {
                        subChunkData.setData(Unpooled.wrappedBuffer(info.blockEntityBytes));
                    } else {
                        subChunkData.setData(Unpooled.EMPTY_BUFFER);
                    }
                } else {
                    // SUCCESS_ALL_AIR or error results
                    subChunkData.setBlobId(0);
                    subChunkData.setData(Unpooled.EMPTY_BUFFER);
                }

                responseData.add(subChunkData);
            }
        } else {
            // Non-cached path
            for (var info : infos) {
                var subChunkData = new SubChunkData();
                subChunkData.setPosition(info.offset);
                subChunkData.setResult(info.result);

                setHeightMaps(subChunkData, info.hMapType, info.hMapData);

                // Data - combine section blob + block entities
                if (info.result == SubChunkRequestResult.SUCCESS) {
                    ByteBuf buffer = Unpooled.wrappedBuffer(info.sectionBlob);
                    if (info.blockEntityBytes != null) {
                        buffer = Unpooled.wrappedBuffer(info.sectionBlob, info.blockEntityBytes);
                    }
                    subChunkData.setData(buffer);
                } else {
                    // SUCCESS_ALL_AIR or error results
                    subChunkData.setData(Unpooled.EMPTY_BUFFER);
                }

                responseData.add(subChunkData);
            }
        }

        allayPlayer.sendPacket(encoder.encodeSubChunk(
                packet.getDimension(),
                centerPosition,
                cacheEnabled,
                responseData
        ));
        return PacketSignal.HANDLED;
    }

    /**
     * Assigns independently releasable buffer views to the logical and render height maps.
     */
    static void setHeightMaps(SubChunkData subChunkData, HeightMapDataType type, byte[] data) {
        subChunkData.setHeightMapType(type);
        subChunkData.setRenderHeightMapType(type);
        if (data == null) {
            subChunkData.setHeightMapData(Unpooled.EMPTY_BUFFER);
            subChunkData.setRenderHeightMapData(Unpooled.EMPTY_BUFFER);
            return;
        }

        ByteBuf heightMap = Unpooled.wrappedBuffer(data);
        subChunkData.setHeightMapData(heightMap);
        subChunkData.setRenderHeightMapData(heightMap.retainedDuplicate());
    }

    private SubChunkInfo collectSubChunkInfo(Player player, DimensionType dimensionType, Vector3i center, Vector3i offset) {
        var entity = player.getControlledEntity();

        int sectionY = center.getY() + offset.getY();
        if (sectionY < dimensionType.minSectionY() || sectionY > dimensionType.maxSectionY()) {
            log.warn("Player {} requested sub-chunk at y={} which is out of bounds ({}, {})", player.getOriginName(), sectionY, dimensionType.minSectionY(), dimensionType.maxSectionY());
            return new SubChunkInfo(offset, SubChunkRequestResult.INDEX_OUT_OF_BOUNDS, HeightMapDataType.NO_DATA, null, null, null);
        }

        int chunkX = center.getX() + offset.getX();
        int chunkZ = center.getZ() + offset.getZ();
        var chunk = entity.getDimension().getChunkManager().getChunk(chunkX, chunkZ);
        if (chunk == null) {
            return new SubChunkInfo(offset, SubChunkRequestResult.CHUNK_NOT_FOUND, HeightMapDataType.NO_DATA, null, null, null);
        }

        // Compute height map
        byte[][] heightMapHolder = new byte[1][];
        HeightMapDataType[] hMapTypeHolder = new HeightMapDataType[1];

        chunk.applyOperation(unsafeChunk -> {
            var heightMap = new byte[256];
            boolean allHigher = true;
            boolean allLower = true;

            for (int x = 0; x < 16; x++) {
                for (int z = 0; z < 16; z++) {
                    int height = unsafeChunk.getHeight(x, z);
                    int heightSectionY = height >> 4;
                    int index = (z << 4) | x;

                    if (heightSectionY > sectionY) {
                        heightMap[index] = 16;
                        allLower = false;
                    } else if (heightSectionY < sectionY) {
                        heightMap[index] = -1;
                        allHigher = false;
                    } else {
                        heightMap[index] = (byte) (height & 0xf);
                        allHigher = false;
                        allLower = false;
                    }
                }
            }

            if (allHigher) {
                hMapTypeHolder[0] = HeightMapDataType.TOO_HIGH;
                heightMapHolder[0] = null;
            } else if (allLower) {
                hMapTypeHolder[0] = HeightMapDataType.TOO_LOW;
                heightMapHolder[0] = null;
            } else {
                hMapTypeHolder[0] = HeightMapDataType.HAS_DATA;
                heightMapHolder[0] = heightMap;
            }
        }, OperationType.READ, OperationType.NONE);

        // Encode section and block entities
        byte[][] sectionBlobHolder = new byte[1][];
        byte[][] blockEntityBytesHolder = new byte[1][];
        SubChunkRequestResult[] resultHolder = new SubChunkRequestResult[1];

        chunk.applyOperationInSection(sectionY, s -> {
            var subChunk = (AllayChunkSection) s;
            if (subChunk.isAirSection()) {
                resultHolder[0] = SubChunkRequestResult.SUCCESS_ALL_AIR;
                sectionBlobHolder[0] = null;
                blockEntityBytesHolder[0] = null;
            } else {
                resultHolder[0] = SubChunkRequestResult.SUCCESS;
                sectionBlobHolder[0] = ChunkEncoder.encodeSectionBlob(subChunk);

                var blockEntities = chunk.getSectionBlockEntities(sectionY);
                if (!blockEntities.isEmpty()) {
                    var buffer = ByteBufAllocator.DEFAULT.ioBuffer();
                    try (var writer = NbtUtils.createNetworkWriter(new ByteBufOutputStream(buffer))) {
                        for (var blockEntity : blockEntities) {
                            writer.writeTag(blockEntity.saveNBT());
                        }
                        byte[] beBytes = new byte[buffer.readableBytes()];
                        buffer.readBytes(beBytes);
                        blockEntityBytesHolder[0] = beBytes;
                    } catch (IOException e) {
                        log.error("Error while encoding block entity in sub-chunk ({}, {}, {})", chunkX, sectionY, chunkZ, e);
                    } finally {
                        buffer.release();
                    }
                } else {
                    blockEntityBytesHolder[0] = null;
                }
            }
        }, OperationType.READ, OperationType.NONE);

        return new SubChunkInfo(
            offset,
            resultHolder[0],
            hMapTypeHolder[0],
            heightMapHolder[0],
            sectionBlobHolder[0],
            blockEntityBytesHolder[0]
        );
    }

    @Override
    public BedrockPacketType getPacketType() {
        return BedrockPacketType.SUB_CHUNK_REQUEST;
    }
}
