//SPDX-License-Identifier: LGPL-3.0-only
//SPDX-FileCopyrightText: Copyright (C) 2026 Catcrafts®

// A Mesh hands its acceleration structure to the deferred-deletion queue
// (issue #101) whenever it lets go of it, because frames in flight may still
// trace it:
//  - Building a Mesh that already has a BLAS used to destroy the old handle at
//    once, although the command buffer being recorded or a frame in flight
//    could still use it, and then built into the old storage with no barrier.
//    The old handle is now queued, and a barrier orders earlier reads of the
//    storage and writes of the scratch before the build. The storage is still
//    reused when the new BLAS fits, which keeps blasAddr stable; storage the
//    new BLAS outgrows is queued behind the handle. Two builds of one Mesh in
//    one command buffer is the case the validation layer sees: the command
//    buffer became invalid because its first build's destination was
//    destroyed, and synchronization validation
//    (VK_KHRONOS_VALIDATION_VALIDATE_SYNC=1) sees the two builds' writes race.
//  - Destroying a Mesh used to leak the handle. The handle and the Mesh's
//    buffers are now queued, the handle first so it goes before its storage.
//  - Moving a Mesh hands the handle over, so only the last owner queues it,
//    including when a std::vector<Mesh> grows.
//
// Needs a Vulkan ray-tracing device (exits 77 without one) and no window: the
// test steps Device::frameCounter itself. The validation checks only mean
// something with CRAFTER_GRAPHICS_VALIDATION set.

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

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

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

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

// `count` unit cubes in a row: enough geometry to outgrow one cube's BLAS.
void Cubes(int count, std::vector<Vector<float, 3, 3>>& verts, std::vector<std::uint32_t>& indices) {
    const std::vector<Vector<float, 3, 3>> cube = CubeVerts(0.5f);
    const std::vector<std::uint32_t> cubeIndices = CubeIndices();
    verts.clear();
    indices.clear();
    for (int i = 0; i < count; ++i) {
        const std::uint32_t base = static_cast<std::uint32_t>(verts.size());
        for (const Vector<float, 3, 3>& v : cube) verts.push_back({v.x + 2.0f * static_cast<float>(i), v.y, v.z});
        for (std::uint32_t index : cubeIndices) indices.push_back(base + index);
    }
}

// Where an entry sits in the queue, or the queue size when it isn't there.
std::size_t QueuePosition(auto matches) {
    const auto it = std::ranges::find_if(Device::deletionQueue, matches);
    return static_cast<std::size_t>(it - Device::deletionQueue.begin());
}

std::size_t QueuedCount(VkAccelerationStructureKHR handle) {
    return static_cast<std::size_t>(std::ranges::count_if(Device::deletionQueue,
        [&](const Device::PendingDeletion& e) { return e.accelerationStructure == handle; }));
}

bool Queued(VkBuffer buffer) {
    return std::ranges::any_of(Device::deletionQueue,
        [&](const Device::PendingDeletion& e) { return e.buffer == buffer; });
}

} // namespace

