/*
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:UI_impl;
import :UI;
import :ComputeShader;
import :Device;
import :Window;
import :DescriptorHeapVulkan;
import :ImageVulkan;
import :VulkanBuffer;
import :FontAtlas;
import :Font;
import std;

using namespace Crafter;

// ─── Initialize ─────────────────────────────────────────────────────────

void UIRenderer::Initialize(Window& window, DescriptorHeapVulkan& heap, VkCommandBuffer initCmd,
                            std::filesystem::path quadsSpv,
                            std::filesystem::path circlesSpv,
                            std::filesystem::path imagesSpv,
                            std::filesystem::path textSpv) {
    window_ = &window;
    heap_   = &heap;

    // Load the four standard pipelines.
    drawQuads.Load(quadsSpv);
    drawCircles.Load(circlesSpv);
    drawImages.Load(imagesSpv);
    drawText.Load(textSpv);

    // Allocate one image slot for the swapchain output. Each per-frame heap
    // copy will hold ITS frame's image at this slot.
    auto outRange = heap_->AllocateImageSlots(1);
    outImageSlot_ = ImageSlot{heap_, outRange.firstElement};

    WriteSwapchainDescriptors();

    // Optional font-atlas registration (user must have called atlas->Initialize
    // already before reaching here, so atlas->image is live).
    if (fontAtlas != nullptr) {
        auto atlasImg = heap_->AllocateImageSlots(1);
        fontAtlasImageSlot_   = ImageSlot{heap_, atlasImg.firstElement};
        fontAtlasSamplerSlot_ = RegisterLinearClampSampler();
        WriteFontAtlasDescriptor();
    }

    // Flush the host-mapped descriptor heaps so the GPU sees what we wrote.
    for (auto& h : heap_->resourceHeap) h.FlushDevice();
    for (auto& h : heap_->samplerHeap)  h.FlushDevice();

    // Re-bind the swapchain image descriptors after every resize. Window
    // already drained the queue and rebuilt the imageViews[] before
    // firing the event, so vkWriteResourceDescriptorsEXT is safe here.
    resizeSub_.SetEvent(&window.onResize, [this]() {
        WriteSwapchainDescriptors();
        for (auto& h : heap_->resourceHeap) h.FlushDevice();
    });

    (void)initCmd; // reserved for future image-layout tweaks
}

// ─── per-frame Record ───────────────────────────────────────────────────

void UIRenderer::Record(VkCommandBuffer cmd, std::uint32_t frameIdx, Window& window) {
    // Reset per-frame state.
    firstDispatchThisFrame_ = true;

    // If text is in use, flush any glyphs that user-side ShapeText calls
    // produced during a previous frame's onBuild. (Ensure() during the
    // current onBuild also marks the atlas dirty; that's flushed on the
    // NEXT Record. For v1 this is fine because the current frame's text
    // dispatch reads whatever's already been uploaded, and brand-new glyphs
    // missing from the atlas this frame will simply render blank for one
    // frame and resolve next frame. To get them this frame, the user can
    // call atlas->Update(cmd) themselves at the top of onBuild.)
    if (fontAtlas != nullptr && fontAtlas->dirty) {
        fontAtlas->Update(cmd);
    }

    onBuild.Invoke({cmd, frameIdx});

    (void)window;
}

// ─── header builder ─────────────────────────────────────────────────────

UIDispatchHeader UIRenderer::FillHeader(std::uint32_t itemBufferSlot,
                                        std::uint32_t itemCount,
                                        std::array<float,4> clipRectPx,
                                        std::uint32_t flags) const noexcept {
    UIDispatchHeader h{};
    h.outImage      = outImageSlot_;
    h.itemBuffer    = itemBufferSlot;
    h.surfaceWidth  = window_->width;
    h.surfaceHeight = window_->height;
    h.clipX = clipRectPx[0];
    h.clipY = clipRectPx[1];
    h.clipW = clipRectPx[2];
    h.clipH = clipRectPx[3];
    h.itemCount = itemCount;
    h.frameIdx  = window_->currentBuffer;
    h.flags     = flags;
    h._pad      = 0;
    return h;
}

// ─── group-count helper ────────────────────────────────────────────────

namespace {
    // Number of 8-pixel tiles needed to cover `dim` pixels (rounded up).
    inline std::uint32_t TilesFor(std::uint32_t dim) {
        return (dim + 7u) / 8u;
    }
}

// ─── standard-shader convenience dispatches ─────────────────────────────
//
// All four standard shaders use the same pixel-tile dispatch model: one
// workgroup per 8×8 screen tile, each thread iterates every item in order
// inside the workgroup, accumulating into a local register. This guarantees
// "items in the buffer render in order" (later items overdraw earlier ones)
// without inter-workgroup races on imageLoad/imageStore — the bug that the
// per-item dispatch model had.

void UIRenderer::DispatchQuads(VkCommandBuffer cmd, std::uint32_t bufferSlot,
                               std::uint32_t itemCount,
                               std::array<float,4> clipRectPx) {
    if (itemCount == 0) return;
    struct PC { UIDispatchHeader hdr; } pc { FillHeader(bufferSlot, itemCount, clipRectPx) };
    Dispatch(cmd, drawQuads, &pc, sizeof(pc),
             TilesFor(window_->width), TilesFor(window_->height), 1u);
}

void UIRenderer::DispatchCircles(VkCommandBuffer cmd, std::uint32_t bufferSlot,
                                 std::uint32_t itemCount,
                                 std::array<float,4> clipRectPx) {
    if (itemCount == 0) return;
    struct PC { UIDispatchHeader hdr; } pc { FillHeader(bufferSlot, itemCount, clipRectPx) };
    Dispatch(cmd, drawCircles, &pc, sizeof(pc),
             TilesFor(window_->width), TilesFor(window_->height), 1u);
}

void UIRenderer::DispatchImages(VkCommandBuffer cmd, std::uint32_t bufferSlot,
                                std::uint32_t itemCount,
                                std::array<float,4> clipRectPx) {
    if (itemCount == 0) return;
    struct PC { UIDispatchHeader hdr; } pc { FillHeader(bufferSlot, itemCount, clipRectPx) };
    Dispatch(cmd, drawImages, &pc, sizeof(pc),
             TilesFor(window_->width), TilesFor(window_->height), 1u);
}

void UIRenderer::DispatchText(VkCommandBuffer cmd, std::uint32_t bufferSlot,
                              std::uint32_t itemCount,
                              std::array<float,4> clipRectPx) {
    if (itemCount == 0) return;
    if (!fontAtlasImageSlot_) {
        throw std::runtime_error("UIRenderer::DispatchText: no FontAtlas registered (set fontAtlas before Initialize)");
    }
    // Flush any glyphs that ShapeText calls (during this onBuild) just
    // rasterised, so the dispatch below sees them.
    if (fontAtlas != nullptr && fontAtlas->dirty) {
        fontAtlas->Update(cmd);
    }
    struct PC {
        UIDispatchHeader hdr;
        std::uint32_t fontTextureSlot;
        std::uint32_t fontSamplerSlot;
        std::uint32_t _p0;
        std::uint32_t _p1;
    } pc {
        FillHeader(bufferSlot, itemCount, clipRectPx),
        fontAtlasImageSlot_,
        fontAtlasSamplerSlot_,
        0, 0
    };
    Dispatch(cmd, drawText, &pc, sizeof(pc),
             TilesFor(window_->width), TilesFor(window_->height), 1u);
}

// ─── generic Dispatch (with barrier) ────────────────────────────────────

void UIRenderer::Dispatch(VkCommandBuffer cmd, const ComputeShader& shader,
                          const void* push, std::uint32_t pushBytes,
                          std::uint32_t gx, std::uint32_t gy, std::uint32_t gz) {
    if (!firstDispatchThisFrame_) {
        VkMemoryBarrier mb {
            .sType         = VK_STRUCTURE_TYPE_MEMORY_BARRIER,
            .srcAccessMask = VK_ACCESS_SHADER_WRITE_BIT,
            .dstAccessMask = VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT,
        };
        vkCmdPipelineBarrier(cmd,
            VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
            VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
            0, 1, &mb, 0, nullptr, 0, nullptr);
    }
    firstDispatchThisFrame_ = false;
    shader.Dispatch(cmd, push, pushBytes, gx, gy, gz);
}

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

void UIRenderer::WriteSwapchainDescriptors() {
    // Each per-frame heap holds ITS swapchain image at outImageSlot_.
    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  = &window_->imageViews[f],
            .layout = VK_IMAGE_LAYOUT_GENERAL,
        };
        resources[f] = {
            .sType = VK_STRUCTURE_TYPE_RESOURCE_DESCRIPTOR_INFO_EXT,
            .type  = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
            .data  = { .pImage = &infos[f] },
        };
        destinations[f] = {
            .address = heap_->resourceHeap[f].value
                     + heap_->ImageByteOffset(outImageSlot_),
            .size    = Device::descriptorHeapProperties.imageDescriptorSize,
        };
    }
    Device::vkWriteResourceDescriptorsEXT(
        Device::device, Window::numFrames, resources.data(), destinations.data()
    );
}

void UIRenderer::WriteFontAtlasDescriptor() {
    atlasViewCreateInfo_ = {
        .sType    = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO,
        .image    = fontAtlas->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,
        },
    };
    WriteSampledImageDescriptor(fontAtlasImageSlot_,
                                atlasViewCreateInfo_,
                                VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);
}

void UIRenderer::WriteSampledImageDescriptor(std::uint16_t slot,
                                             const VkImageViewCreateInfo& viewInfo,
                                             VkImageLayout layout) {
    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  = &viewInfo,
            .layout = layout,
        };
        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(slot),
            .size    = Device::descriptorHeapProperties.imageDescriptorSize,
        };
    }
    Device::vkWriteResourceDescriptorsEXT(
        Device::device, Window::numFrames, resources.data(), destinations.data()
    );
    for (auto& h : heap_->resourceHeap) h.FlushDevice();
}

void UIRenderer::WriteBufferDescriptor(std::uint16_t slot, VkDeviceAddress address, std::uint32_t size) {
    std::array<VkDeviceAddressRangeEXT, Window::numFrames>     ranges{};
    std::array<VkResourceDescriptorInfoEXT, Window::numFrames> resources{};
    std::array<VkHostAddressRangeEXT, Window::numFrames>       destinations{};

    for (std::uint32_t f = 0; f < Window::numFrames; ++f) {
        ranges[f] = { .address = address, .size = size };
        resources[f] = {
            .sType = VK_STRUCTURE_TYPE_RESOURCE_DESCRIPTOR_INFO_EXT,
            .type  = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
            .data  = { .pAddressRange = &ranges[f] },
        };
        destinations[f] = {
            .address = heap_->resourceHeap[f].value + heap_->BufferByteOffset(slot),
            .size    = Device::descriptorHeapProperties.bufferDescriptorSize,
        };
    }
    Device::vkWriteResourceDescriptorsEXT(
        Device::device, Window::numFrames, resources.data(), destinations.data()
    );
    for (auto& h : heap_->resourceHeap) h.FlushDevice();
}

SamplerSlot UIRenderer::RegisterSampler(const VkSamplerCreateInfo& info) {
    auto range = heap_->AllocateSamplerSlots(1);
    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(range.firstElement),
            .size    = Device::descriptorHeapProperties.samplerDescriptorSize,
        };
    }
    Device::vkWriteSamplerDescriptorsEXT(
        Device::device, Window::numFrames, infos.data(), destinations.data()
    );
    for (auto& h : heap_->samplerHeap) h.FlushDevice();
    return SamplerSlot{heap_, range.firstElement};
}

SamplerSlot UIRenderer::RegisterLinearClampSampler() {
    VkSamplerCreateInfo s {
        .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,
    };
    return RegisterSampler(s);
}

// ─── ShapeText ─────────────────────────────────────────────────────────

std::uint32_t UIRenderer::ShapeText(Font& font, float pxSize,
                                    float x, float baselineY,
                                    std::string_view utf8,
                                    std::array<float,4> color,
                                    GlyphItem* out, std::uint32_t outCapacity,
                                    float* outAdvance) {
    if (fontAtlas == nullptr) {
        throw std::runtime_error("UIRenderer::ShapeText: no FontAtlas (set fontAtlas before Initialize)");
    }

    const float scale = pxSize / FontAtlas::kBaseSize;
    float cursor = x;
    std::uint32_t written = 0;

    std::size_t i = 0;
    while (i < utf8.size() && written < outCapacity) {
        std::uint32_t cp = DecodeUtf8(utf8, i);
        if (cp == 0) break;
        if (cp == '\n') { /* single-line shaper — ignore */ continue; }

        fontAtlas->Ensure(font, cp);
        const Glyph* g = fontAtlas->Lookup(font, cp);
        if (g == nullptr) continue;

        // Empty glyph (whitespace) — advance only.
        if (g->w > 0 && g->h > 0) {
            GlyphItem& gi = out[written++];
            gi.x  = cursor + g->xoff * scale;
            gi.y  = baselineY + g->yoff * scale;
            gi.w  = g->w * scale;
            gi.h  = g->h * scale;
            gi.u0 = g->u0; gi.v0 = g->v0;
            gi.u1 = g->u1; gi.v1 = g->v1;
            gi.r = color[0]; gi.g = color[1]; gi.b = color[2]; gi.a = color[3];
        }
        cursor += g->advance * scale;
    }

    if (outAdvance) *outAdvance = cursor - x;
    return written;
}