//SPDX-License-Identifier: MIT
//SPDX-FileCopyrightText: Copyright (C) 2026 Catcrafts®

// RTVolume — procedural (AABB) ray tracing on both backends. Demonstrates
// the two features this example was written to exercise:
//
//   * VK_GEOMETRY_TYPE_AABBS_KHR — a BLAS built from AABBs
//     (Mesh::BuildProcedural) whose surface is supplied by an intersection
//     shader (here an analytic ray–sphere test). The boxes are unit cubes
//     [-1,1]^3; the intersection shader turns each into a sphere.
//
//   * any-hit — the spheres are registered non-opaque, and an any-hit
//     shader punches a spherical checkerboard of holes by ignoring the
//     intersection for half the cells. Without any-hit the spheres are
//     solid; with it you can see the background (and other spheres)
//     through the cut-out cells.
//
// A 3×3×3 grid of these procedural spheres is shaded by surface normal +
// a fixed sun. The Vulkan path runs the same scene through hardware RT
// (PROCEDURAL_HIT_GROUP_KHR with GLSL intersection / any-hit shaders, a
// fixed camera in raygen.glsl); the WebGPU path is the software wavefront
// tracer with a host-driven free camera.

#ifndef CRAFTER_GRAPHICS_WINDOW_DOM
#include "vulkan/vulkan.h"

import Crafter.Graphics;
import Crafter.Math;
import std;

using namespace Crafter;

namespace {
    constexpr int   kGrid    = 3;
    constexpr float kSpacing = 3.0f;
}

