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

// RenderingElement3D::BuildTLAS with no elements registered, i.e. an app that
// builds its TLAS before the scene has anything in it. The TLAS instance and
// metadata buffers used to be sized to the instance count, so the first build
// of an empty scene created zero-size buffers: invalid usage, which NVIDIA
// answers with VK_ERROR_OUT_OF_DEVICE_MEMORY, so BuildTLAS threw. An empty
// TLAS is valid to trace against (every ray misses), so it has to build, and
// adding and removing instances afterwards has to keep working.
//
// Needs a Vulkan RT device, like BLASBuildOptions; without RT it skips.

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

// One-time command buffer helpers: record a 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);
}

// Builds frame 0's TLAS from whatever is registered; false if it threw.
bool BuildFrame0() {
    VkCommandBuffer cmd = BeginCmd();
    try {
        RenderingElement3D::BuildTLAS(cmd, 0);
    } catch (const std::exception& e) {
        std::println("      BuildTLAS threw: {}", e.what());
        vkEndCommandBuffer(cmd);
        vkFreeCommandBuffers(Device::device, Device::commandPool, 1, &cmd);
        return false;
    }
    SubmitWait(cmd);
    return true;
}

} // namespace

int main() {
    Device::Initialize();
    if (!Device::rayTracingSupported) {
        std::println("SKIP  no Vulkan ray-tracing device");
        return 77;
    }
    Device::validationErrorCount = 0;
    const TlasWithBuffer& tlas = RenderingElement3D::tlases[0];

    Check(RenderingElement3D::elements.empty(), "the scene starts empty");
    if (!BuildFrame0()) {
        Check(false, "an empty scene's TLAS builds");
        std::println("check(s) failed");
        return EXIT_FAILURE;
    }
    Check(true, "an empty scene's TLAS builds");
    Check(tlas.accelerationStructure != VK_NULL_HANDLE && tlas.address != 0,
          "it is a real acceleration structure with a device address");
    Check(tlas.builtInstanceCount == 0, "built for zero instances");

    // One instance added to the scene afterwards.
    Mesh triangle;
    {
        VkCommandBuffer cmd = BeginCmd();
        std::array<Vector<float, 3, 3>, 3> verts {{{-1, -1, 0}, {0, 1, 0}, {1, -1, 0}}};
        std::array<std::uint32_t, 3> index {{2, 1, 0}};
        triangle.Build(verts, index, cmd);
        SubmitWait(cmd);
    }
    RenderingElement3D element;
    element.instance = {
        .instanceCustomIndex                    = 0,
        .mask                                   = 0xFF,
        .instanceShaderBindingTableRecordOffset = 0,
        .flags                                  = VK_GEOMETRY_INSTANCE_FORCE_OPAQUE_BIT_KHR,
        .accelerationStructureReference         = triangle.blasAddr,
    };
    MatrixRowMajor<float, 4, 3, 1>::Identity().Store(reinterpret_cast<float*>(element.instance.transform.matrix));
    RenderingElement3D::Add(&element);
    Check(BuildFrame0(), "the TLAS rebuilds once an instance is added");
    Check(tlas.builtInstanceCount == 1, "built for one instance");

    // And back to empty.
    RenderingElement3D::Remove(&element);
    Check(BuildFrame0(), "the TLAS rebuilds once the scene is empty again");
    Check(tlas.builtInstanceCount == 0, "built for zero instances again");

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

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