/*
Crafter®.Graphics
Copyright (C) 2026 Catcrafts®
catcrafts.net
*/

// WebGPU UIRenderer implementation — DOM mode parallel to
// Crafter.Graphics-UI.cpp. Compute pipelines and bind groups live JS-side
// in additional/dom-webgpu.js; this file just translates UIRenderer's
// public method calls into wgpu* import calls.

module Crafter.Graphics:UI_webgpu_impl;
import :UI;
import :Window;
import :Font;
import :FontAtlas;
import :WebGPU;
import :WebGPUBuffer;
import :DescriptorHeapWebGPU;
import :GraphicsTypes;
import std;

using namespace Crafter;

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*/) {
    (void)initCmd;
    window_ = &window;
    heap_   = &heap;

    // The JS bridge owns the compute pipelines (4 of them, precompiled at
    // page load). The C++ side has nothing to load. We do reserve one
    // image slot for the swapchain output (kept symmetrical with Vulkan)
    // and register the font atlas if one was set.
    auto outRange = heap_->AllocateImageSlots(1);
    outImageSlot_ = ImageSlot{heap_, outRange.firstElement};
    // No JS texture handle stored at this slot — the canvas's
    // getCurrentTexture() is per-frame and the JS bridge resolves it
    // internally. The slot value is only meaningful through
    // UIDispatchHeader.outImage, which the WGSL shaders ignore.

    if (fontAtlas != nullptr) {
        auto atlasImg = heap_->AllocateImageSlots(1);
        fontAtlasImageSlot_ = ImageSlot{heap_, atlasImg.firstElement};
        heap_->imageTable[atlasImg.firstElement] = fontAtlas->textureHandle;
        fontAtlasSamplerSlot_ = RegisterLinearClampSampler();
    }

    resizeSub_.SetEvent(&window.onResize, [this]() {
        // Storage textures + bind groups are recreated JS-side at next
        // wgpuFrameBegin (via the ensureSized path). Nothing to do here.
    });
}

void UIRenderer::Record(GraphicsCommandBuffer cmd, std::uint32_t frameIdx, Window& window) {
    if (fontAtlas != nullptr && fontAtlas->dirty) {
        fontAtlas->Update(cmd);
    }
    onBuild.Invoke({cmd, frameIdx});
    if (fontAtlas != nullptr && fontAtlas->dirty) {
        fontAtlas->Update(cmd);
    }
    (void)window;
}

namespace {
    inline std::uint32_t TilesFor(std::uint32_t dim) { return (dim + 7u) / 8u; }
}

void UIRenderer::DispatchQuads(GraphicsCommandBuffer /*cmd*/, std::uint32_t bufferSlot,
                               std::uint32_t itemCount,
                               std::array<float,4> clipRectPx) {
    if (itemCount == 0) return;
    UIDispatchHeader hdr = FillHeader(bufferSlot, itemCount, clipRectPx);
    auto handle = heap_->bufferTable[bufferSlot];
    WebGPU::wgpuDispatchQuads(handle, &hdr,
        static_cast<std::int32_t>(TilesFor(window_->width)),
        static_cast<std::int32_t>(TilesFor(window_->height)));
}

void UIRenderer::DispatchCircles(GraphicsCommandBuffer /*cmd*/, std::uint32_t bufferSlot,
                                 std::uint32_t itemCount,
                                 std::array<float,4> clipRectPx) {
    if (itemCount == 0) return;
    UIDispatchHeader hdr = FillHeader(bufferSlot, itemCount, clipRectPx);
    auto handle = heap_->bufferTable[bufferSlot];
    WebGPU::wgpuDispatchCircles(handle, &hdr,
        static_cast<std::int32_t>(TilesFor(window_->width)),
        static_cast<std::int32_t>(TilesFor(window_->height)));
}

void UIRenderer::DispatchImages(GraphicsCommandBuffer /*cmd*/, std::uint32_t bufferSlot,
                                std::uint32_t itemCount,
                                std::array<float,4> clipRectPx) {
    if (itemCount == 0) return;
    UIDispatchHeader hdr = FillHeader(bufferSlot, itemCount, clipRectPx);
    auto handle = heap_->bufferTable[bufferSlot];
    // Backward-compatible fallback: callers that don't pass a texture
    // get the font atlas. Useful for tests, useless for real content.
    // New code should use the 6-arg overload below.
    if (fontAtlasImageSlot_) {
        auto texHandle  = heap_->imageTable[fontAtlasImageSlot_];
        auto sampHandle = heap_->samplerTable[fontAtlasSamplerSlot_];
        WebGPU::wgpuDispatchImages(handle, &hdr,
            static_cast<std::int32_t>(TilesFor(window_->width)),
            static_cast<std::int32_t>(TilesFor(window_->height)),
            texHandle, sampHandle);
    }
}

