// 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.

#include "vulkan/vulkan.h"

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 byte-for-byte.
struct InverseCircleItem {
    float cx, cy, radius, _pad;
};

int main() {
    Device::Initialize();
    Window window(1280, 720, "Custom Shader");

    VkCommandBuffer init = window.StartInit();

    DescriptorHeapVulkan 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. Same wrapper as the four shipped with
    // the library — there is no privileged path.
    ComputeShader inverseCircle;
    inverseCircle.Load("inverse-circle.comp.spv");

    // User-owned buffers.
    VulkanBuffer<QuadItem,           true> quadsBuf;
    VulkanBuffer<InverseCircleItem,  true> invBuf;
    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);

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

    EventListener<UIBuildArgs> buildSub(&ui.onBuild, [&](UIBuildArgs a) {
        VkCommandBuffer 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) {
            quadsBuf.FlushDevice(cmd, VK_ACCESS_SHADER_READ_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT);
            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.
        if (ic > 0) {
            invBuf.FlushDevice(cmd, VK_ACCESS_SHADER_READ_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT);
            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();
}