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

// A TlasWithBuffer hands its acceleration structure to the deferred-deletion
// queue (issue #101) when it lets go of it, the way a Mesh does its BLAS
// (MeshLifetime). Its special members used to be implicit, so only
// BuildTLAS's replacement of the handle went through the queue:
//  - Destroying one dropped the handle while its buffers queued themselves,
//    leaking the acceleration structure over freed storage. The handle is now
//    queued too, ahead of the storage it lives in.
//  - Assigning over one, the way to reset a scene's TLAS, dropped the handle it
//    held the same way. It is now queued the same way.
//  - Moving one copied the handle, leaving two owners, and left the moved-from
//    TlasWithBuffer with its count but none of its buffers, so its next
//    same-count BuildTLAS refit through buffers it no longer had. Moving now
//    hands the handle over, and the moved-from one builds a new TLAS.
//
// Needs a Vulkan ray-tracing device (exits 77 without one) and no window: the
// test steps Device::frameCounter itself. The validation check only means
// 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);
}

// Build the TLASes of the given frames for the registered elements.
void Build(std::initializer_list<std::uint32_t> frames) {
    VkCommandBuffer cmd = BeginCmd();
    for (std::uint32_t frame : frames) RenderingElement3D::BuildTLAS(cmd, frame);
    SubmitWait(cmd);
}

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

// The handle is queued exactly once, ahead of the storage it lives in.
bool QueuedAhead(VkAccelerationStructureKHR handle, VkBuffer storage) {
    const std::size_t handleAt = QueuePosition([&](const Device::PendingDeletion& e) { return e.accelerationStructure == handle; });
    const std::size_t storageAt = QueuePosition([&](const Device::PendingDeletion& e) { return e.buffer == storage; });
    return QueuedCount(handle) == 1 && storageAt < Device::deletionQueue.size() && handleAt < storageAt;
}

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

    // A cube BLAS and a few instances of it, so the TLASes have something in
    // them.
    Mesh cube;
    {
        std::vector<Vector<float, 3, 3>> verts {
            {-1,-1,-1}, { 1,-1,-1}, { 1, 1,-1}, {-1, 1,-1},
            {-1,-1, 1}, { 1,-1, 1}, { 1, 1, 1}, {-1, 1, 1},
        };
        std::vector<std::uint32_t> indices {
            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,
        };
        VkCommandBuffer cmd = BeginCmd();
        cube.Build(verts, indices, cmd);
        SubmitWait(cmd);
    }
    std::vector<RenderingElement3D> elements(4);
    for (RenderingElement3D& element : elements) {
        element.instance = {
            .mask                           = 0xFF,
            .accelerationStructureReference = cube.blasAddr,
        };
        MatrixRowMajor<float, 4, 3, 1>::Identity().Store(reinterpret_cast<float*>(element.instance.transform.matrix));
        RenderingElement3D::Add(&element);
    }
    std::array<TlasWithBuffer, Window::numFrames>& tlases = RenderingElement3D::tlases;
    Device::DrainDeletions();

    // ── Moving hands the TLAS over; destroying it queues the handle first ────
    {
        Build({0});
        const VkAccelerationStructureKHR handle = tlases[0].accelerationStructure;
        const VkDeviceAddress address = tlases[0].address;
        const VkBuffer storage = tlases[0].buffer.buffer;
        const VkBuffer scratch = tlases[0].scratchBuffer.buffer;
        const VkBuffer instances = tlases[0].instanceBuffer.buffer;
        const VkBuffer metadata = tlases[0].metadataBuffer.buffer;
        {
            TlasWithBuffer taken(std::move(tlases[0]));
            Check(tlases[0].accelerationStructure == VK_NULL_HANDLE && tlases[0].address == 0,
                  "a moved-from TlasWithBuffer no longer holds the acceleration structure");
            Check(taken.accelerationStructure == handle && taken.address == address && taken.buffer.buffer == storage,
                  "the one it moved to does");
            Check(Device::deletionQueue.empty(), "moving queues nothing");
        }
        Check(QueuedCount(handle) == 1, "destroying a TlasWithBuffer queues its acceleration structure");
        Check(Queued(storage) && Queued(scratch) && Queued(instances) && Queued(metadata),
              "and its storage, scratch, instance and metadata buffers");
        Check(!Device::deletionQueue.empty() && Device::deletionQueue.front().accelerationStructure == handle,
              "the acceleration structure is queued before the storage it lives in");

        // The moved-from TLAS still has its count; it has to build afresh
        // rather than refit through the buffers it gave away.
        Build({0});
        Check(tlases[0].accelerationStructure != VK_NULL_HANDLE && tlases[0].accelerationStructure != handle
              && tlases[0].address != 0 && tlases[0].builtInstanceCount == elements.size(),
              "the moved-from TlasWithBuffer builds a new TLAS");

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

    // ── Assigning an empty TlasWithBuffer resets the TLAS ───────────────────
    {
        const VkAccelerationStructureKHR handle = tlases[0].accelerationStructure;
        const VkBuffer storage = tlases[0].buffer.buffer;
        tlases[0] = TlasWithBuffer{};
        Check(tlases[0].accelerationStructure == VK_NULL_HANDLE && tlases[0].address == 0
              && tlases[0].buffer.buffer == VK_NULL_HANDLE,
              "assigning an empty TlasWithBuffer leaves no TLAS");
        Check(QueuedAhead(handle, storage),
              "and queues the acceleration structure it held once, ahead of its storage");
        Build({0});
        Check(tlases[0].accelerationStructure != VK_NULL_HANDLE && tlases[0].builtInstanceCount == elements.size(),
              "the reset TLAS builds again");
        Device::DrainDeletions();
    }

    // ── Move assignment hands the TLAS over and queues the one it replaces ───
    {
        Build({1});
        const VkAccelerationStructureKHR kept = tlases[0].accelerationStructure;
        const VkDeviceAddress keptAddress = tlases[0].address;
        const VkAccelerationStructureKHR replaced = tlases[1].accelerationStructure;
        const VkBuffer replacedStorage = tlases[1].buffer.buffer;

        tlases[1] = std::move(tlases[0]);
        Check(tlases[0].accelerationStructure == VK_NULL_HANDLE && tlases[0].address == 0
              && tlases[1].accelerationStructure == kept && tlases[1].address == keptAddress,
              "move assignment hands the acceleration structure over");
        Check(QueuedAhead(replaced, replacedStorage) && QueuedCount(kept) == 0,
              "and queues the one it replaces once, ahead of its storage");

        Build({0, 1});
        Check(tlases[0].accelerationStructure != VK_NULL_HANDLE && tlases[1].accelerationStructure == kept,
              "both TLASes build again, the one moved in in place");
        Device::DrainDeletions();
    }

    for (RenderingElement3D& element : elements) RenderingElement3D::Remove(&element);

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