void UIRenderer::DispatchImages(GraphicsCommandBuffer /*cmd*/, std::uint32_t bufferSlot,
                                std::uint32_t itemCount,
                                std::uint16_t imageSlot, std::uint16_t samplerSlot,
                                std::array<float,4> clipRectPx) {
    if (itemCount == 0) return;
    UIDispatchHeader hdr = FillHeader(bufferSlot, itemCount, clipRectPx);
    auto handle     = heap_->bufferTable[bufferSlot];
    auto texHandle  = heap_->imageTable[imageSlot];
    auto sampHandle = heap_->samplerTable[samplerSlot];
    WebGPU::wgpuDispatchImages(handle, &hdr,
        static_cast<std::int32_t>(TilesFor(window_->width)),
        static_cast<std::int32_t>(TilesFor(window_->height)),
        texHandle, sampHandle);
}

void UIRenderer::DispatchText(GraphicsCommandBuffer /*cmd*/, std::uint32_t bufferSlot,
                              std::uint32_t itemCount,
                              std::array<float,4> clipRectPx) {
    if (itemCount == 0) return;
    if (!fontAtlasImageSlot_) {
        std::println("UIRenderer::DispatchText: no FontAtlas registered (set fontAtlas before Initialize)");
        std::abort();
    }
    UIDispatchHeader hdr = FillHeader(bufferSlot, itemCount, clipRectPx);
    auto bufHandle  = heap_->bufferTable[bufferSlot];
    auto texHandle  = heap_->imageTable[fontAtlasImageSlot_];
    auto sampHandle = heap_->samplerTable[fontAtlasSamplerSlot_];
    WebGPU::wgpuDispatchText(bufHandle, &hdr,
        static_cast<std::int32_t>(TilesFor(window_->width)),
        static_cast<std::int32_t>(TilesFor(window_->height)),
        texHandle, sampHandle);
}

SamplerSlot UIRenderer::RegisterLinearClampSampler() {
    auto range = heap_->AllocateSamplerSlots(1);
    heap_->samplerTable[range.firstElement] = WebGPU::wgpuCreateLinearClampSampler();
    return SamplerSlot{heap_, range.firstElement};
}

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) {
    // For each user-declared binding, read the slot uint32 out of push
    // data at the recorded offset, look up the GPU handle in the heap,
    // and assemble a list of handles in the same order the JS bridge
    // expects (matching shader.customBindings).
    std::vector<std::uint32_t> handles;
    handles.reserve(shader.customBindings.size());
    const std::uint8_t* p = static_cast<const std::uint8_t*>(push);
    for (const auto& b : shader.customBindings) {
        if (b.pushOffset + sizeof(std::uint32_t) > pushBytes) {
            std::println("UIRenderer::Dispatch: binding pushOffset {} out of bounds (push={})",
                         b.pushOffset, pushBytes);
            return;
        }
        std::uint32_t slot;
        std::memcpy(&slot, p + b.pushOffset, sizeof(slot));
        std::uint32_t handle = 0;
        switch (b.kind) {
            case UICustomBindingKind::Buffer:
                if (slot < heap_->bufferTable.size()) handle = heap_->bufferTable[slot];
                break;
            case UICustomBindingKind::SampledTexture:
                if (slot < heap_->imageTable.size()) handle = heap_->imageTable[slot];
                break;
            case UICustomBindingKind::Sampler:
                if (slot < heap_->samplerTable.size()) handle = heap_->samplerTable[slot];
                break;
            default: break;
        }
        handles.push_back(handle);
    }
    WebGPU::wgpuDispatchCustom(shader.pipelineHandle,
                               push, static_cast<std::int32_t>(pushBytes),
                               handles.data(), static_cast<std::int32_t>(handles.size()),
                               static_cast<std::int32_t>(gx),
                               static_cast<std::int32_t>(gy),
                               static_cast<std::int32_t>(gz));
}