/*
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  0215-1301 USA
*/

module;

#ifndef CRAFTER_GRAPHICS_WINDOW_DOM
#include "vulkan/vulkan.h"

#endif // !CRAFTER_GRAPHICS_WINDOW_DOM
export module Crafter.Graphics:ImageVulkan;
#ifndef CRAFTER_GRAPHICS_WINDOW_DOM
import std;
import Crafter.Asset;
import :Decompress;
import :VulkanBuffer;

export namespace Crafter {
    template <typename PixelType>
    class ImageVulkan {
    public:
        std::uint16_t width;
        std::uint16_t height;
        std::uint8_t mipLevels;
        VkImage image;
        VkDeviceMemory imageMemory;
        VulkanBuffer<PixelType, true> buffer;
        // Lives until the compressed Update path's cmd buffer completes.
        // Same lifetime contract as Mesh::compressedStaging — caller must
        // not destroy / re-Update before the submit fence is signaled.
        VulkanBuffer<std::byte, true> compressedStaging;
        VkImageView imageView;
        VkDescriptorImageInfo descriptor;

        void Create(std::uint16_t width, std::uint16_t height, std::uint8_t mipLevels, VkCommandBuffer cmd, VkFormat format, VkImageCreateFlags flags, VkImageLayout layout) {
            this->width = width;
            this->height = height;
            this->mipLevels = mipLevels;
            buffer.Create(
                VK_BUFFER_USAGE_TRANSFER_SRC_BIT | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT,
                VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT,
                width * height
            );
            
            VkImageCreateInfo imageInfo = {};
            imageInfo.sType = VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO;
            imageInfo.imageType = VK_IMAGE_TYPE_2D;
            imageInfo.extent.width = width;
            imageInfo.extent.height = height;
            imageInfo.extent.depth = 1;
            imageInfo.mipLevels = mipLevels;
            imageInfo.arrayLayers = 1;
            imageInfo.format = format;
            imageInfo.tiling = VK_IMAGE_TILING_OPTIMAL;
            imageInfo.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;
            imageInfo.usage = flags;
            imageInfo.samples = VK_SAMPLE_COUNT_1_BIT;
            imageInfo.sharingMode = VK_SHARING_MODE_EXCLUSIVE;
            
            Device::CheckVkResult(vkCreateImage(Device::device, &imageInfo, nullptr, &image));

            VkMemoryRequirements memRequirements;
            vkGetImageMemoryRequirements(Device::device, image, &memRequirements);
            
            VkMemoryAllocateInfo allocInfo = {};
            allocInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
            allocInfo.allocationSize = memRequirements.size;
            allocInfo.memoryTypeIndex = Device::GetMemoryType(memRequirements.memoryTypeBits, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);
            
            Device::CheckVkResult(vkAllocateMemory(Device::device, &allocInfo, nullptr, &imageMemory));
            vkBindImageMemory(Device::device, image, imageMemory, 0);

            VkImageViewCreateInfo viewInfo = {};
            viewInfo.sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO;
            viewInfo.image = image;
            viewInfo.viewType = VK_IMAGE_VIEW_TYPE_2D;
            viewInfo.format = format;
            viewInfo.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
            viewInfo.subresourceRange.baseMipLevel = 0;
            viewInfo.subresourceRange.levelCount = mipLevels;
            viewInfo.subresourceRange.baseArrayLayer = 0;
            viewInfo.subresourceRange.layerCount = 1;

            Device::CheckVkResult(vkCreateImageView(Device::device, &viewInfo, nullptr, &imageView));

            // Final transition to shader read-only layout
            TransitionImageLayout(cmd, image, VK_IMAGE_LAYOUT_UNDEFINED, layout, VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT, VK_PIPELINE_STAGE_RAY_TRACING_SHADER_BIT_KHR, 0, VK_ACCESS_SHADER_READ_BIT, 0, mipLevels);

            descriptor = { .imageView = imageView, .imageLayout = layout };
        }