int main() {
    const int instanceCount = kGrid * kGrid * kGrid;
    std::println("[RTVolume] grid {}^3 = {} procedural spheres", kGrid, instanceCount);

    Device::Initialize();
    Window window(1280, 720, "RTVolume");
    VkCommandBuffer cmd = window.StartInit();

    DescriptorHeapVulkan descriptorHeap;
    descriptorHeap.Initialize(/*images*/ 1, /*buffers*/ 1, /*samplers*/ 0);

    // Specialization constant: the TLAS slot offset (same pattern as the
    // Sponza example). Camera is fixed in raygen.glsl — no user buffers.
    VkSpecializationMapEntry raygenEntry = { .constantID = 0, .offset = 0, .size = sizeof(std::uint16_t) };
    VkSpecializationInfo raygenSpec = {
        .mapEntryCount = 1, .pMapEntries = &raygenEntry,
        .dataSize = sizeof(std::uint16_t), .pData = &descriptorHeap.bufferStartElement,
    };

    auto imgSlots = descriptorHeap.AllocateImageSlots(1);
    auto bufSlots = descriptorHeap.AllocateBufferSlots(1);

    // SBT order fixes the shader indices used by the groups below.
    std::array<VulkanShader, 5> shaders {{
        { "raygen.spv",       "main", VK_SHADER_STAGE_RAYGEN_BIT_KHR,       &raygenSpec },
        { "miss.spv",         "main", VK_SHADER_STAGE_MISS_BIT_KHR,         nullptr     },
        { "closesthit.spv",   "main", VK_SHADER_STAGE_CLOSEST_HIT_BIT_KHR,  nullptr     },
        { "anyhit.spv",       "main", VK_SHADER_STAGE_ANY_HIT_BIT_KHR,      nullptr     },
        { "intersection.spv", "main", VK_SHADER_STAGE_INTERSECTION_BIT_KHR, nullptr     },
    }};
    ShaderBindingTableVulkan shaderTable;
    shaderTable.Init(shaders);

    std::array<VkRayTracingShaderGroupCreateInfoKHR, 1> raygenGroups {{ {
        .sType = VK_STRUCTURE_TYPE_RAY_TRACING_SHADER_GROUP_CREATE_INFO_KHR,
        .type = VK_RAY_TRACING_SHADER_GROUP_TYPE_GENERAL_KHR,
        .generalShader = 0, .closestHitShader = VK_SHADER_UNUSED_KHR,
        .anyHitShader = VK_SHADER_UNUSED_KHR, .intersectionShader = VK_SHADER_UNUSED_KHR,
    } }};
    std::array<VkRayTracingShaderGroupCreateInfoKHR, 1> missGroups {{ {
        .sType = VK_STRUCTURE_TYPE_RAY_TRACING_SHADER_GROUP_CREATE_INFO_KHR,
        .type = VK_RAY_TRACING_SHADER_GROUP_TYPE_GENERAL_KHR,
        .generalShader = 1, .closestHitShader = VK_SHADER_UNUSED_KHR,
        .anyHitShader = VK_SHADER_UNUSED_KHR, .intersectionShader = VK_SHADER_UNUSED_KHR,
    } }};
    // One procedural hit group: closest-hit + any-hit + intersection.
    std::array<VkRayTracingShaderGroupCreateInfoKHR, 1> hitGroups {{ {
        .sType = VK_STRUCTURE_TYPE_RAY_TRACING_SHADER_GROUP_CREATE_INFO_KHR,
        .type = VK_RAY_TRACING_SHADER_GROUP_TYPE_PROCEDURAL_HIT_GROUP_KHR,
        .generalShader = VK_SHADER_UNUSED_KHR, .closestHitShader = 2,
        .anyHitShader = 3, .intersectionShader = 4,
    } }};

    PipelineRTVulkan pipeline;
    pipeline.Init(cmd, raygenGroups, missGroups, hitGroups, shaderTable);

    // ── One procedural unit-box BLAS, built from a DEVICE buffer the GPU
    //    writes (issue #37). A compute shader (aabbs.comp.glsl) writes the
    //    box [-r,r]^3 into a device-local buffer; that buffer is fed straight
    //    into Mesh::BuildProcedural by device address — no host copy of the
    //    AABBs. The intersection shader then treats the box as the bounding
    //    volume of a sphere centred at the object origin. opaque=false so the
    //    any-hit cut-out runs. ──────────────────────────────────────────────
    constexpr std::uint32_t kBoxCount = 1;
    VulkanBuffer<RTAabb, false> aabbDevice;
    aabbDevice.Resize(
        VK_BUFFER_USAGE_2_ACCELERATION_STRUCTURE_BUILD_INPUT_READ_ONLY_BIT_KHR
        | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT
        | VK_BUFFER_USAGE_STORAGE_BUFFER_BIT,
        VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, kBoxCount);

    // The producer reaches its output through a buffer_reference (device
    // address pushed below). It binds no heap descriptors, but a
    // descriptor-heap compute pipeline still expects a heap bound for its
    // push-data path — the render loop binds it every frame, but the init
    // command buffer hasn't, so bind it here before dispatching.
    {
        VkBindHeapInfoEXT resourceHeapInfo {
            .sType = VK_STRUCTURE_TYPE_BIND_HEAP_INFO_EXT,
            .heapRange = {
                .address = descriptorHeap.resourceHeap[window.currentBuffer].address,
                .size = static_cast<std::uint32_t>(descriptorHeap.resourceHeap[window.currentBuffer].size),
            },
            .reservedRangeOffset = (descriptorHeap.resourceHeap[window.currentBuffer].size - Device::descriptorHeapProperties.minResourceHeapReservedRange) & ~(Device::descriptorHeapProperties.imageDescriptorAlignment - 1),
            .reservedRangeSize = Device::descriptorHeapProperties.minResourceHeapReservedRange,
        };
        Device::vkCmdBindResourceHeapEXT(cmd, &resourceHeapInfo);
    }

    ComputeShader aabbWriter;
    aabbWriter.Load("aabbs.comp.spv");
    struct AabbPush { VkDeviceAddress boxes; std::uint32_t count; float time; } push {
        aabbDevice.address, kBoxCount, 0.0f };
    aabbWriter.Dispatch(cmd, &push, sizeof(push), (kBoxCount + 63) / 64);

    // Order the producing writes before the BLAS build reads the buffer.
    VkMemoryBarrier aabbBarrier {
        .sType = VK_STRUCTURE_TYPE_MEMORY_BARRIER,
        .srcAccessMask = VK_ACCESS_SHADER_WRITE_BIT,
        .dstAccessMask = VK_ACCESS_ACCELERATION_STRUCTURE_READ_BIT_KHR,
    };
    vkCmdPipelineBarrier(cmd,
        VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
        VK_PIPELINE_STAGE_ACCELERATION_STRUCTURE_BUILD_BIT_KHR,
        0, 1, &aabbBarrier, 0, nullptr, 0, nullptr);

    Mesh sphere;
    // allowUpdate so the GPU-written boxes can be refit in place each frame
    // (the per-frame companion to this build — see the headless
    // BLASBuildOptions test for the device-buffer refit path).
    sphere.BuildProcedural(aabbDevice.address, kBoxCount, /*opaque*/ false, cmd,
        RTBuildOptions{ .allowUpdate = true });

    // ── Instance grid. ─────────────────────────────────────────────────
    static std::vector<RenderingElement3D> renderers;
    renderers.reserve(static_cast<std::size_t>(instanceCount));
    const float origin0 = -0.5f * static_cast<float>(kGrid - 1) * kSpacing;
    for (int x = 0; x < kGrid; ++x)
    for (int y = 0; y < kGrid; ++y)
    for (int z = 0; z < kGrid; ++z) {
        renderers.emplace_back();
        RenderingElement3D& r = renderers.back();
        auto& tx = r.instance.transform.matrix;
        tx[0][0] = 1; tx[0][1] = 0; tx[0][2] = 0; tx[0][3] = origin0 + float(x) * kSpacing;
        tx[1][0] = 0; tx[1][1] = 1; tx[1][2] = 0; tx[1][3] = origin0 + float(y) * kSpacing;
        tx[2][0] = 0; tx[2][1] = 0; tx[2][2] = 1; tx[2][3] = origin0 + float(z) * kSpacing;
        r.instance.instanceCustomIndex                    = static_cast<std::uint32_t>(renderers.size() - 1);
        r.instance.mask                                   = 0xFF;
        r.instance.instanceShaderBindingTableRecordOffset = 0;
        // flags = 0: do NOT force opaque, so the any-hit shader runs.
        r.instance.flags                                  = 0;
        r.instance.accelerationStructureReference         = sphere.blasAddr;
        RenderingElement3D::Add(&r);
    }
    for (std::uint32_t f = 0; f < Window::numFrames; ++f) {
        RenderingElement3D::BuildTLAS(cmd, f);
    }

    window.FinishInit();

    // Write descriptors: TLAS at bufSlots[0], output image at imgSlots[0].
    // Per-frame replicated — same pattern as the Sponza example.
    VkDeviceAddressRangeKHR tlasRanges[Window::numFrames];
    VkImageDescriptorInfoEXT outImgInfos[Window::numFrames];
    for (std::uint32_t f = 0; f < Window::numFrames; ++f) {
        tlasRanges[f] = { .address = RenderingElement3D::tlases[f].address };
        outImgInfos[f] = {
            .sType = VK_STRUCTURE_TYPE_IMAGE_DESCRIPTOR_INFO_EXT,
            .pView = &window.imageViews[f],
            .layout = VK_IMAGE_LAYOUT_GENERAL,
        };
    }

    std::vector<VkResourceDescriptorInfoEXT> resources;
    std::vector<VkHostAddressRangeEXT>       destinations;
    resources.reserve(Window::numFrames * 2);
    destinations.reserve(Window::numFrames * 2);
    for (std::uint32_t f = 0; f < Window::numFrames; ++f) {
        resources.push_back({
            .sType = VK_STRUCTURE_TYPE_RESOURCE_DESCRIPTOR_INFO_EXT,
            .type = VK_DESCRIPTOR_TYPE_ACCELERATION_STRUCTURE_KHR,
            .data = { .pAddressRange = &tlasRanges[f] },
        });
        destinations.push_back({
            .address = descriptorHeap.resourceHeap[f].value
                     + descriptorHeap.BufferByteOffset(bufSlots.firstElement),
            .size = Device::descriptorHeapProperties.bufferDescriptorSize,
        });
        resources.push_back({
            .sType = VK_STRUCTURE_TYPE_RESOURCE_DESCRIPTOR_INFO_EXT,
            .type = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
            .data = { .pImage = &outImgInfos[f] },
        });
        destinations.push_back({
            .address = descriptorHeap.resourceHeap[f].value
                     + descriptorHeap.ImageByteOffset(imgSlots.firstElement),
            .size = Device::descriptorHeapProperties.imageDescriptorSize,
        });
    }
    Device::vkWriteResourceDescriptorsEXT(Device::device,
        static_cast<std::uint32_t>(resources.size()),
        resources.data(), destinations.data());
    for (std::uint32_t f = 0; f < Window::numFrames; ++f) {
        descriptorHeap.resourceHeap[f].FlushDevice();
    }

    window.descriptorHeap = &descriptorHeap;
    RTPass rtPass(&pipeline);
    window.passes.push_back(&rtPass);

    window.Render();
    window.StartSync();
    return 0;
}
#else