int main() {
    Device::Initialize();
    if (!Device::rayTracingSupported) {
        std::println("SKIP  no Vulkan ray-tracing device");
        return 77;
    }
    Device::validationErrorCount = 0;
    Device::framesInFlight = Window::numFrames;
    Device::frameCounter = 0;
    Device::DrainDeletions();

    std::vector<Vector<float, 3, 3>> verts = CubeVerts(1.0f);
    std::vector<std::uint32_t> indices = CubeIndices();
    const RTBuildOptions refittable { .allowUpdate = true };

    // ── A rebuild queues the old handle and reuses storage that fits ─────────
    {
        Mesh mesh;
        VkCommandBuffer cmd = BeginCmd();
        mesh.Build(verts, indices, cmd, refittable);
        const VkAccelerationStructureKHR first = mesh.accelerationStructure;
        const VkBuffer firstStorage = mesh.blasBuffer.buffer;

        // Same geometry, so the new BLAS fits the old storage.
        mesh.Build(verts, indices, cmd, refittable);
        Check(QueuedCount(first) == 1, "a rebuild queues the old acceleration structure instead of destroying it while recorded");
        Check(mesh.accelerationStructure != VK_NULL_HANDLE && mesh.accelerationStructure != first,
              "and creates a new one");
        Check(mesh.blasBuffer.buffer == firstStorage && !Queued(firstStorage), "in the same storage, which still fits");
        SubmitWait(cmd);
        if (std::getenv("CRAFTER_GRAPHICS_VALIDATION") != nullptr) {
            Check(Device::validationErrorCount == 0,
                  std::format("two builds of one Mesh in one command buffer are validation-clean ({} errors)", Device::validationErrorCount));
        }

        // A BLAS that outgrows the storage moves to new storage; the old
        // storage is queued behind the handle that lives in it.
        std::vector<Vector<float, 3, 3>> manyVerts;
        std::vector<std::uint32_t> manyIndices;
        Cubes(64, manyVerts, manyIndices);
        const VkAccelerationStructureKHR second = mesh.accelerationStructure;
        VkCommandBuffer growCmd = BeginCmd();
        mesh.Build(manyVerts, manyIndices, growCmd, refittable);
        const std::size_t handleAt = QueuePosition([&](const Device::PendingDeletion& e) { return e.accelerationStructure == second; });
        const std::size_t storageAt = QueuePosition([&](const Device::PendingDeletion& e) { return e.buffer == firstStorage; });
        Check(mesh.blasBuffer.buffer != firstStorage && storageAt < Device::deletionQueue.size(),
              "a rebuild that outgrows the storage queues it and gets new storage");
        Check(handleAt < storageAt, "the old handle is queued before the storage it lives in");
        SubmitWait(growCmd);

        // A refit after the rebuilds still updates the latest BLAS in place.
        VkCommandBuffer refitCmd = BeginCmd();
        const VkAccelerationStructureKHR rebuilt = mesh.accelerationStructure;
        manyVerts[0] = {-1.5f, -1.5f, -1.5f};
        mesh.Refit(manyVerts, manyIndices, refitCmd);
        SubmitWait(refitCmd);
        Check(mesh.accelerationStructure == rebuilt, "a refit after the rebuilds keeps the handle");
    }
    Device::DrainDeletions();

    // ── Destroying a Mesh queues its BLAS ahead of its buffers ───────────────
    {
        VkAccelerationStructureKHR handle = VK_NULL_HANDLE;
        VkBuffer storage = VK_NULL_HANDLE;
        VkBuffer vertices = VK_NULL_HANDLE;
        VkBuffer scratch = VK_NULL_HANDLE;
        {
            Mesh mesh;
            VkCommandBuffer cmd = BeginCmd();
            mesh.Build(verts, indices, cmd, refittable);
            SubmitWait(cmd);
            Device::DrainDeletions();
            handle = mesh.accelerationStructure;
            storage = mesh.blasBuffer.buffer;
            vertices = mesh.vertexBuffer.buffer;
            scratch = mesh.scratchBuffer.buffer;
        }
        Check(QueuedCount(handle) == 1, "destroying a Mesh queues its acceleration structure");
        Check(Queued(storage) && Queued(vertices) && Queued(scratch),
              "and its BLAS storage, vertex and scratch buffers");
        Check(!Device::deletionQueue.empty() && Device::deletionQueue.front().accelerationStructure == handle,
              "the acceleration structure is queued before the storage it lives in");

        Device::frameCounter += Window::numFrames;
        Device::ReclaimDeletions();
        Check(Device::deletionQueue.empty(), "everything a destroyed Mesh queued is freed on its retire frame");
    }

    // ── A never-built Mesh has nothing to queue ──────────────────────────────
    {
        { Mesh empty; }
        Check(Device::deletionQueue.empty(), "destroying a Mesh that was never built queues nothing");
    }

    // ── Moving hands the BLAS over; only the last owner queues it ────────────
    {
        VkAccelerationStructureKHR handle = VK_NULL_HANDLE;
        {
            Mesh source;
            VkCommandBuffer cmd = BeginCmd();
            source.Build(verts, indices, cmd);
            SubmitWait(cmd);
            Device::DrainDeletions();
            handle = source.accelerationStructure;
            const VkDeviceAddress address = source.blasAddr;

            Mesh moved(std::move(source));
            Check(source.accelerationStructure == VK_NULL_HANDLE && source.blasAddr == 0,
                  "a moved-from Mesh no longer holds the BLAS");
            Check(moved.accelerationStructure == handle && moved.blasAddr == address,
                  "the Mesh it moved to does");
        }
        Check(QueuedCount(handle) == 1, "the BLAS is queued once, by the Mesh it moved to");
        Device::DrainDeletions();
    }

    // ── Move assignment hands the BLAS over and queues the one it replaces ───
    {
        VkAccelerationStructureKHR kept = VK_NULL_HANDLE;
        VkAccelerationStructureKHR replaced = VK_NULL_HANDLE;
        {
            Mesh source;
            Mesh target;
            VkCommandBuffer cmd = BeginCmd();
            source.Build(verts, indices, cmd);
            target.Build(verts, indices, cmd);
            SubmitWait(cmd);
            Device::DrainDeletions();
            kept = source.accelerationStructure;
            replaced = target.accelerationStructure;

            target = std::move(source);
            Check(source.accelerationStructure == VK_NULL_HANDLE && target.accelerationStructure == kept,
                  "move assignment hands the BLAS over");
            Check(QueuedCount(replaced) == 1 && QueuedCount(kept) == 0, "and queues the BLAS it replaces, once");
        }
        Check(QueuedCount(kept) == 1, "the BLAS moved in is queued once, when its new Mesh goes");
        Device::DrainDeletions();

        // Erasing from a vector moves the later Meshes down over it.
        std::vector<Mesh> meshes(3);
        VkCommandBuffer cmd = BeginCmd();
        for (Mesh& mesh : meshes) mesh.Build(verts, indices, cmd);
        SubmitWait(cmd);
        const VkAccelerationStructureKHR erased = meshes[0].accelerationStructure;
        const VkAccelerationStructureKHR last = meshes[2].accelerationStructure;
        meshes.erase(meshes.begin());
        Check(QueuedCount(erased) == 1 && QueuedCount(last) == 0 && meshes[1].accelerationStructure == last,
              "erasing a Mesh from a vector queues its BLAS and keeps the others");
        meshes.clear();
        Device::DrainDeletions();
    }

    // ── A growing std::vector<Mesh> moves its elements without queueing ──────
    {
        std::vector<VkAccelerationStructureKHR> handles;
        {
            std::vector<Mesh> meshes;
            VkCommandBuffer cmd = BeginCmd();
            for (int i = 0; i < 5; ++i) {
                meshes.emplace_back().Build(verts, indices, cmd);
            }
            SubmitWait(cmd);
            for (const Mesh& mesh : meshes) handles.push_back(mesh.accelerationStructure);
            Check(std::ranges::none_of(handles, [](VkAccelerationStructureKHR h) { return QueuedCount(h) != 0; }),
                  "moving the meshes as the vector grew queued none of their BLASes");
            Device::DrainDeletions();
        }
        Check(std::ranges::all_of(handles, [](VkAccelerationStructureKHR h) { return QueuedCount(h) == 1; }),
              "each BLAS in the vector is queued exactly once when the vector goes");
        Device::DrainDeletions();
    }

    if (std::getenv("CRAFTER_GRAPHICS_VALIDATION") != nullptr) {
        Check(Device::validationErrorCount == 0,
              std::format("no Vulkan validation errors across the run ({} seen)", Device::validationErrorCount));
    }
    std::println("{}", failures == 0 ? "all checks passed" : "check(s) failed");
    return failures == 0 ? EXIT_SUCCESS : EXIT_FAILURE;
}