        void Update(VkCommandBuffer cmd, VkImageLayout layout) {
            buffer.FlushDevice(cmd, VK_ACCESS_MEMORY_READ_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT);
            TransitionImageLayout(cmd, image, layout, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, VK_PIPELINE_STAGE_RAY_TRACING_SHADER_BIT_KHR, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_ACCESS_SHADER_READ_BIT, VK_ACCESS_TRANSFER_WRITE_BIT, 0, mipLevels);

            VkBufferImageCopy region{};
            region.bufferOffset = 0;
            region.bufferRowLength = 0;
            region.bufferImageHeight = 0;
            region.imageSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
            region.imageSubresource.mipLevel = 0;
            region.imageSubresource.baseArrayLayer = 0;
            region.imageSubresource.layerCount = 1;
            region.imageOffset = {0, 0, 0};
            region.imageExtent = { width, height, 1};
            vkCmdCopyBufferToImage(
                cmd,
                buffer.buffer,
                image,
                VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
                1,
                &region
            );

            if(mipLevels > 1) {
                TransitionImageLayout(cmd, image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_ACCESS_TRANSFER_WRITE_BIT, VK_ACCESS_TRANSFER_READ_BIT, 0, 1);

                for (std::uint16_t i = 1; i < mipLevels; ++i) {
                    std::uint16_t mipWidth = width >> i;
                    std::uint16_t mipHeight = height >> i;

                    std::uint16_t previousMipWidth = width >> (i - std::uint16_t(1));
                    std::uint16_t previousMipHeight = height >> (i - std::uint16_t(1));

                    VkImageBlit blit = {};
                    blit.srcSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
                    blit.srcSubresource.mipLevel = i - 1;
                    blit.srcSubresource.baseArrayLayer = 0;
                    blit.srcSubresource.layerCount = 1;
                    blit.srcOffsets[0] = { 0, 0, 0 };
                    blit.srcOffsets[1] = { (int32_t)previousMipWidth, (int32_t)previousMipHeight, 1 };

                    blit.dstSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
                    blit.dstSubresource.mipLevel = i;
                    blit.dstSubresource.baseArrayLayer = 0;
                    blit.dstSubresource.layerCount = 1;
                    blit.dstOffsets[0] = { 0, 0, 0 };
                    blit.dstOffsets[1] = { (int32_t)mipWidth, (int32_t)mipHeight, 1 };

                    vkCmdBlitImage(cmd, image, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL, image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 1, &blit, VK_FILTER_LINEAR);
                    TransitionImageLayout(cmd, image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_ACCESS_TRANSFER_WRITE_BIT, VK_ACCESS_TRANSFER_READ_BIT, i, 1);
                }

                TransitionImageLayout(cmd, image, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL, layout, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_RAY_TRACING_SHADER_BIT_KHR, VK_ACCESS_TRANSFER_READ_BIT, VK_ACCESS_SHADER_READ_BIT, 0, mipLevels);
            } else {
                TransitionImageLayout(cmd, image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, layout, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_RAY_TRACING_SHADER_BIT_KHR, VK_ACCESS_TRANSFER_WRITE_BIT, VK_ACCESS_SHADER_READ_BIT, 0, mipLevels);
            }
        }

