/*
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 :GraphicsTypes;
import std;

using namespace Crafter;

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

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

    // Load the four standard pipelines plus the fused uber-kernel (issue #47).
    drawQuads.Load(quadsSpv);
    drawCircles.Load(circlesSpv);
    drawImages.Load(imagesSpv);
    drawText.Load(textSpv);
    drawFused.Load(fusedSpv);

    // 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();
    }

    // No whole-heap flush here: WriteSwapchainDescriptors / WriteFontAtlasDescriptor
    // / RegisterSampler each flush their own written byte range above, so the GPU
    // already sees everything written during setup.

    // 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.
    // WriteSwapchainDescriptors flushes its own range, so no extra flush needed.
    resizeSub_.SetEvent(&window.onResize, [this]() {
        WriteSwapchainDescriptors();
    });

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

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

void UIRenderer::Record(GraphicsCommandBuffer 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;
}

// ─── 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(GraphicsCommandBuffer 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(GraphicsCommandBuffer 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(GraphicsCommandBuffer 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(GraphicsCommandBuffer 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);
}

// ─── fused dispatch (issue #47) ─────────────────────────────────────────
//
// One uber-kernel composites quads → circles → images → text into the output
// image with a single load and a single store. Routed through the generic
// Dispatch() so it still gets the standard write-after-write barrier when it
// is not the first dispatch of the frame, and a following dispatch barriers
// against it — the only barriers it removes are the INTRA-fuse ones a run of
// Dispatch* calls would have inserted between its categories.

void UIRenderer::DispatchFused(GraphicsCommandBuffer cmd,
                               const FusedBatch& quads,
                               const FusedBatch& circles,
                               const FusedBatch& images,
                               const FusedBatch& text) {
    // Nothing to draw → no dispatch (and, crucially, no barrier bump).
    if (quads.itemCount == 0 && circles.itemCount == 0 &&
        images.itemCount == 0 && text.itemCount == 0) {
        return;
    }

    if (text.itemCount != 0 && !fontAtlasImageSlot_) {
        throw std::runtime_error("UIRenderer::DispatchFused: text batch given but no FontAtlas registered (set fontAtlas before Initialize)");
    }
    // Flush any glyphs ShapeText rasterised during this onBuild, mirroring
    // DispatchText — only relevant when there is text to draw.
    if (text.itemCount != 0 && fontAtlas != nullptr && fontAtlas->dirty) {
        fontAtlas->Update(cmd);
    }

    UIFusedHeader pc{};
    pc.itemBuffers[0] = quads.bufferSlot;
    pc.itemBuffers[1] = circles.bufferSlot;
    pc.itemBuffers[2] = images.bufferSlot;
    pc.itemBuffers[3] = text.bufferSlot;
    pc.itemCounts[0]  = quads.itemCount;
    pc.itemCounts[1]  = circles.itemCount;
    pc.itemCounts[2]  = images.itemCount;
    pc.itemCounts[3]  = text.itemCount;
    pc.outImage       = outImageSlot_;
    pc.fontTexture    = fontAtlasImageSlot_;
    pc.fontSampler    = fontAtlasSamplerSlot_;
    // Per-category clip-active bits, reusing the same surface-coverage test the
    // standard path uses (ClipFlags). A category whose clip already covers the
    // surface clears its bit so the kernel skips that category's per-pixel clip
    // compares (the common full-surface case) — pixel-identical either way.
    auto clipBit = [&](const FusedBatch& b, std::uint32_t bit) -> std::uint32_t {
        return (ClipFlags(b.clipRectPx, window_->width, window_->height, 0) & kUIFlagClip) ? bit : 0u;
    };
    pc.flags          = clipBit(quads,   0x1u) | clipBit(circles, 0x2u)
                      | clipBit(images,  0x4u) | clipBit(text,    0x8u);
    pc.surfaceWidth   = window_->width;
    pc.surfaceHeight  = window_->height;
    pc.frameIdx       = window_->currentBuffer;
    pc._pad           = 0;
    std::copy(quads.clipRectPx.begin(),   quads.clipRectPx.end(),   pc.clipQuads);
    std::copy(circles.clipRectPx.begin(), circles.clipRectPx.end(), pc.clipCircles);
    std::copy(images.clipRectPx.begin(),  images.clipRectPx.end(),  pc.clipImages);
    std::copy(text.clipRectPx.begin(),    text.clipRectPx.end(),    pc.clipText);

    Dispatch(cmd, drawFused, &pc, sizeof(pc),
             TilesFor(window_->width), TilesFor(window_->height), 1u);
}

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

void UIRenderer::Dispatch(GraphicsCommandBuffer cmd, const GraphicsComputeShader& shader,
                          const void* push, std::uint32_t pushBytes,
                          std::uint32_t gx, std::uint32_t gy, std::uint32_t gz) {
    if (!firstDispatchThisFrame_) {
        // Every UI pass read-modify-writes the same swapchain storage image
        // (later passes overdraw earlier ones), so each dispatch must observe
        // the previous dispatch's writes. The only resource the hazard touches
        // is that one image, so scope the dependency to its subresource with a
        // VkImageMemoryBarrier rather than a queue-wide VkMemoryBarrier: same
        // correctness, but only this image's shader caches are flushed —
        // unrelated buffers/images are no longer invalidated. The image stays
        // in VK_IMAGE_LAYOUT_GENERAL (bound as a storage image), so this is a
        // pure memory dependency, no layout transition.
        //
        // The execution dependency is still COMPUTE->COMPUTE over the whole
        // queue, so the passes do NOT overlap — this is a narrower cache flush
        // only. Eliminating the barriers entirely requires fusing the passes
        // into one dispatch (tracked in #47); do not attempt that here.
        VkImageMemoryBarrier ib {
            .sType               = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER,
            .srcAccessMask       = VK_ACCESS_SHADER_WRITE_BIT,
            .dstAccessMask       = VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT,
            .oldLayout           = VK_IMAGE_LAYOUT_GENERAL,
            .newLayout           = VK_IMAGE_LAYOUT_GENERAL,
            .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
            .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
            .image               = window_->images[window_->currentBuffer],
            .subresourceRange    = {
                .aspectMask     = VK_IMAGE_ASPECT_COLOR_BIT,
                .baseMipLevel   = 0,
                .levelCount     = 1,
                .baseArrayLayer = 0,
                .layerCount     = 1,
            },
        };
        vkCmdPipelineBarrier(cmd,
            VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
            VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
            0, 0, nullptr, 0, nullptr, 1, &ib);
    }
    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()
    );
    // Only the one image descriptor at outImageSlot_ changed in each per-frame
    // heap (same slot offset in every frame), so flush just that byte range
    // rather than the whole multi-KB heap.
    const std::uint32_t off = heap_->ImageByteOffset(outImageSlot_);
    const std::uint32_t sz  = Device::descriptorHeapProperties.imageDescriptorSize;
    for (auto& h : heap_->resourceHeap) h.FlushDevice(off, sz);
}

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()
    );
    // One image descriptor written per heap, at the same slot offset — flush
    // just that range rather than the whole heap.
    const std::uint32_t off = heap_->ImageByteOffset(slot);
    const std::uint32_t sz  = Device::descriptorHeapProperties.imageDescriptorSize;
    for (auto& h : heap_->resourceHeap) h.FlushDevice(off, sz);
}

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()
    );
    // One buffer descriptor written per heap, at the same slot offset — flush
    // just that range rather than the whole heap.
    const std::uint32_t off = heap_->BufferByteOffset(slot);
    const std::uint32_t sz  = Device::descriptorHeapProperties.bufferDescriptorSize;
    for (auto& h : heap_->resourceHeap) h.FlushDevice(off, sz);
}

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()
    );
    // One sampler descriptor written per heap, at the same slot offset — flush
    // just that range rather than the whole sampler heap.
    const std::uint32_t off = heap_->SamplerByteOffset(range.firstElement);
    const std::uint32_t sz  = Device::descriptorHeapProperties.samplerDescriptorSize;
    for (auto& h : heap_->samplerHeap) h.FlushDevice(off, sz);
    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);
}