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

export namespace Crafter {
    // ─── push-constant header ───────────────────────────────────────────
    // Mirrors shaders/ui-shared.glsl::UIDispatchHeader byte-for-byte. User
    // shaders MUST embed this as the first member of their push-constant
    // struct so UIRenderer::FillHeader works.
    struct UIDispatchHeader {
        std::uint32_t outImage;
        std::uint32_t itemBuffer;
        std::uint32_t surfaceWidth;
        std::uint32_t surfaceHeight;
        float         clipX, clipY, clipW, clipH;
        std::uint32_t itemCount;
        std::uint32_t frameIdx;
        std::uint32_t flags;
        std::uint32_t _pad;
    };
    static_assert(sizeof(UIDispatchHeader) == 48);

    // ─── standard item PODs (match GLSL std430) ─────────────────────────
    struct QuadItem {
        float x, y, w, h;
        float r, g, b, a;
        float cTL, cTR, cBR, cBL;       // per-corner radius in px
        float outline, oR, oG, oB;      // outline thickness + RGB
    };
    static_assert(sizeof(QuadItem) == 64);

    struct CircleItem {
        float cx, cy, radius, _p0;
        float r, g, b, a;
        float outline, oR, oG, oB;
    };
    static_assert(sizeof(CircleItem) == 48);

    struct ImageItem {
        float x, y, w, h;
        float u0, v0, u1, v1;
        float tR, tG, tB, tA;
        std::uint32_t texSlot, sampSlot, _p1, _p2;
    };
    static_assert(sizeof(ImageItem) == 64);

    struct GlyphItem {
        float x, y, w, h;
        float u0, v0, u1, v1;
        float r, g, b, a;
    };
    static_assert(sizeof(GlyphItem) == 48);

    // ─── tiny rect-carving helper ───────────────────────────────────────
    // Pure value semantics. No engine, just convenience. Skip if you'd rather
    // compute pixels yourself.
    struct Rect {
        float x = 0, y = 0, w = 0, h = 0;

        enum class Anchor { Top, Bottom, Left, Right };

        // Returns a sub-rect of `size` along the given anchor edge of self.
        // Does not modify `*this`. (Use `.Inset(...)` to drop a margin first.)
        Rect SubRect(float size, Anchor a) const noexcept {
            switch (a) {
                case Anchor::Top:    return { x, y,                 w,    std::min(size, h) };
                case Anchor::Bottom: return { x, y + h - std::min(size, h), w, std::min(size, h) };
                case Anchor::Left:   return { x, y,                 std::min(size, w), h };
                case Anchor::Right:  return { x + w - std::min(size, w), y, std::min(size, w), h };
            }
            return *this;
        }

        Rect Inset(float padding) const noexcept {
            return Inset(padding, padding, padding, padding);
        }
        Rect Inset(float top, float right, float bottom, float left) const noexcept {
            float nw = std::max(0.0f, w - left - right);
            float nh = std::max(0.0f, h - top - bottom);
            return { x + left, y + top, nw, nh };
        }

        bool Contains(float px, float py) const noexcept {
            return px >= x && px < x + w && py >= y && py < y + h;
        }

        static Rect FromWindow(const Window& win) noexcept {
            return { 0, 0, static_cast<float>(win.width), static_cast<float>(win.height) };
        }
    };

    // ─── per-frame callback args ────────────────────────────────────────
    struct UIBuildArgs {
        VkCommandBuffer cmd;
        std::uint32_t   frameIdx;
    };

    // ─── UIRenderer ─────────────────────────────────────────────────────
    // One per Window (typically). Owns the four standard compute shaders,
    // pre-allocates heap slots for the swapchain images, and exposes a thin
    // dispatch helper for both the standard shaders and user-supplied ones.
    //
    // Workflow:
    //   1. Construct, configure (set fontAtlas if drawing text).
    //   2. Initialize(window, heap, initCmd) — once, after window.descriptorHeap
    //      is set and before window.FinishInit().
    //   3. window.passes.push_back(&ui).
    //   4. Listen on `onBuild`. Inside the callback, fill your item buffers,
    //      flush them, and call DispatchQuads / DispatchCircles / DispatchImages
    //      / DispatchText / Dispatch as needed. Library inserts a SHADER_WRITE
    //      → SHADER_READ|WRITE memory barrier between consecutive dispatches.
    class UIRenderer : public RenderPass {
    public:
        // Pre-loaded standard shaders (public so users can call Dispatch
        // directly with them if they want to embed extra push-constant fields
        // beyond the standard header).
        ComputeShader drawQuads;
        ComputeShader drawCircles;
        ComputeShader drawImages;
        ComputeShader drawText;