import Crafter.Graphics;
import Crafter.Math;
import Crafter.Event;
import std;

using namespace Crafter;
namespace fs = std::filesystem;

namespace {
    constexpr int   kGrid    = 3;
    constexpr float kSpacing = 3.0f;

    struct CameraGPU {
        float origin[3];   float pad0;
        float right[3];    float tanHalf;
        float up[3];       float aspect;
        float forward[3];  float pad1;
    };
    static_assert(sizeof(CameraGPU) == 64);
}

int main() {
    const int instanceCount = kGrid * kGrid * kGrid;
    std::println("[RTVolume] grid {}^3 = {} procedural spheres", kGrid, instanceCount);

    Device::Initialize();
    static Window window(1280, 720, "RTVolume");
    auto cmd = window.StartInit();

    DescriptorHeapWebGPU heap;
    heap.Initialize(/*images*/ 1, /*buffers*/ 2, /*samplers*/ 1);

    // SBT order fixes the shader indices used by the groups below.
    std::array<WebGPUShader, 6> shaders {{
        WebGPUShader(fs::path("raygen.wgsl"),       "raygen_main",       WebGPURTStage::Raygen),
        WebGPUShader(fs::path("miss.wgsl"),         "miss_main",         WebGPURTStage::Miss),
        WebGPUShader(fs::path("closesthit.wgsl"),   "closesthit_main",   WebGPURTStage::ClosestHit),
        WebGPUShader(fs::path("anyhit.wgsl"),       "anyhit_main",       WebGPURTStage::AnyHit),
        WebGPUShader(fs::path("intersection.wgsl"), "intersection_main", WebGPURTStage::Intersection),
        WebGPUShader(fs::path("resolve.wgsl"),      "resolve_main",      WebGPURTStage::Resolve),
    }};
    ShaderBindingTableWebGPU sbt;
    sbt.Init(shaders);

    std::array<RTShaderGroup, 1> raygenGroups {{ { .type = RTShaderGroupType::General, .generalShader = 0 } }};
    std::array<RTShaderGroup, 1> missGroups   {{ { .type = RTShaderGroupType::General, .generalShader = 1 } }};
    // One procedural hit group: closest-hit + any-hit + intersection.
    std::array<RTShaderGroup, 1> hitGroups    {{ {
        .type               = RTShaderGroupType::ProceduralHitGroup,
        .closestHitShader   = 2,
        .anyHitShader       = 3,
        .intersectionShader = 4,
    } }};

    std::array<UICustomBinding, 1> bindings {{
        { .group = 3, .binding = 0, .kind = UICustomBindingKind::Buffer, .pushOffset = 0 },
    }};

    PipelineRTWebGPU pipeline;
    pipeline.Init(cmd, raygenGroups, missGroups, hitGroups, sbt, bindings);

    // ── One procedural unit-box BLAS, built from a DEVICE buffer the GPU
    //    writes (issue #37). A compute shader (aabbs.comp.wgsl) writes the
    //    box [-r,r]^3 into a device storage buffer; that buffer is fed
    //    straight into Mesh::BuildProcedural by handle — no host copy of the
    //    AABBs — and RefitProcedural re-consumes it each frame so the boxes
    //    can breathe. The intersection shader treats the box as the bounding
    //    volume of a sphere; opaque=false so the any-hit cut-out runs. ──────
    constexpr std::uint32_t kBoxCount = 1;
    // worldBounds must enclose every box the producer can write (r ≤ 1.15).
    constexpr RTAabb kWorldBounds { .min = {-1.2f, -1.2f, -1.2f}, .max = {1.2f, 1.2f, 1.2f} };

    static WebGPUBuffer<RTAabb, false> aabbBuf;
    aabbBuf.Create(kBoxCount);

    // GPU producer: writes the boxes into aabbBuf. No rayQuery → the storage
    // buffer binding lives at group 1.
    struct AabbParams { std::uint32_t count; float time; std::uint32_t _0, _1; };
    std::array<UICustomBinding, 1> aabbBindings {{
        { .group = 1, .binding = 0, .kind = UICustomBindingKind::BufferReadWrite, .pushOffset = 0 },
    }};
    static PlainComputeShader aabbWriter;
    aabbWriter.Load(fs::path("aabbs.comp.wgsl"), sizeof(AabbParams), aabbBindings);
    static std::array<std::uint32_t, 1> aabbHandles { aabbBuf.handle };

    {
        AabbParams params { kBoxCount, 0.0f, 0, 0 };
        aabbWriter.Dispatch(&params, sizeof(params), aabbHandles, (kBoxCount + 63u) / 64u);
    }

    static Mesh sphere;
    sphere.BuildProcedural(aabbBuf.handle, kBoxCount, kWorldBounds, /*opaque*/ false, cmd);

    // ── Camera buffer + handle array. ─────────────────────────────────
    WebGPUBuffer<CameraGPU, true> cameraBuf;
    cameraBuf.Create(1);
    static std::array<std::uint32_t, 1> userHandles { cameraBuf.handle };

    // ── Instance grid. ────────────────────────────────────────────────
    static std::vector<RenderingElement3D> renderers;
    renderers.reserve(static_cast<std::size_t>(instanceCount));
    const float origin0 = -0.5f * static_cast<float>(kGrid - 1) * kSpacing;
    for (int x = 0; x < kGrid; ++x)
    for (int y = 0; y < kGrid; ++y)
    for (int z = 0; z < kGrid; ++z) {
        renderers.emplace_back();
        RenderingElement3D& r = renderers.back();
        auto& tx = r.instance.transform.matrix;
        tx[0][0] = 1; tx[0][1] = 0; tx[0][2] = 0; tx[0][3] = origin0 + float(x) * kSpacing;
        tx[1][0] = 0; tx[1][1] = 1; tx[1][2] = 0; tx[1][3] = origin0 + float(y) * kSpacing;
        tx[2][0] = 0; tx[2][1] = 0; tx[2][2] = 1; tx[2][3] = origin0 + float(z) * kSpacing;
        r.instance.instanceCustomIndex                    = static_cast<std::uint32_t>(renderers.size() - 1);
        r.instance.mask                                   = 0xFF;
        r.instance.instanceShaderBindingTableRecordOffset = 0;
        // flags = 0: do NOT force opaque, so the any-hit shader runs.
        r.instance.flags                                  = 0;
        r.instance.accelerationStructureReference         = sphere.blasAddr;
        RenderingElement3D::Add(&r);
    }
    RenderingElement3D::BuildTLAS(cmd, 0);

    window.descriptorHeap = &heap;
    window.FinishInit();

    RTPass rtPass(&pipeline);
    rtPass.handlesPtr   = userHandles.data();
    rtPass.handlesCount = static_cast<std::uint32_t>(userHandles.size());
    rtPass.maxDepth     = 1;   // primary only
    window.passes.push_back(&rtPass);

    // ── Free camera framing the grid. ─────────────────────────────────
    const float ext = float(kGrid - 1) * kSpacing;
    struct CamState {
        Vector<float, 3, 4> position;
        float yaw;
        float pitch;
    } cam {
        Vector<float, 3, 4>{ ext * 1.1f, ext * 0.8f, ext * 1.6f },
        0.0f, 0.0f,
    };
    {
        Vector<float, 3, 4> d { -cam.position.x, -cam.position.y, -cam.position.z };
        const float len = std::sqrt(d.x*d.x + d.y*d.y + d.z*d.z);
        cam.yaw   = std::atan2(d.z, d.x);
        cam.pitch = std::asin(d.y / len);
    }

    Input::Map inputMap;
    Input::Action& moveAct = inputMap.AddAction("Move", Input::ActionType::Vector2);
    Input::Action& lookAct = inputMap.AddAction("Look", Input::ActionType::Vector2);
    moveAct.bindings = { Input::WASDBind{
        Key(CrafterKeys::W), Key(CrafterKeys::S), Key(CrafterKeys::A), Key(CrafterKeys::D) } };
    lookAct.bindings = { Input::MouseDeltaBind{ 1.0f } };
    inputMap.Attach(window);

    const float kMoveSpeed = ext * 0.8f + 1.0f;
    const float kLookSens  = 0.05f;
    const float kDt        = 1.0f / 60.0f;

    EventListener<void> camTick(&window.onBeforeUpdate, [&]() {
        inputMap.Tick();
        cam.yaw   += lookAct.vector2.x * kLookSens;
        cam.pitch -= lookAct.vector2.y * kLookSens;
        cam.pitch  = std::clamp(cam.pitch, -1.55f, 1.55f);

        const float cp = std::cos(cam.pitch), sp = std::sin(cam.pitch);
        const float cy = std::cos(cam.yaw),   sy = std::sin(cam.yaw);
        Vector<float, 3, 4> forward { cp * cy, sp, cp * sy };
        Vector<float, 3, 4> worldUp { 0.0f, 1.0f, 0.0f };
        Vector<float, 3, 4> right { forward.y*worldUp.z - forward.z*worldUp.y,
                                    forward.z*worldUp.x - forward.x*worldUp.z,
                                    forward.x*worldUp.y - forward.y*worldUp.x };
        const float rLen = std::sqrt(right.x*right.x + right.y*right.y + right.z*right.z);
        right.x /= rLen; right.y /= rLen; right.z /= rLen;
        Vector<float, 3, 4> up { right.y*forward.z - right.z*forward.y,
                                 right.z*forward.x - right.x*forward.z,
                                 right.x*forward.y - right.y*forward.x };

        const float dx = moveAct.vector2.x * kMoveSpeed * kDt;
        const float dy = moveAct.vector2.y * kMoveSpeed * kDt;
        cam.position.x += right.x*dx + forward.x*dy;
        cam.position.y += right.y*dx + forward.y*dy;
        cam.position.z += right.z*dx + forward.z*dy;

        CameraGPU& g = cameraBuf.value[0];
        g.origin[0]=cam.position.x; g.origin[1]=cam.position.y; g.origin[2]=cam.position.z; g.pad0=0;
        g.right[0]=right.x;   g.right[1]=right.y;   g.right[2]=right.z;
        g.up[0]=up.x;         g.up[1]=up.y;         g.up[2]=up.z;
        g.forward[0]=forward.x; g.forward[1]=forward.y; g.forward[2]=forward.z;
        g.aspect  = float(window.width) / float(window.height);
        g.tanHalf = std::tan(70.0f * 3.14159265f / 360.0f);
        g.pad1    = 0;
        cameraBuf.FlushDevice();
    });

    // ── Per-frame procedural refit from the device buffer (issue #37). The
    //    GPU producer rewrites the breathing box into aabbBuf each frame and
    //    RefitProcedural re-consumes it by handle — no host copy, and the
    //    BLAS handle (blasAddr) stays stable so the TLAS instances built
    //    above keep referencing it. worldBounds is constant so the TLAS need
    //    not rebuild. ───────────────────────────────────────────────────────
    EventListener<void> aabbTick(&window.onBeforeUpdate, [&]() {
        static float t = 0.0f;
        t += kDt;
        AabbParams params { kBoxCount, t, 0, 0 };
        aabbWriter.Dispatch(&params, sizeof(params), aabbHandles, (kBoxCount + 63u) / 64u);
        sphere.RefitProcedural(aabbBuf.handle, kBoxCount, kWorldBounds);
    });

    window.Render();
    window.StartUpdate();
    window.StartSync();
    return 0;
}
#endif