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

// Issue #36: BLAS build options — fast-build / fast-trace preference and
// in-place refit (UPDATE-mode rebuild). This exercises the real hardware
// path: it spins up a headless Vulkan device (no swapchain / window needed
// — a BLAS build only touches the queue + command pool), records BLAS
// builds and refits into one-time command buffers, and submits them.
//
// What is asserted:
//   - Build() with RTBuildOptions records the requested preference + the
//     ALLOW_UPDATE bit, and produces a non-zero device address.
//   - Refit() on an allowUpdate BLAS keeps the SAME acceleration-structure
//     handle and blasAddr (proof it took the in-place UPDATE path, so TLAS
//     instances referencing it stay valid).
//   - Refit() without allowUpdate, or after a topology change, falls back to
//     a fresh build (new handle / address are fine; it must still succeed).
//   - The procedural (AABB) path supports the same options + refit.
//   - The Vulkan validation layer reports ZERO errors across all of the
//     above (Device::validationErrorCount) — the strongest check that the
//     UPDATE builds are spec-correct (scratch sizing, ALLOW_UPDATE present,
//     src==dst, matching topology, …).
//
// Validation layers are required for the last check to be meaningful; the
// build marks this test as needing the SDK layers.

#include "vulkan/vulkan.h"
#include <cstdlib>

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

using namespace Crafter;

namespace {

int failures = 0;

void Check(bool ok, std::string_view what) {
    std::println("{}  {}", ok ? "PASS" : "FAIL", what);
    if (!ok) ++failures;
}

// One-time command buffer helpers — record a BLAS build, submit, block.
VkCommandBuffer BeginCmd() {
    VkCommandBufferAllocateInfo allocInfo {
        .sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO,
        .commandPool = Device::commandPool,
        .level = VK_COMMAND_BUFFER_LEVEL_PRIMARY,
        .commandBufferCount = 1,
    };
    VkCommandBuffer cmd = VK_NULL_HANDLE;
    Device::CheckVkResult(vkAllocateCommandBuffers(Device::device, &allocInfo, &cmd));
    VkCommandBufferBeginInfo beginInfo {
        .sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO,
        .flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT,
    };
    Device::CheckVkResult(vkBeginCommandBuffer(cmd, &beginInfo));
    return cmd;
}

void SubmitWait(VkCommandBuffer cmd) {
    Device::CheckVkResult(vkEndCommandBuffer(cmd));
    VkSubmitInfo submitInfo {
        .sType = VK_STRUCTURE_TYPE_SUBMIT_INFO,
        .commandBufferCount = 1,
        .pCommandBuffers = &cmd,
    };
    Device::CheckVkResult(vkQueueSubmit(Device::queue, 1, &submitInfo, VK_NULL_HANDLE));
    Device::CheckVkResult(vkQueueWaitIdle(Device::queue));
    vkFreeCommandBuffers(Device::device, Device::commandPool, 1, &cmd);
}

// A unit cube (8 verts, 12 triangles) — enough topology for a real BLAS.
std::vector<Vector<float, 3, 3>> CubeVerts(float s) {
    return {
        {-s,-s,-s}, { s,-s,-s}, { s, s,-s}, {-s, s,-s},
        {-s,-s, s}, { s,-s, s}, { s, s, s}, {-s, s, s},
    };
}
std::vector<std::uint32_t> CubeIndices() {
    return {
        0,1,2, 0,2,3,  4,6,5, 4,7,6,
        0,4,5, 0,5,1,  3,2,6, 3,6,7,
        1,5,6, 1,6,2,  0,3,7, 0,7,4,
    };
}

} // namespace

