/*
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"
module Crafter.Graphics:UIRenderer_impl;
import :UIRenderer;
import :Device;
import :Window;
import :DescriptorHeapVulkan;
import :VulkanBuffer;
import :ShaderVulkan;
import :ImageVulkan;
import :UIDrawList;
import :UIAtlas;
import std;

using namespace Crafter;
using namespace Crafter::UI;

namespace {
    // Push-constant block — must match shaders/ui.comp.glsl. The shader's
    // `vec2 surfaceSize` field has 8-byte alignment under std430, so we
    // insert explicit padding after `itemCount` to keep the C++ and GLSL
    // layouts byte-identical (40 bytes total).
    struct PC {
        std::uint32_t itemCount;            //  0
        std::uint32_t _pad0;                //  4
        float         surfaceSize[2];       //  8
        float         scale;                // 16
        std::uint32_t outImageHeapIdx;      // 20
        std::uint32_t itemBufHeapIdx;       // 24
        std::uint32_t atlasTextureHeapIdx;  // 28
        std::uint32_t bindlessBaseHeapIdx;  // 32
        std::uint32_t linearSamplerHeapIdx; // 36
    };
    static_assert(sizeof(PC) == 40, "PC layout must match shader push-constant block");
    static_assert(sizeof(PC) <= 128, "Push-constant block exceeds the spec-mandated minimum (128 bytes)");
}

void UIRenderer::Initialize(Window& window,
                            VkCommandBuffer initCmd,
                            const std::filesystem::path& spvPath,
                            std::uint16_t bindlessImageCount) {
    if (!window.descriptorHeap) {
        throw std::runtime_error("UIRenderer::Initialize: window.descriptorHeap must be set first");
    }
    window_ = &window;
    bindlessCount_ = bindlessImageCount;
    auto& heap = *window.descriptorHeap;

    // Slot allocation. Layout in the resource heap (image-typed indexing):
    //   [outImageBase_ ..]    : Window::numFrames swapchain views (storage)
    //   [atlasImageSlot_]     : 1 sampled SDF atlas
    //   [bindlessBase_  ..]   : bindlessImageCount user image slots
    auto imgSlots = heap.AllocateImageSlots(
        Window::numFrames + 1 + bindlessImageCount
    );
    outImageBase_   = imgSlots.firstElement;
    atlasImageSlot_ = imgSlots.firstElement + Window::numFrames;
    bindlessBase_   = imgSlots.firstElement + Window::numFrames + 1;

    // One SSBO per swapchain frame.
    auto bufSlots = heap.AllocateBufferSlots(Window::numFrames);
    itemBufBase_ = bufSlots.firstElement;

    // One linear sampler.
    auto sampSlots = heap.AllocateSamplerSlots(1);
    linearSamplerSlot_ = sampSlots.firstElement;

    // Initial item-buffer capacity (grows on demand).
    GrowItemBuffersIfNeeded(256);

    // Atlas image — Initialize records a layout transition into initCmd.
    atlas.Initialize(initCmd);

    CreatePipeline(spvPath);
    WriteSwapchainDescriptors();
    WriteAtlasDescriptor();
    WriteSamplerDescriptors();
    WriteItemBufferDescriptors();

    for (auto& h : heap.resourceHeap) h.FlushDevice();
    for (auto& h : heap.samplerHeap)  h.FlushDevice();
}

void UIRenderer::GrowItemBuffersIfNeeded(std::uint32_t needed) {
    if (needed <= itemCapacity_) return;
    std::uint32_t newCap = itemCapacity_ ? itemCapacity_ * 2 : 256;
    while (newCap < needed) newCap *= 2;
    itemCapacity_ = static_cast<std::uint16_t>(std::min<std::uint32_t>(newCap, 65535));

    for (auto& b : itemBufs_) {
        b.Resize(
            VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT,
            VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT,
            itemCapacity_
        );
    }
    // Item buffer descriptors point at the buffers' device addresses, so
    // they must be re-written after Resize.
    if (window_) WriteItemBufferDescriptors();
}

void UIRenderer::SetItems(std::span<const UIItem> items) {
    if (items.size() > itemCapacity_) {
        GrowItemBuffersIfNeeded(static_cast<std::uint32_t>(items.size()));
    }
    pendingItemCount = static_cast<std::uint32_t>(items.size());
    auto& buf = itemBufs_[window_->currentBuffer];
    if (!items.empty()) {
        std::memcpy(buf.value, items.data(), items.size() * sizeof(UIItem));
    }
    buf.FlushDevice();
}

void UIRenderer::Record(VkCommandBuffer cmd, std::uint32_t frameIdx, Window& window) {
    // Make sure any glyph rasterisation done during Emit lands on the GPU
    // before we sample the atlas.
    atlas.Update(cmd);

    vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, pipeline_);

    PC pc{};
    pc.itemCount             = pendingItemCount;
    pc.surfaceSize[0]        = static_cast<float>(window.width);
    pc.surfaceSize[1]        = static_cast<float>(window.height);
    #ifdef CRAFTER_GRAPHICS_WINDOW_WAYLAND
    pc.scale                 = window.scale;
    #else
    pc.scale                 = 1.0f;
    #endif
    pc.outImageHeapIdx       = outImageBase_ + frameIdx;
    // Buffer-typed shader views index the *whole* heap in buffer-descriptor
    // units, so we offset past the image region: bufferStartElement is the
    // first element index where buffer descriptors actually live.
    pc.itemBufHeapIdx        = window.descriptorHeap->bufferStartElement
                             + itemBufBase_ + frameIdx;
    pc.atlasTextureHeapIdx   = atlasImageSlot_;
    pc.bindlessBaseHeapIdx   = bindlessBase_;
    pc.linearSamplerHeapIdx  = linearSamplerSlot_;

    // Pipelines created with VK_PIPELINE_CREATE_2_DESCRIPTOR_HEAP_BIT_EXT
    // use vkCmdPushDataEXT for push constants (the spec requires layout to
    // be VK_NULL_HANDLE in that mode, which means vkCmdPushConstants has
    // nowhere to attach to).
    VkPushDataInfoEXT pushInfo{
        .sType  = VK_STRUCTURE_TYPE_PUSH_DATA_INFO_EXT,
        .offset = 0,
        .data   = { .address = &pc, .size = sizeof(PC) },
    };
    Device::vkCmdPushDataEXT(cmd, &pushInfo);

    std::uint32_t gx = (window.width  + 15) / 16;
    std::uint32_t gy = (window.height + 15) / 16;
    vkCmdDispatch(cmd, gx, gy, 1);
}

void UIRenderer::CreatePipeline(const std::filesystem::path& spvPath) {
    VulkanShader shader(spvPath, "main", VK_SHADER_STAGE_COMPUTE_BIT, nullptr);

    // Spec: "If VkPipelineCreateFlags2CreateInfoKHR::flags includes
    //       VK_PIPELINE_CREATE_2_DESCRIPTOR_HEAP_BIT_EXT, layout must be
    //       VK_NULL_HANDLE." Push constants are then attached via
    //       vkCmdPushDataEXT at draw time, not via the layout.
    VkPipelineCreateFlags2CreateInfo flags2{
        .sType = VK_STRUCTURE_TYPE_PIPELINE_CREATE_FLAGS_2_CREATE_INFO,
        .flags = VK_PIPELINE_CREATE_2_DESCRIPTOR_HEAP_BIT_EXT,
    };
    VkComputePipelineCreateInfo info{
        .sType  = VK_STRUCTURE_TYPE_COMPUTE_PIPELINE_CREATE_INFO,
        .pNext  = &flags2,
        .stage  = {
            .sType  = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO,
            .stage  = VK_SHADER_STAGE_COMPUTE_BIT,
            .module = shader.shader,
            .pName  = "main",
        },
        .layout = VK_NULL_HANDLE,
    };
    Device::CheckVkResult(vkCreateComputePipelines(
        Device::device, VK_NULL_HANDLE, 1, &info, nullptr, &pipeline_));
}

// ─── descriptor writes ───────────────────────────────────────────────────

void UIRenderer::WriteSwapchainDescriptors() {
    auto& heap = *window_->descriptorHeap;

    // One write per (frame, frame index) pairing — same swapchain view per
    // frame index for each per-frame heap copy.
    std::array<VkImageDescriptorInfoEXT, Window::numFrames * Window::numFrames> infos{};
    std::array<VkResourceDescriptorInfoEXT, Window::numFrames * Window::numFrames> resources{};
    std::array<VkHostAddressRangeEXT, Window::numFrames * Window::numFrames> destinations{};

    std::size_t k = 0;
    for (std::uint32_t heapFrame = 0; heapFrame < Window::numFrames; ++heapFrame) {
        for (std::uint32_t imgFrame = 0; imgFrame < Window::numFrames; ++imgFrame) {
            infos[k] = {
                .sType  = VK_STRUCTURE_TYPE_IMAGE_DESCRIPTOR_INFO_EXT,
                .pView  = &window_->imageViews[imgFrame],
                .layout = VK_IMAGE_LAYOUT_GENERAL,
            };
            resources[k] = {
                .sType = VK_STRUCTURE_TYPE_RESOURCE_DESCRIPTOR_INFO_EXT,
                .type  = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
                .data  = { .pImage = &infos[k] },
            };
            destinations[k] = {
                .address = heap.resourceHeap[heapFrame].value
                         + heap.ImageByteOffset(static_cast<std::uint16_t>(outImageBase_ + imgFrame)),
                .size    = Device::descriptorHeapProperties.imageDescriptorSize,
            };
            ++k;
        }
    }
    Device::vkWriteResourceDescriptorsEXT(
        Device::device, static_cast<std::uint32_t>(k),
        resources.data(), destinations.data()
    );
}

void UIRenderer::WriteAtlasDescriptor() {
    auto& heap = *window_->descriptorHeap;

    // Build a stable VkImageViewCreateInfo for the atlas. ImageVulkan
    // pre-creates a VkImageView, but the descriptor-heap path needs a
    // pointer to a create-info — keep one on the renderer so the
    // pointers we hand to vkWriteResourceDescriptorsEXT stay valid.
    atlasViewCreateInfo_ = {
        .sType    = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO,
        .image    = atlas.image.image,
        .viewType = VK_IMAGE_VIEW_TYPE_2D,
        .format   = VK_FORMAT_R8_UNORM,
        .components = {
            VK_COMPONENT_SWIZZLE_IDENTITY, VK_COMPONENT_SWIZZLE_IDENTITY,
            VK_COMPONENT_SWIZZLE_IDENTITY, VK_COMPONENT_SWIZZLE_IDENTITY,
        },
        .subresourceRange = {
            .aspectMask     = VK_IMAGE_ASPECT_COLOR_BIT,
            .baseMipLevel   = 0,
            .levelCount     = 1,
            .baseArrayLayer = 0,
            .layerCount     = 1,
        },
    };

    std::array<VkImageDescriptorInfoEXT,    Window::numFrames> infos{};
    std::array<VkResourceDescriptorInfoEXT, Window::numFrames> resources{};
    std::array<VkHostAddressRangeEXT,       Window::numFrames> destinations{};

    for (std::uint32_t f = 0; f < Window::numFrames; ++f) {
        infos[f] = {
            .sType  = VK_STRUCTURE_TYPE_IMAGE_DESCRIPTOR_INFO_EXT,
            .pView  = &atlasViewCreateInfo_,
            .layout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
        };
        resources[f] = {
            .sType = VK_STRUCTURE_TYPE_RESOURCE_DESCRIPTOR_INFO_EXT,
            .type  = VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE,
            .data  = { .pImage = &infos[f] },
        };
        destinations[f] = {
            .address = heap.resourceHeap[f].value
                     + heap.ImageByteOffset(atlasImageSlot_),
            .size    = Device::descriptorHeapProperties.imageDescriptorSize,
        };
    }
    Device::vkWriteResourceDescriptorsEXT(
        Device::device, Window::numFrames, resources.data(), destinations.data()
    );
}

void UIRenderer::WriteSamplerDescriptors() {
    auto& heap = *window_->descriptorHeap;

    VkSamplerCreateInfo info{
        .sType         = VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO,
        .magFilter     = VK_FILTER_LINEAR,
        .minFilter     = VK_FILTER_LINEAR,
        .mipmapMode    = VK_SAMPLER_MIPMAP_MODE_LINEAR,
        .addressModeU  = VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE,
        .addressModeV  = VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE,
        .addressModeW  = VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE,
        .maxAnisotropy = 1.0f,
        .minLod        = 0.0f,
        .maxLod        = VK_LOD_CLAMP_NONE,
    };
    std::array<VkSamplerCreateInfo, Window::numFrames> infos{};
    std::array<VkHostAddressRangeEXT, Window::numFrames> destinations{};
    for (std::uint32_t f = 0; f < Window::numFrames; ++f) {
        infos[f] = info;
        destinations[f] = {
            .address = heap.samplerHeap[f].value
                     + heap.SamplerByteOffset(linearSamplerSlot_),
            .size    = Device::descriptorHeapProperties.samplerDescriptorSize,
        };
    }
    Device::vkWriteSamplerDescriptorsEXT(
        Device::device, Window::numFrames, infos.data(), destinations.data()
    );
}

void UIRenderer::WriteItemBufferDescriptors() {
    auto& heap = *window_->descriptorHeap;

    std::array<VkDeviceAddressRangeEXT, Window::numFrames * Window::numFrames> ranges{};
    std::array<VkResourceDescriptorInfoEXT, Window::numFrames * Window::numFrames> resources{};
    std::array<VkHostAddressRangeEXT, Window::numFrames * Window::numFrames> destinations{};

    std::size_t k = 0;
    for (std::uint32_t heapFrame = 0; heapFrame < Window::numFrames; ++heapFrame) {
        for (std::uint32_t bufFrame = 0; bufFrame < Window::numFrames; ++bufFrame) {
            ranges[k] = {
                .address = itemBufs_[bufFrame].address,
                .size    = itemBufs_[bufFrame].size,
            };
            resources[k] = {
                .sType = VK_STRUCTURE_TYPE_RESOURCE_DESCRIPTOR_INFO_EXT,
                .type  = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
                .data  = { .pAddressRange = &ranges[k] },
            };
            destinations[k] = {
                .address = heap.resourceHeap[heapFrame].value + heap.BufferByteOffset(static_cast<std::uint16_t>(itemBufBase_ + bufFrame)),
                .size    = Device::descriptorHeapProperties.bufferDescriptorSize,
            };
            ++k;
        }
    }
    Device::vkWriteResourceDescriptorsEXT(
        Device::device, static_cast<std::uint32_t>(k),
        resources.data(), destinations.data()
    );
}

void UIRenderer::CreateLinearSampler() {
    // Not used — VK_EXT_descriptor_heap writes the sampler create-info
    // directly into the heap (see WriteSamplerDescriptors).
}