/*
Crafter®.Graphics
Copyright (C) 2026 Catcrafts®
catcrafts.net

This library is free software; you can redistribute it and/or
modify it under the terms of the GNU Lesser General Public
License version 3.0 as published by the Free Software Foundation;

This library 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
Lesser General Public License for more details.

You should have received a copy of the GNU Lesser General Public
License along with this library; if not, write to the Free Software
Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301 USA
*/

module;
#include "vulkan/vulkan.h"

export module Crafter.Graphics:Decompress;
import :Device;
import std;

export namespace Crafter::Decompress {
    // GDeflate 1.0 requires each VkDecompressMemoryRegionEXT to cover one
    // 64 KiB tile. CompressedBlob stores one stream per RegionMeta — this
    // helper walks the stream's tile-stream header and emits per-tile
    // regions ready for vkCmdDecompressMemoryEXT.
    //
    // streamBytes is the raw stream bytes (CompressedBlob::bytes.subspan(
    // r.srcOffset, r.compressedSize)). srcBase is its GPU device address.
    // dstBase is the GPU device address of the first decompressed byte.
    inline void ExpandStreamToTileRegions(
        std::span<const std::byte> streamBytes,
        VkDeviceAddress srcBase,
        VkDeviceAddress dstBase,
        std::vector<VkDecompressMemoryRegionEXT>& out)
    {
        if (streamBytes.empty()) return;

        constexpr std::size_t kTileSize = 64 * 1024;
        constexpr std::size_t kHeaderSize = 8;

        // TileStream wire layout (see vendored TileStream.h):
        //   u8 id; u8 magic; u16 numTiles;
        //   { u2 tileSizeIdx; u18 lastTileSize; u12 reserved; } packed into u32
        const std::uint8_t* p = reinterpret_cast<const std::uint8_t*>(streamBytes.data());
        std::uint8_t id = p[0];
        std::uint8_t magic = p[1];
        if (id != (magic ^ 0xff)) {
            throw std::runtime_error("GDeflate tile-stream: bad id/magic");
        }
        std::uint16_t numTiles =
              static_cast<std::uint16_t>(p[2])
            | static_cast<std::uint16_t>(p[3]) << 8;
        std::uint32_t packed =
              static_cast<std::uint32_t>(p[4])
            | static_cast<std::uint32_t>(p[5]) << 8
            | static_cast<std::uint32_t>(p[6]) << 16
            | static_cast<std::uint32_t>(p[7]) << 24;
        std::uint32_t lastTileSize = (packed >> 2) & 0x3FFFFu;

        const std::uint32_t* tileOffsets = reinterpret_cast<const std::uint32_t*>(
            streamBytes.data() + kHeaderSize);
        const std::size_t dataStart = kHeaderSize + std::size_t(numTiles) * sizeof(std::uint32_t);

        out.reserve(out.size() + numTiles);
        for (std::uint32_t k = 0; k < numTiles; ++k) {
            std::size_t tileOff = (k > 0) ? tileOffsets[k] : 0;
            std::size_t compressedSize = (k + 1 < numTiles)
                ? std::size_t(tileOffsets[k + 1]) - tileOff
                : std::size_t(tileOffsets[0]);
            std::size_t decompressedSize = (k + 1 < numTiles)
                ? kTileSize
                : (lastTileSize != 0 ? std::size_t(lastTileSize) : kTileSize);
            out.push_back(VkDecompressMemoryRegionEXT {
                .srcAddress      = srcBase + dataStart + tileOff,
                .dstAddress      = dstBase + std::size_t(k) * kTileSize,
                .compressedSize  = compressedSize,
                .decompressedSize = decompressedSize,
            });
        }
    }

    // Records vkCmdDecompressMemoryEXT into `cmd` for the supplied regions
    // (caller pre-resolves srcAddress/dstAddress) and follows it with a
    // synchronization2 memory barrier so the consumer at (dstStage, dstAccess)
    // observes the decompressed bytes.
    //
    // Method is fixed at GDeflate 1.0 — Crafter.Asset's compressed file format
    // emits GDeflate streams. Caller must guarantee Device::memoryDecompressionSupported.
    inline void DecompressOnGPU(
        VkCommandBuffer cmd,
        std::span<const VkDecompressMemoryRegionEXT> regions,
        VkPipelineStageFlags2 dstStage,
        VkAccessFlags2 dstAccess)
    {
        if (regions.empty()) return;

        VkDecompressMemoryInfoEXT info {
            .sType = VK_STRUCTURE_TYPE_DECOMPRESS_MEMORY_INFO_EXT,
            .pNext = nullptr,
            .decompressionMethod = VK_MEMORY_DECOMPRESSION_METHOD_GDEFLATE_1_0_BIT_EXT,
            .regionCount = static_cast<std::uint32_t>(regions.size()),
            .pRegions = regions.data(),
        };
        Device::vkCmdDecompressMemoryEXT(cmd, &info);

        VkMemoryBarrier2 barrier {
            .sType = VK_STRUCTURE_TYPE_MEMORY_BARRIER_2,
            .pNext = nullptr,
            .srcStageMask  = VK_PIPELINE_STAGE_2_MEMORY_DECOMPRESSION_BIT_EXT,
            .srcAccessMask = VK_ACCESS_2_MEMORY_DECOMPRESSION_WRITE_BIT_EXT,
            .dstStageMask  = dstStage,
            .dstAccessMask = dstAccess,
        };
        VkDependencyInfo dep {
            .sType = VK_STRUCTURE_TYPE_DEPENDENCY_INFO,
            .pNext = nullptr,
            .dependencyFlags = 0,
            .memoryBarrierCount = 1,
            .pMemoryBarriers = &barrier,
            .bufferMemoryBarrierCount = 0,
            .pBufferMemoryBarriers = nullptr,
            .imageMemoryBarrierCount = 0,
            .pImageMemoryBarriers = nullptr,
        };
        vkCmdPipelineBarrier2(cmd, &dep);
    }
}