        // Optional. If set before Initialize, the atlas is registered into a
        // sampled-image slot + linear sampler slot, and Update(cmd) is called
        // at the top of every Record() so any glyphs ensured during onBuild
        // make it to the GPU before the text dispatch reads them.
        FontAtlas* fontAtlas = nullptr;

        // User callback. Subscribe by holding a Crafter::EventListener<UIBuildArgs>:
        //     EventListener<UIBuildArgs> sub(&ui.onBuild, [&](UIBuildArgs a) { ... });
        // Listener lifetime governs the subscription.
        Crafter::Event<UIBuildArgs> onBuild;

        UIRenderer() = default;
        UIRenderer(const UIRenderer&) = delete;
        UIRenderer& operator=(const UIRenderer&) = delete;

        // Default shader paths assume Crafter.Build placed the .spv files
        // alongside the consumer binary (this is what cfg.shaders does).
        void Initialize(Window& window, DescriptorHeapVulkan& heap, VkCommandBuffer initCmd,
                        std::filesystem::path quadsSpv   = "ui-quads.comp.spv",
                        std::filesystem::path circlesSpv = "ui-circles.comp.spv",
                        std::filesystem::path imagesSpv  = "ui-images.comp.spv",
                        std::filesystem::path textSpv    = "ui-text.comp.spv");

        // RenderPass interface — invoked from Window::Render.
        void Record(VkCommandBuffer cmd, std::uint32_t frameIdx, Window& window) override;

        // ─── helpers used inside `onBuild` ─────────────────────────────

        // Builds a populated header. `clipRectPx` defaults to "no clip".
        UIDispatchHeader FillHeader(std::uint32_t itemBufferSlot,
                                    std::uint32_t itemCount,
                                    std::array<float,4> clipRectPx = {0.0f, 0.0f, 1e9f, 1e9f},
                                    std::uint32_t flags = 0) const noexcept;

        // Convenience: dispatches the named standard shader. Group count is
        // computed from the window's surface size — the standard shaders
        // dispatch one workgroup per 8×8 screen tile and iterate every item
        // in the buffer in order, so item ORDER in the buffer == draw order
        // on screen (later items overdraw earlier ones, race-free).
        void DispatchQuads(VkCommandBuffer cmd, std::uint32_t bufferSlot, std::uint32_t itemCount,
                           std::array<float,4> clipRectPx = {0.0f, 0.0f, 1e9f, 1e9f});
        void DispatchCircles(VkCommandBuffer cmd, std::uint32_t bufferSlot, std::uint32_t itemCount,
                             std::array<float,4> clipRectPx = {0.0f, 0.0f, 1e9f, 1e9f});
        void DispatchImages(VkCommandBuffer cmd, std::uint32_t bufferSlot, std::uint32_t itemCount,
                            std::array<float,4> clipRectPx = {0.0f, 0.0f, 1e9f, 1e9f});
        // For DispatchText, the font atlas image+sampler slots are taken from
        // UIRenderer's Initialize-time registration (see fontAtlasImageSlot()
        // / fontAtlasSamplerSlot()). Set `fontAtlas` before Initialize.
        void DispatchText(VkCommandBuffer cmd, std::uint32_t bufferSlot, std::uint32_t itemCount,
                          std::array<float,4> clipRectPx = {0.0f, 0.0f, 1e9f, 1e9f});

        // Generic dispatch — for user-authored shaders. Inserts the standard
        // pre-dispatch barrier (skipped on the first call per frame).
        void Dispatch(VkCommandBuffer cmd, const ComputeShader& shader,
                      const void* push, std::uint32_t pushBytes,
                      std::uint32_t gx, std::uint32_t gy = 1, std::uint32_t gz = 1);

        // Allocates a heap slot for the buffer and writes its descriptor into
        // every per-frame heap. The user's mapped buffer is shared across
        // frames — fine because Window::Render currently waits idle before
        // submitting the next frame. Returns a move-only BufferSlot handle
        // whose destructor returns the slot to the heap. Implicitly converts
        // to the absolute heap index when passed to FillHeader / Dispatch*.
        template<typename T, bool Mapped>
        BufferSlot RegisterBuffer(VulkanBuffer<T, Mapped>& buffer);

