// Tier 1 demo: a user-authored compute shader dispatched alongside the
// standard ones. The custom shader inverts RGB in the area covered by a
// list of circles. The mouse-tracking circle moves; two static ones sit
// on a striped background drawn with the standard DrawQuads shader.
//
// Works on both Vulkan (native) and WebGPU (DOM). The shader source is
// in two files — inverse-circle.comp.glsl (native, SPIR-V) and
// inverse-circle.comp.wgsl (DOM). The C++ surface differs only at the
// shader-load / buffer-flag sites; the per-frame UI building logic is
// identical.

#ifndef CRAFTER_GRAPHICS_WINDOW_DOM
#include "vulkan/vulkan.h"
#endif
#include <cstddef>  // offsetof — a macro, so not visible via `import std;`

import Crafter.Graphics;
import Crafter.Event;
import std;
using namespace Crafter;

// Application-side item POD. Matches `struct InverseCircleItem { vec4
// centerRadius; }` in inverse-circle.comp.{glsl,wgsl} byte-for-byte.
struct InverseCircleItem {
    float cx, cy, radius, _pad;
};

int main() {
    Device::Initialize();
#ifdef CRAFTER_GRAPHICS_WINDOW_DOM
    static Window window(1280, 720, "Custom Shader");
#else
    Window window(1280, 720, "Custom Shader");
#endif

    auto init = window.StartInit();

    GraphicsDescriptorHeap heap;
    heap.Initialize(/*images*/ 8, /*buffers*/ 8, /*samplers*/ 4);
    window.descriptorHeap = &heap;

    UIRenderer ui;
    ui.Initialize(window, heap, init);
    window.passes.push_back(&ui);

    // Load the user-authored shader. On native it's an offline-compiled
    // SPIR-V from .comp.glsl; on DOM it's WGSL source compiled at
    // startup by the device. The DOM Load takes an extra `bindings` arg
    // declaring resources the renderer should bind at dispatch time —
    // here just one entry: the items SSBO at group 2, with its heap slot
    // read from `hdr.itemBuffer` in the push data.
    GraphicsComputeShader inverseCircle;
#ifndef CRAFTER_GRAPHICS_WINDOW_DOM
    inverseCircle.Load("inverse-circle.comp.spv");
#else
    UICustomBinding invBindings[] = {
        { .group   = 2,
          .binding = 0,
          .kind    = UICustomBindingKind::Buffer,
          ._pad    = 0,
          .pushOffset = static_cast<std::uint32_t>(offsetof(Crafter::UIDispatchHeader, itemBuffer)) },
    };
    inverseCircle.Load(std::filesystem::path("inverse-circle.comp.wgsl"), invBindings);
#endif

    // User-owned buffers.
    GraphicsBuffer<QuadItem,          true> quadsBuf;
    GraphicsBuffer<InverseCircleItem, true> invBuf;
#ifndef CRAFTER_GRAPHICS_WINDOW_DOM
    quadsBuf.Create(
        VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT,
        VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, 64);
    invBuf.Create(
        VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT,
        VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, 16);
#else
    quadsBuf.Create(64);
    invBuf.Create(16);
#endif

    auto quadsSlot = ui.RegisterBuffer(quadsBuf);
    auto invSlot   = ui.RegisterBuffer(invBuf);

    EventListener<UIBuildArgs> buildSub(&ui.onBuild, [&](UIBuildArgs a) {
        auto cmd = a.cmd;

        Rect canvas = Rect::FromWindow(window);

        // Six vertical stripes covering the canvas — gives the inverse
        // circles something visibly different to invert.
        std::array<std::array<float, 4>, 6> palette = {{
            {0.95f, 0.30f, 0.30f, 1.0f},
            {0.95f, 0.65f, 0.20f, 1.0f},
            {0.95f, 0.95f, 0.20f, 1.0f},
            {0.30f, 0.85f, 0.30f, 1.0f},
            {0.20f, 0.55f, 0.95f, 1.0f},
            {0.65f, 0.30f, 0.95f, 1.0f},
        }};
        std::uint32_t qc = 0;
        float stripeW = canvas.w / 6.0f;
        for (int i = 0; i < 6; ++i) {
            quadsBuf.value[qc++] = QuadItem{
                i * stripeW, 0, stripeW, canvas.h,
                palette[i][0], palette[i][1], palette[i][2], palette[i][3],
                0, 0, 0, 0,
                0, 0, 0, 0,
            };
        }

        // Three inverse circles: one tracking the mouse, two stationary.
        std::uint32_t ic = 0;
        invBuf.value[ic++] = { window.currentMousePos.x, window.currentMousePos.y, 100.0f, 0.0f };
        invBuf.value[ic++] = { canvas.w * 0.25f, canvas.h * 0.5f,  60.0f, 0.0f };
        invBuf.value[ic++] = { canvas.w * 0.75f, canvas.h * 0.5f,  80.0f, 0.0f };

        // Standard dispatch first — paints the stripes.
        if (qc > 0) {
#ifndef CRAFTER_GRAPHICS_WINDOW_DOM
            quadsBuf.FlushDevice(cmd, VK_ACCESS_SHADER_READ_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT);
#else
            quadsBuf.FlushDevice();
#endif
            ui.DispatchQuads(cmd, quadsSlot, qc);
        }

        // Custom dispatch second — reads the stripes, inverts under
        // circles, writes back. The library inserts the inter-dispatch
        // SHADER_WRITE → SHADER_READ|WRITE barrier automatically on
        // Vulkan; WebGPU's compute-pass tracking does it for free.
        if (ic > 0) {
#ifndef CRAFTER_GRAPHICS_WINDOW_DOM
            invBuf.FlushDevice(cmd, VK_ACCESS_SHADER_READ_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT);
#else
            invBuf.FlushDevice();
#endif
            struct PC { UIDispatchHeader hdr; } pc { ui.FillHeader(invSlot, ic) };
            std::uint32_t gx = (window.width  + 7) / 8;
            std::uint32_t gy = (window.height + 7) / 8;
            ui.Dispatch(cmd, inverseCircle, &pc, sizeof(pc), gx, gy, 1);
        }
    });

    window.FinishInit();
    window.Render();
    window.StartUpdate();
    window.StartSync();
}