/*
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;
#ifndef CRAFTER_GRAPHICS_WINDOW_DOM
#include "vulkan/vulkan.h"
#endif
export module Crafter.Graphics:UI;
import std;
import Crafter.Event;
import :Window;
import :RenderPass;
import :GraphicsTypes;
import :FontAtlas;
import :Font;
#ifndef CRAFTER_GRAPHICS_WINDOW_DOM
import :Device;
import :DescriptorHeapVulkan;
import :ImageVulkan;
import :VulkanBuffer;
import :ComputeShader;
#else
import :DescriptorHeapWebGPU;
import :WebGPU;
import :WebGPUBuffer;
import :WebGPUComputeShader;
#endif

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);

    // Reserved `flags` bit (mirrors shaders/ui-shared.glsl::UI_FLAG_CLIP).
    // FillHeader sets it when the clip rect is narrower than the surface so
    // the standard shaders can skip the per-pixel clip compares otherwise.
    // The remaining low bits stay free for user-defined feature flags.
    inline constexpr std::uint32_t kUIFlagClip = 0x80000000u;

    // ─── fused-dispatch push-constant header (issue #47) ────────────────
    // Mirrors the PC block in shaders/ui-fused.comp.glsl byte-for-byte: every
    // member is vec4-aligned, no padding holes, exactly 128 bytes (the
    // guaranteed push-constant minimum). This is a SEPARATE contract — the
    // per-element Dispatch* path and the frozen 48-byte UIDispatchHeader are
    // untouched. DispatchFused builds this; users never fill it directly.
    struct UIFusedHeader {
        std::uint32_t itemBuffers[4];  // heap slots:  quads, circles, images, text
        std::uint32_t itemCounts[4];   // item counts: quads, circles, images, text
        std::uint32_t outImage;        // swapchain image heap slot
        std::uint32_t fontTexture;     // font-atlas image slot (text phase)
        std::uint32_t fontSampler;     // font-atlas sampler slot (text phase)
        std::uint32_t flags;           // per-category clip-active bits (0x1 quads .. 0x8 text)
        std::uint32_t surfaceWidth;
        std::uint32_t surfaceHeight;
        std::uint32_t frameIdx;
        std::uint32_t _pad;            // reserved — keep zeroed
        float clipQuads[4];            // (x, y, w, h) clip per category
        float clipCircles[4];
        float clipImages[4];
        float clipText[4];
    };
    static_assert(sizeof(UIFusedHeader) == 128);

    // One category's input to UIRenderer::DispatchFused. A zero itemCount
    // makes that category a free no-op (its kernel loop is zero-trip).
    // clipRectPx defaults to "no clip", matching the Dispatch* default.
    struct FusedBatch {
        std::uint32_t       bufferSlot = 0;
        std::uint32_t       itemCount  = 0;
        std::array<float,4> clipRectPx = {0.0f, 0.0f, 1e9f, 1e9f};
    };

    // ─── standard item PODs (match GLSL std430) ─────────────────────────
    struct QuadItem {
        float x, y, w, h;
        float r, g, b, a;
        float cTL, cTR, cBR, cBL;
        float outline, oR, oG, oB;
    };
    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);

    // Where the X coordinate sits relative to the emitted glyph run. Used by
    // UIRenderer::ShapeText and the Tier 3 text components (DrawText). Lives
    // here (not in :UIComponents) so ShapeText can fold the alignment shift
    // into its single pass.
    enum class TextAlign : std::uint8_t { Left, Center, Right };

    // ─── tiny rect-carving helper ───────────────────────────────────────
    struct Rect {
        float x = 0, y = 0, w = 0, h = 0;

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

        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 {
        GraphicsCommandBuffer cmd;
        std::uint32_t         frameIdx;
    };

    // ─── UIRenderer ─────────────────────────────────────────────────────
    class UIRenderer : public RenderPass {
    public:
        GraphicsComputeShader drawQuads;
        GraphicsComputeShader drawCircles;
        GraphicsComputeShader drawImages;
        GraphicsComputeShader drawText;
        GraphicsComputeShader drawFused;

        FontAtlas* fontAtlas = nullptr;

        Crafter::Event<UIBuildArgs> onBuild;

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

        void Initialize(Window& window, GraphicsDescriptorHeap& heap, GraphicsCommandBuffer 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",
                        std::filesystem::path fusedSpv   = "ui-fused.comp.spv");

        void Record(GraphicsCommandBuffer cmd, std::uint32_t frameIdx, Window& window) override;

#ifndef CRAFTER_GRAPHICS_WINDOW_DOM
        // A UI pass draws into the swapchain image entirely from compute
        // shaders (see Dispatch), so the frame loop's inter-pass / frame-edge
        // barriers only need to synchronise the COMPUTE_SHADER stage for it.
        VkPipelineStageFlags SwapchainStage() const override {
            return VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT;
        }
#endif // !CRAFTER_GRAPHICS_WINDOW_DOM

        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;

        // Builds the header `flags` word: the caller's user flags with the
        // reserved kUIFlagClip bit set iff `clipRectPx` does not already cover
        // the whole surface. When the clip rect spans the surface, the four
        // per-pixel clip compares in uiResolveScreenPixel can never reject an
        // in-surface pixel, so the bit stays clear and the shaders skip them.
        // Static + dimension-parameterised so the decision is unit-testable
        // without a live Window; FillHeader feeds it the current window size.
        static std::uint32_t ClipFlags(std::array<float,4> clipRectPx,
                                       std::uint32_t surfaceWidth,
                                       std::uint32_t surfaceHeight,
                                       std::uint32_t flags) noexcept;

        void DispatchQuads(GraphicsCommandBuffer cmd, std::uint32_t bufferSlot, std::uint32_t itemCount,
                           std::array<float,4> clipRectPx = {0.0f, 0.0f, 1e9f, 1e9f});
        void DispatchCircles(GraphicsCommandBuffer cmd, std::uint32_t bufferSlot, std::uint32_t itemCount,
                             std::array<float,4> clipRectPx = {0.0f, 0.0f, 1e9f, 1e9f});
        void DispatchImages(GraphicsCommandBuffer cmd, std::uint32_t bufferSlot, std::uint32_t itemCount,
                            std::array<float,4> clipRectPx = {0.0f, 0.0f, 1e9f, 1e9f});
#ifdef CRAFTER_GRAPHICS_WINDOW_DOM
        // WebGPU-only overload. WebGPU bind groups can only carry one
        // texture/sampler per dispatch, so all items in `bufferSlot`
        // share the same texture (`imageSlot`) and sampler (`samplerSlot`).
        // The per-item `slots` field in ImageItem is ignored on this
        // backend. On Vulkan the bindless heap resolves per-item slots,
        // so the cross-backend path is to call the 4-arg overload above
        // on native and this 6-arg overload on DOM.
        void DispatchImages(GraphicsCommandBuffer cmd, std::uint32_t bufferSlot, std::uint32_t itemCount,
                            std::uint16_t imageSlot, std::uint16_t samplerSlot,
                            std::array<float,4> clipRectPx = {0.0f, 0.0f, 1e9f, 1e9f});
#endif
        void DispatchText(GraphicsCommandBuffer cmd, std::uint32_t bufferSlot, std::uint32_t itemCount,
                          std::array<float,4> clipRectPx = {0.0f, 0.0f, 1e9f, 1e9f});

        // Fused UI dispatch (issue #47). Composites up to four standard
        // categories — quads → circles → images → text, canonical
        // back-to-front order — in ONE compute dispatch that loads the
        // destination image once and stores once, eliminating the per-pass
        // load+store and the inter-pass memory barriers a run of consecutive
        // Dispatch* calls would pay. The win materialises at >= 2 non-empty
        // categories; an empty category (itemCount 0) is a free no-op.
        //
        // Additive: the per-element Dispatch* calls and UIDispatchHeader are
        // untouched. Keep using those for single-category or fine-grained
        // work. To interleave a custom ui.Dispatch() between standard passes,
        // bracket it with two DispatchFused calls — the dispatch boundary is
        // the explicit, app-declared flush point.
        //
        // The text category uses the registered font atlas (set fontAtlas
        // before Initialize); passing a non-empty text batch without one
        // throws, exactly like DispatchText. On the WebGPU backend this falls
        // back to issuing the per-element Dispatch* calls in canonical order
        // (one texture/sampler per dispatch there — see DispatchImages),
        // so the result matches; the load/store fusion is Vulkan-only.
        void DispatchFused(GraphicsCommandBuffer cmd,
                           const FusedBatch& quads,
                           const FusedBatch& circles,
                           const FusedBatch& images,
                           const FusedBatch& text);

        // Generic dispatch for user-authored shaders. On Vulkan, `shader` is
        // a SPIR-V compute pipeline (bindless via VK_EXT_descriptor_heap, so
        // any resource indices baked into push data resolve through the
        // global heap). On DOM, `shader` carries a UICustomBinding list
        // declared at Load time; the renderer reads the listed slot uints
        // out of `push`, resolves them against heap.bufferTable /
        // imageTable / samplerTable, and builds the bind groups before
        // dispatching.
        void Dispatch(GraphicsCommandBuffer cmd, const GraphicsComputeShader& 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 registers the GPU handle.
        // Returns a move-only BufferSlot RAII handle.
        template<typename T, bool Mapped>
        BufferSlot RegisterBuffer(GraphicsBuffer<T, Mapped>& buffer);

#ifndef CRAFTER_GRAPHICS_WINDOW_DOM
        template<typename Pixel>
        ImageSlot RegisterImage(ImageVulkan<Pixel>& image, VkFormat format,
                                VkImageLayout layout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);

        SamplerSlot RegisterSampler(const VkSamplerCreateInfo& info);
#endif
        SamplerSlot RegisterLinearClampSampler();

        // Shapes `utf8` into `out` (up to `outCapacity` glyphs), positioning
        // the run at `(x, baselineY)` with horizontal `align`. Returns the
        // number of glyphs written; `*outAdvance` (if non-null) receives the
        // full advance width regardless of truncation.
        //
        // The origin-relative shaped run is memoized in a CPU cache keyed by
        // (Font*, pxSize, color, utf8); a hit skips UTF-8 decode, glyph
        // lookup, and layout, translating the cached run into `out` instead.
        // Caching is byte-equivalent — the glyph buffer is fully rewritten
        // and re-flushed every frame regardless. Call InvalidateFont before
        // destroying a Font (the cache keys on the Font pointer).
        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,
                                TextAlign align = TextAlign::Left);

        // Drop every cached shaped run that references `font`. Must be called
        // before a Font is destroyed: the run cache (and FontAtlas's glyph
        // cache) key on the Font pointer, so a Font freed and a new one
        // reallocated at the same address would otherwise alias stale runs.
        void InvalidateFont(const Font& font) noexcept;

        // Shaped-run cache introspection (issue #123). Hits/misses are
        // byte-identical and a re-shape of an already-rasterised string leaves
        // the atlas clean, so the LRU eviction policy isn't observable through
        // output alone — these let tests assert that the cache stays bounded
        // and that the hot set survives an overflow. IsShapedRunCached is a
        // pure query: it does not bump recency.
        std::size_t ShapedRunCacheSize() const noexcept { return shapedRuns_.size(); }
        bool IsShapedRunCached(const Font& font, float pxSize,
                               std::array<float,4> color,
                               std::string_view utf8) const;

        std::uint16_t FontAtlasImageSlot()   const noexcept { return fontAtlasImageSlot_; }
        std::uint16_t FontAtlasSamplerSlot() const noexcept { return fontAtlasSamplerSlot_; }

        std::uint16_t OutImageSlot() const noexcept { return outImageSlot_; }

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

        // ─── shaped-run cache ───────────────────────────────────────────
        // Glyphs are stored origin-relative (pen at x=0, baseline=0); a hit
        // only needs a translate by (x + alignShift, baselineY). The full
        // string is part of the key (compared on lookup) so hash collisions
        // can't return the wrong run.
        struct ShapedRunKey {
            const Font*          font;
            float                pxSize;
            std::array<float, 4> color;
            std::string          text;
            bool operator==(const ShapedRunKey&) const = default;
        };
        struct ShapedRunKeyHash {
            std::size_t operator()(const ShapedRunKey& k) const noexcept {
                std::size_t h = std::hash<const void*>{}(k.font);
                auto mix = [&h](std::size_t v) noexcept {
                    h ^= v + 0x9e3779b97f4a7c15ULL + (h << 6) + (h >> 2);
                };
                mix(std::hash<std::uint32_t>{}(std::bit_cast<std::uint32_t>(k.pxSize)));
                for (float c : k.color)
                    mix(std::hash<std::uint32_t>{}(std::bit_cast<std::uint32_t>(c)));
                mix(std::hash<std::string_view>{}(k.text));
                return h;
            }
        };
        struct ShapedRun {
            std::vector<GlyphItem> glyphs;       // origin-relative
            float                  advance = 0;
            // Position of this run's key in shapedRunsLru_ (front = most
            // recently used). Lets a hit splice the entry to the front and an
            // overflow pop the least-recently-used entry, both O(1).
            std::list<const ShapedRunKey*>::iterator lruIt{};
        };
        // Soft cap. A pathological caller drawing a fresh string every frame
        // would grow this without bound; on overflow we evict the single
        // least-recently-used run (issue #123) rather than clearing the whole
        // cache. The hot set of stable labels — reshaped every frame, so always
        // near the front — stays resident, while a churn of unique strings
        // (FPS counters, timers) only recycles the cold tail. Correctness is
        // unaffected either way; this just avoids the periodic full-UI reshape
        // spike the old clear() caused.
        static constexpr std::size_t kMaxShapedRuns = 8192;
        std::unordered_map<ShapedRunKey, ShapedRun, ShapedRunKeyHash> shapedRuns_;
        // LRU recency order for shapedRuns_; front = most recently used. Holds
        // pointers to the keys owned by the map — stable across rehash because
        // unordered_map is node-based — so eviction never re-hashes the world.
        std::list<const ShapedRunKey*> shapedRunsLru_;

        ImageSlot   outImageSlot_;
        ImageSlot   fontAtlasImageSlot_;
        SamplerSlot fontAtlasSamplerSlot_;

#ifndef CRAFTER_GRAPHICS_WINDOW_DOM
        VkImageViewCreateInfo atlasViewCreateInfo_{};
        bool firstDispatchThisFrame_ = true;
#endif

        Crafter::EventListener<void> resizeSub_;

#ifndef CRAFTER_GRAPHICS_WINDOW_DOM
        void WriteSwapchainDescriptors();
        void WriteFontAtlasDescriptor();
        void WriteBufferDescriptor(std::uint16_t slot, VkDeviceAddress address, std::uint32_t size);
        void WriteSampledImageDescriptor(std::uint16_t slot,
                                         const VkImageViewCreateInfo& viewInfo,
                                         VkImageLayout layout);
#endif
    };

    // ─── template-method implementations ────────────────────────────────
    template<typename T, bool Mapped>
    BufferSlot UIRenderer::RegisterBuffer(GraphicsBuffer<T, Mapped>& buffer) {
        auto range = heap_->AllocateBufferSlots(1);
#ifndef CRAFTER_GRAPHICS_WINDOW_DOM
        WriteBufferDescriptor(range.firstElement, buffer.address, buffer.size);
#else
        heap_->bufferTable[range.firstElement] = buffer.handle;
#endif
        return BufferSlot{heap_, range.firstElement};
    }

#ifndef CRAFTER_GRAPHICS_WINDOW_DOM
    template<typename Pixel>
    ImageSlot UIRenderer::RegisterImage(ImageVulkan<Pixel>& image, VkFormat format,
                                        VkImageLayout layout) {
        auto range = heap_->AllocateImageSlots(1);
        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};
    }
#endif
}