int main() {
    Device::Initialize();
    Device::validationErrorCount = 0;

    // ── Triangle BLAS: fast-trace + allow-update, then in-place refit. ──
    Mesh tri;
    {
        auto verts = CubeVerts(1.0f);
        auto idx   = CubeIndices();
        VkCommandBuffer cmd = BeginCmd();
        tri.Build(verts, idx, cmd, RTBuildOptions{
            .preference  = RTBuildPreference::FastTrace,
            .allowUpdate = true,
        });
        SubmitWait(cmd);
    }
    Check(tri.blasAddr != 0, "triangle Build produced a non-zero blasAddr");
    Check(tri.accelerationStructure != VK_NULL_HANDLE, "triangle Build created an AS handle");
    Check((tri.buildFlags & VK_BUILD_ACCELERATION_STRUCTURE_PREFER_FAST_TRACE_BIT_KHR) != 0,
          "FastTrace preference → PREFER_FAST_TRACE bit set");
    Check((tri.buildFlags & VK_BUILD_ACCELERATION_STRUCTURE_ALLOW_UPDATE_BIT_KHR) != 0,
          "allowUpdate=true → ALLOW_UPDATE bit set");
    Check(tri.builtPrimitiveCount == 12, "triangle BLAS reports 12 primitives");

    const VkDeviceAddress             triAddrBefore = tri.blasAddr;
    const VkAccelerationStructureKHR  triHandleBefore = tri.accelerationStructure;
    {
        // Same topology, deformed positions → must take the UPDATE path.
        auto verts = CubeVerts(1.5f);
        auto idx   = CubeIndices();
        VkCommandBuffer cmd = BeginCmd();
        tri.Refit(verts, idx, cmd);
        SubmitWait(cmd);
    }
    Check(tri.accelerationStructure == triHandleBefore,
          "Refit kept the same AS handle (in-place UPDATE)");
    Check(tri.blasAddr == triAddrBefore,
          "Refit kept the same blasAddr (instances stay valid)");

    // ── Triangle BLAS: fast-build, no update → flags reflect the choice. ─
    Mesh triFast;
    {
        auto verts = CubeVerts(1.0f);
        auto idx   = CubeIndices();
        VkCommandBuffer cmd = BeginCmd();
        triFast.Build(verts, idx, cmd, RTBuildOptions{ .preference = RTBuildPreference::FastBuild });
        SubmitWait(cmd);
    }
    Check((triFast.buildFlags & VK_BUILD_ACCELERATION_STRUCTURE_PREFER_FAST_BUILD_BIT_KHR) != 0,
          "FastBuild preference → PREFER_FAST_BUILD bit set");
    Check((triFast.buildFlags & VK_BUILD_ACCELERATION_STRUCTURE_ALLOW_UPDATE_BIT_KHR) == 0,
          "allowUpdate=false → ALLOW_UPDATE bit clear");

    // Refit without allowUpdate must still succeed via the rebuild fallback.
    {
        auto verts = CubeVerts(0.5f);
        auto idx   = CubeIndices();
        VkCommandBuffer cmd = BeginCmd();
        triFast.Refit(verts, idx, cmd);
        SubmitWait(cmd);
    }
    Check(triFast.blasAddr != 0, "Refit fallback rebuild still produced a valid BLAS");

    // ── Procedural (AABB) BLAS: build with options, then refit. ─────────
    Mesh proc;
    {
        std::array<RTAabb, 2> boxes {{
            { .min = {-1,-1,-1}, .max = {1,1,1} },
            { .min = { 2, 2, 2}, .max = {3,3,3} },
        }};
        VkCommandBuffer cmd = BeginCmd();
        proc.BuildProcedural(boxes, /*opaque*/ false, cmd, RTBuildOptions{
            .preference = RTBuildPreference::FastTrace, .allowUpdate = true });
        SubmitWait(cmd);
    }
    Check(proc.blasAddr != 0, "procedural Build produced a non-zero blasAddr");
    Check(proc.builtPrimitiveCount == 2, "procedural BLAS reports 2 primitives");

    const VkDeviceAddress            procAddrBefore   = proc.blasAddr;
    const VkAccelerationStructureKHR procHandleBefore = proc.accelerationStructure;
    {
        std::array<RTAabb, 2> boxes {{
            { .min = {-2,-2,-2}, .max = {2,2,2} },
            { .min = { 4, 4, 4}, .max = {5,5,5} },
        }};
        VkCommandBuffer cmd = BeginCmd();
        proc.RefitProcedural(boxes, cmd);
        SubmitWait(cmd);
    }
    Check(proc.accelerationStructure == procHandleBefore && proc.blasAddr == procAddrBefore,
          "RefitProcedural kept the same AS handle + blasAddr (in-place UPDATE)");

    // ── Procedural (AABB) BLAS from a DEVICE buffer (issue #37). The boxes
    //    live in a device-local buffer a GPU pass would write — here filled
    //    via a staging copy + barrier, standing in for a compute dispatch —
    //    and fed straight into BuildProcedural/RefitProcedural by device
    //    address, with no host memcpy through Mesh::aabbBuffer. ────────────
    constexpr std::uint32_t kDevCount = 3;
    // Device-local AABB build input: the producer's buffer. Needs the AS
    // build-input + device-address usage bits, plus TRANSFER_DST so the
    // staging upload (our stand-in GPU writer) can fill it.
    VulkanBuffer<RTAabb, false> devAabb;
    devAabb.Resize(
        VK_BUFFER_USAGE_2_ACCELERATION_STRUCTURE_BUILD_INPUT_READ_ONLY_BIT_KHR
        | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT
        | VK_BUFFER_USAGE_TRANSFER_DST_BIT,
        VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, kDevCount);
    VulkanBuffer<RTAabb, true> devStaging;
    // SHADER_DEVICE_ADDRESS because VulkanBuffer::Create always queries the
    // buffer device address (the staging buffer's address itself is unused).
    devStaging.Resize(VK_BUFFER_USAGE_TRANSFER_SRC_BIT | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT,
                      VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT, kDevCount);

    auto uploadBoxes = [&](VkCommandBuffer cmd, std::span<const RTAabb> boxes) {
        std::memcpy(devStaging.value, boxes.data(), boxes.size() * sizeof(RTAabb));
        devStaging.FlushDevice();
        VkBufferCopy region { .srcOffset = 0, .dstOffset = 0, .size = devAabb.size };
        vkCmdCopyBuffer(cmd, devStaging.buffer, devAabb.buffer, 1, &region);
        // Order the producing write before the AS build reads the buffer —
        // exactly the barrier a real compute producer would need.
        VkBufferMemoryBarrier barrier {
            .sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER,
            .srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT,
            .dstAccessMask = VK_ACCESS_ACCELERATION_STRUCTURE_READ_BIT_KHR,
            .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
            .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
            .buffer = devAabb.buffer, .offset = 0, .size = VK_WHOLE_SIZE,
        };
        vkCmdPipelineBarrier(cmd, VK_PIPELINE_STAGE_TRANSFER_BIT,
            VK_PIPELINE_STAGE_ACCELERATION_STRUCTURE_BUILD_BIT_KHR,
            0, 0, nullptr, 1, &barrier, 0, nullptr);
    };

    Mesh devProc;
    {
        std::array<RTAabb, kDevCount> boxes {{
            { .min = {-1,-1,-1}, .max = {1,1,1} },
            { .min = { 2, 2, 2}, .max = {3,3,3} },
            { .min = {-3,-3,-3}, .max = {-2,-2,-2} },
        }};
        VkCommandBuffer cmd = BeginCmd();
        uploadBoxes(cmd, boxes);
        devProc.BuildProcedural(devAabb.address, kDevCount, /*opaque*/ false, cmd,
            RTBuildOptions{ .preference = RTBuildPreference::FastTrace, .allowUpdate = true });
        SubmitWait(cmd);
    }
    Check(devProc.blasAddr != 0, "device-buffer BuildProcedural produced a non-zero blasAddr");
    Check(devProc.builtPrimitiveCount == kDevCount, "device-buffer BLAS reports 3 primitives");
    // The host-side aabbBuffer must be untouched — proof there was no memcpy.
    Check(devProc.aabbBuffer.buffer == VK_NULL_HANDLE,
          "device-buffer path never allocated the host aabbBuffer (zero-copy)");

    const VkDeviceAddress            devAddrBefore   = devProc.blasAddr;
    const VkAccelerationStructureKHR devHandleBefore = devProc.accelerationStructure;
    {
        // Same count, moved boxes (re-written into the same device buffer) →
        // in-place UPDATE straight from the device address.
        std::array<RTAabb, kDevCount> boxes {{
            { .min = {-2,-2,-2}, .max = {2,2,2} },
            { .min = { 4, 4, 4}, .max = {5,5,5} },
            { .min = {-5,-5,-5}, .max = {-3,-3,-3} },
        }};
        VkCommandBuffer cmd = BeginCmd();
        uploadBoxes(cmd, boxes);
        devProc.RefitProcedural(devAabb.address, kDevCount, cmd);
        SubmitWait(cmd);
    }
    Check(devProc.accelerationStructure == devHandleBefore && devProc.blasAddr == devAddrBefore,
          "device-buffer RefitProcedural kept the same AS handle + blasAddr (in-place UPDATE)");
    Check(devProc.aabbBuffer.buffer == VK_NULL_HANDLE,
          "device-buffer refit still never touched the host aabbBuffer");

    Check(Device::validationErrorCount == 0,
          std::format("no Vulkan validation errors ({} seen)", Device::validationErrorCount));

    if (failures != 0) {
        std::println("{} check(s) failed", failures);
        return EXIT_FAILURE;
    }
    std::println("all checks passed");
    return EXIT_SUCCESS;
}
