/*
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");
    // allowUpdate=true keeps the scratch buffer alive so the in-place UPDATE
    // refit below can reuse it (build scratch ≥ update scratch, no resize).
    Check(tri.scratchBuffer.buffer != VK_NULL_HANDLE,
          "allowUpdate=true → scratch retained for refit (issue #66)");
    // Geometry is placed in device-local memory regardless of the upload
    // strategy taken (issue #73): the direct ReBAR/UMA path lands a
    // HOST_VISIBLE | DEVICE_LOCAL type, the staged path a pure DEVICE_LOCAL
    // one — either way the DEVICE_LOCAL bit is set, so the BLAS build and any
    // hit-shader fetch read from VRAM, not system RAM.
    Check((tri.vertexBuffer.memoryPropertyFlagsChosen & VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT) != 0,
          "triangle vertexBuffer placed in device-local memory (#73)");
    Check((tri.indexBuffer.memoryPropertyFlagsChosen & VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT) != 0,
          "triangle indexBuffer placed in device-local memory (#73)");

    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");
    // A static mesh can never refit, so its scratch is released right after
    // the build completes rather than kept as persistent VRAM (issue #66).
    // DeferredClear nulls the handle immediately (the allocation itself is
    // retired by the fence-keyed queue once the build's frame passes).
    Check(triFast.scratchBuffer.buffer == VK_NULL_HANDLE,
          "allowUpdate=false → scratch released after build (issue #66)");

    // 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");
    // The fallback rebuild re-Created the scratch (Resize saw the nulled
    // handle) and, still being a static mesh, released it again afterwards.
    Check(triFast.scratchBuffer.buffer == VK_NULL_HANDLE,
          "static-mesh refit rebuild re-released its scratch (issue #66)");

    // ── 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");
    Check(proc.scratchBuffer.buffer != VK_NULL_HANDLE,
          "procedural allowUpdate=true → scratch retained for refit (issue #66)");
    Check((proc.aabbBuffer.memoryPropertyFlagsChosen & VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT) != 0,
          "procedural aabbBuffer placed in device-local memory (#73)");

    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");

    // ── Force the STAGED upload path (issue #73). On this ReBAR dev host the
    //    direct map path is normally taken; temporarily zero the cached upload
    //    budget so PreferDirectDeviceWrite returns false, exercising Mesh's
    //    pure-DEVICE_LOCAL + transient-staging-buffer + vkCmdCopyBuffer path
    //    (and its deferred-deletion of the staging buffer). The validation
    //    layer check below is the real assertion that the staged copy and its
    //    barriers are spec-correct on a fresh build and an in-place refit. ───
    {
        const VkDeviceSize savedBudget = Device::directWriteBudget;
        Device::directWriteBudget = 0;   // 0 → PreferDirectDeviceWrite is always false

        Mesh staged;
        auto verts = CubeVerts(1.0f);
        auto idx   = CubeIndices();
        {
            VkCommandBuffer cmd = BeginCmd();
            staged.Build(verts, idx, cmd, RTBuildOptions{
                .preference = RTBuildPreference::FastTrace, .allowUpdate = true });
            SubmitWait(cmd);
        }
        Check(staged.blasAddr != 0, "staged-upload Build produced a non-zero blasAddr (#73)");
        // Staged geometry lives in pure DEVICE_LOCAL — no HOST_VISIBLE bit.
        Check((staged.vertexBuffer.memoryPropertyFlagsChosen & VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT) == 0,
              "staged vertexBuffer is device-local-only, not host-visible (#73)");

        const VkAccelerationStructureKHR stagedHandle = staged.accelerationStructure;
        {
            // Same topology, deformed positions → in-place UPDATE, re-staged.
            auto verts2 = CubeVerts(1.5f);
            VkCommandBuffer cmd = BeginCmd();
            staged.Refit(verts2, idx, cmd);
            SubmitWait(cmd);
        }
        Check(staged.accelerationStructure == stagedHandle,
              "staged-upload Refit kept the same AS handle (in-place UPDATE, #73)");

        Device::directWriteBudget = savedBudget;
    }

    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;
}