        // GPU compressed-asset Update: stage compressed bytes, decompress
        // into `buffer` via VK_EXT_memory_decompression, then copy buffer→image
        // and transition to `layout`. Falls back to CPU decode + the existing
        // Update path when Device::memoryDecompressionSupported is false.
        // Caller is responsible for the dimensions matching: asset.sizeX/sizeY
        // must equal this->width/height (set by Create), and asset.pixelStride
        // must equal sizeof(PixelType).
        void Update(const CompressedTextureAsset& asset, VkCommandBuffer cmd, VkImageLayout layout) {
            if (asset.pixelStride != sizeof(PixelType)) {
                throw std::runtime_error("ImageVulkan::Update(compressed): pixel stride mismatch");
            }
            if (!Device::memoryDecompressionSupported) {
                std::span<PixelType> dst{ buffer.value, static_cast<std::size_t>(width) * height };
                std::array<std::span<std::byte>, 1> outputs = {
                    std::as_writable_bytes(dst),
                };
                Compression::DecompressCPU(asset.blob, outputs);
                Update(cmd, layout);
                return;
            }

            // Re-create the staging-into-image buffer with MEMORY_DECOMPRESSION
            // permission so the GPU codec can write into it. Keeps it
            // HOST_VISIBLE (matches the existing path) — on UMA / ReBAR that's
            // a fast path, on older systems the decompress writes traverse
            // PCIe but correctness is unchanged.
            buffer.Resize(
                VK_BUFFER_USAGE_TRANSFER_SRC_BIT
                | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT
                | VK_BUFFER_USAGE_2_MEMORY_DECOMPRESSION_BIT_EXT,
                VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT,
                static_cast<std::uint32_t>(width) * height);

            compressedStaging.Resize(
                VK_BUFFER_USAGE_TRANSFER_SRC_BIT
                | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT
                | VK_BUFFER_USAGE_2_MEMORY_DECOMPRESSION_BIT_EXT,
                VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT,
                static_cast<std::uint32_t>(asset.blob.bytes.size()));
            std::memcpy(compressedStaging.value, asset.blob.bytes.data(), asset.blob.bytes.size());
            compressedStaging.FlushDevice();

            std::vector<VkDecompressMemoryRegionEXT> regions;
            for (const Compression::RegionMeta& r : asset.blob.regions) {
                if (r.decompressedSize == 0) continue;
                std::span<const std::byte> streamBytes(
                    asset.blob.bytes.data() + r.srcOffset,
                    static_cast<std::size_t>(r.compressedSize));
                Decompress::ExpandStreamToTileRegions(
                    streamBytes,
                    compressedStaging.address + r.srcOffset,
                    buffer.address,
                    regions);
            }
            Decompress::DecompressOnGPU(
                cmd,
                regions,
                VK_PIPELINE_STAGE_2_COPY_BIT,
                VK_ACCESS_2_TRANSFER_READ_BIT);

            // Continue with the existing buffer→image upload + layout transitions.
            // We've already inserted the decompress→transfer-read barrier,
            // so we skip the FlushDevice host-write barrier the regular Update
            // would emit (no host write happened).
            TransitionImageLayout(cmd, image, layout, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, VK_PIPELINE_STAGE_RAY_TRACING_SHADER_BIT_KHR, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_ACCESS_SHADER_READ_BIT, VK_ACCESS_TRANSFER_WRITE_BIT, 0, mipLevels);

            VkBufferImageCopy region{};
            region.bufferOffset = 0;
            region.bufferRowLength = 0;
            region.bufferImageHeight = 0;
            region.imageSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
            region.imageSubresource.mipLevel = 0;
            region.imageSubresource.baseArrayLayer = 0;
            region.imageSubresource.layerCount = 1;
            region.imageOffset = {0, 0, 0};
            region.imageExtent = { width, height, 1 };
            vkCmdCopyBufferToImage(cmd, buffer.buffer, image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 1, &region);

            if (mipLevels > 1) {
                TransitionImageLayout(cmd, image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_ACCESS_TRANSFER_WRITE_BIT, VK_ACCESS_TRANSFER_READ_BIT, 0, 1);
                for (std::uint16_t i = 1; i < mipLevels; ++i) {
                    std::uint16_t mipWidth = width >> i;
                    std::uint16_t mipHeight = height >> i;
                    std::uint16_t previousMipWidth = width >> (i - std::uint16_t(1));
                    std::uint16_t previousMipHeight = height >> (i - std::uint16_t(1));

                    VkImageBlit blit = {};
                    blit.srcSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
                    blit.srcSubresource.mipLevel = i - 1;
                    blit.srcSubresource.baseArrayLayer = 0;
                    blit.srcSubresource.layerCount = 1;
                    blit.srcOffsets[0] = { 0, 0, 0 };
                    blit.srcOffsets[1] = { (int32_t)previousMipWidth, (int32_t)previousMipHeight, 1 };
                    blit.dstSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
                    blit.dstSubresource.mipLevel = i;
                    blit.dstSubresource.baseArrayLayer = 0;
                    blit.dstSubresource.layerCount = 1;
                    blit.dstOffsets[0] = { 0, 0, 0 };
                    blit.dstOffsets[1] = { (int32_t)mipWidth, (int32_t)mipHeight, 1 };
                    vkCmdBlitImage(cmd, image, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL, image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 1, &blit, VK_FILTER_LINEAR);
                    TransitionImageLayout(cmd, image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_ACCESS_TRANSFER_WRITE_BIT, VK_ACCESS_TRANSFER_READ_BIT, i, 1);
                }
                TransitionImageLayout(cmd, image, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL, layout, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_RAY_TRACING_SHADER_BIT_KHR, VK_ACCESS_TRANSFER_READ_BIT, VK_ACCESS_SHADER_READ_BIT, 0, mipLevels);
            } else {
                TransitionImageLayout(cmd, image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, layout, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_RAY_TRACING_SHADER_BIT_KHR, VK_ACCESS_TRANSFER_WRITE_BIT, VK_ACCESS_SHADER_READ_BIT, 0, mipLevels);
            }
        }

        void Destroy() {
            vkDestroyImageView(Device::device, imageView, nullptr);
            vkDestroyImage(Device::device, image, nullptr);
            vkFreeMemory(Device::device, imageMemory, nullptr);
        }

        private:
        void TransitionImageLayout(VkCommandBuffer cmd, VkImage image, VkImageLayout oldLayout, VkImageLayout newLayout, VkPipelineStageFlags sourceStage, VkPipelineStageFlags destinationStage, VkAccessFlags srcAccessMask, VkAccessFlags dstAccessMask, std::uint32_t mipLevel, std::uint32_t count) {
            VkImageMemoryBarrier barrier = {};
            barrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
            barrier.oldLayout = oldLayout;
            barrier.newLayout = newLayout;
            barrier.image = image;
            barrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
            barrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
            barrier.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
            barrier.subresourceRange.baseMipLevel = mipLevel;
            barrier.subresourceRange.levelCount = count;
            barrier.subresourceRange.baseArrayLayer = 0;
            barrier.subresourceRange.layerCount = 1;
            barrier.srcAccessMask = srcAccessMask;
            barrier.dstAccessMask = dstAccessMask;

            vkCmdPipelineBarrier(cmd, sourceStage, destinationStage, 0, 0, nullptr, 0, nullptr, 1, &barrier);
        }
    };
}
#endif // !CRAFTER_GRAPHICS_WINDOW_DOM