        // Same for an ImageVulkan-managed sampled image (e.g. a user texture).
        // Caller specifies the layout the image will be sampled in (typically
        // VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL).
        template<typename Pixel>
        ImageSlot RegisterImage(ImageVulkan<Pixel>& image, VkFormat format,
                                VkImageLayout layout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);

        // Allocates a sampler slot and writes a VkSamplerCreateInfo into
        // every per-frame sampler heap. v1 takes the create-info inline.
        SamplerSlot RegisterSampler(const VkSamplerCreateInfo& info);

        // Convenience: a linear-filter, clamp-to-edge sampler. Returns a
        // SamplerSlot handle. Useful for the FontAtlas and most plain image
        // sampling.
        SamplerSlot RegisterLinearClampSampler();

        // Shapes a UTF-8 string into glyph quads at (x, y) baseline. Calls
        // FontAtlas::Ensure for each codepoint (rasterising on first use),
        // emits one GlyphItem per visible glyph, returns the count written.
        // Use this to fill a GlyphItem buffer that you then dispatch.
        // Cursor advances along +X. No line-wrap, no kerning — single line.
        std::uint32_t 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 = nullptr);

        // Read after Initialize: the slot the font atlas was registered into.
        // 0xFFFF means "no atlas" (set fontAtlas before Initialize).
        std::uint16_t FontAtlasImageSlot()   const noexcept { return fontAtlasImageSlot_; }
        std::uint16_t FontAtlasSamplerSlot() const noexcept { return fontAtlasSamplerSlot_; }

    private:
        Window* window_ = nullptr;
        DescriptorHeapVulkan* heap_ = nullptr;

        // One image slot used for the swapchain output. In each per-frame
        // heap, that slot points at THAT frame's swapchain image. So the
        // shader's `uiImages[hdr.outImage]` is always the current frame's
        // swapchain image regardless of which heap is bound.
        ImageSlot outImageSlot_;

        // Stable VkImageViewCreateInfos for the descriptor heap to ingest.
        // These must outlive the write call.
        VkImageViewCreateInfo atlasViewCreateInfo_{};

        ImageSlot   fontAtlasImageSlot_;
        SamplerSlot fontAtlasSamplerSlot_;

        bool firstDispatchThisFrame_ = true;

        void WriteSwapchainDescriptors();
        void WriteFontAtlasDescriptor();

        // Helper used by RegisterBuffer template (defined in impl). Writes the
        // address-range descriptor at `slot` into all per-frame heaps.
        void WriteBufferDescriptor(std::uint16_t slot, VkDeviceAddress address, std::uint32_t size);

        // Helper used by RegisterImage template — writes a sampled-image at
        // `slot` referring to a stable VkImageViewCreateInfo (caller stores).
        void WriteSampledImageDescriptor(std::uint16_t slot,
                                         const VkImageViewCreateInfo& viewInfo,
                                         VkImageLayout layout);
    };

    // ─── template-method implementations ────────────────────────────────
    template<typename T, bool Mapped>
    BufferSlot UIRenderer::RegisterBuffer(VulkanBuffer<T, Mapped>& buffer) {
        auto range = heap_->AllocateBufferSlots(1);
        WriteBufferDescriptor(range.firstElement, buffer.address, buffer.size);
        // BufferSlot's operator uint16_t() folds in heap_->bufferStartElement,
        // so callers receive the absolute heap index when they convert.
        return BufferSlot{heap_, range.firstElement};
    }

    template<typename Pixel>
    ImageSlot UIRenderer::RegisterImage(ImageVulkan<Pixel>& image, VkFormat format,
                                        VkImageLayout layout) {
        auto range = heap_->AllocateImageSlots(1);

        // Build a stable view-create-info that lives as long as the heap reads
        // it. We co-locate it on the renderer for the font atlas; for arbitrary
        // user images we lean on the fact that vkWriteResourceDescriptorsEXT
        // copies the view descriptor immediately. (Validated by the heap spec:
        // the descriptor is materialised at write time, the create-info need
        // not persist past the call.)
        VkImageViewCreateInfo info {
            .sType    = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO,
            .image    = image.image,
            .viewType = VK_IMAGE_VIEW_TYPE_2D,
            .format   = format,
            .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     = image.mipLevels,
                .baseArrayLayer = 0,
                .layerCount     = 1,
            },
        };
        WriteSampledImageDescriptor(range.firstElement, info, layout);
        return ImageSlot{heap_, range.firstElement};
    }
}