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

// ComputeShader and PipelineRTVulkan release their pipelines through the
// deferred-deletion queue (issue #101), like the other owners of Vulkan
// handles. Both used to destroy the pipeline on the spot, while a frame in
// flight could still be running it, and a static one made a Vulkan call at
// exit. Neither can be copied, as two copies would both release the pipeline;
// moves hand it over, and a move assignment queues what it replaces, a ray-
// tracing pipeline's binding table too. A PipelineRTVulkan that was never
// initialised releases nothing: its handle used to start out uninitialised and
// be destroyed regardless. Loading or initialising one again queues the
// pipeline (and binding table) it had, which used to leak.
//
// Needs a Vulkan device and no window: the test steps Device::frameCounter
// itself. The ray-tracing half skips without ray tracing. The validation check
// only means something with CRAFTER_GRAPHICS_VALIDATION set.

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

import Crafter.Graphics;
import std;

using namespace Crafter;
namespace fs = std::filesystem;

namespace {

int failures = 0;

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

fs::path ExeDir() {
    return fs::read_symlink("/proc/self/exe").parent_path();
}

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

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

std::size_t QueuedPipelines() {
    return static_cast<std::size_t>(std::ranges::count_if(Device::deletionQueue,
        [](const Device::PendingDeletion& e) { return e.pipeline != VK_NULL_HANDLE; }));
}

// Records into a one-off command buffer and runs it to completion.
void Run(const std::function<void(VkCommandBuffer)>& record) {
    const VkCommandBufferAllocateInfo allocation {
        .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, &allocation, &cmd));
    const VkCommandBufferBeginInfo begin {
        .sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO,
        .flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT,
    };
    Device::CheckVkResult(vkBeginCommandBuffer(cmd, &begin));
    record(cmd);
    Device::CheckVkResult(vkEndCommandBuffer(cmd));
    const VkSubmitInfo submit { .sType = VK_STRUCTURE_TYPE_SUBMIT_INFO, .commandBufferCount = 1, .pCommandBuffers = &cmd };
    Device::CheckVkResult(vkQueueSubmit(Device::queue, 1, &submit, VK_NULL_HANDLE));
    Device::CheckVkResult(vkQueueWaitIdle(Device::queue));
    vkFreeCommandBuffers(Device::device, Device::commandPool, 1, &cmd);
}

} // namespace

static_assert(!std::is_copy_constructible_v<ComputeShader> && !std::is_copy_assignable_v<ComputeShader>,
              "two copies of a ComputeShader would both release its pipeline");
static_assert(!std::is_copy_constructible_v<PipelineRTVulkan> && !std::is_copy_assignable_v<PipelineRTVulkan>,
              "two copies of a PipelineRTVulkan would both release its pipeline");
static_assert(std::is_nothrow_move_constructible_v<PipelineRTVulkan> && std::is_move_assignable_v<PipelineRTVulkan>,
              "a PipelineRTVulkan moves, so it can live in a container");

int main() {
    Device::Initialize();
    Device::validationErrorCount = 0;
    Device::framesInFlight = Window::numFrames;
    Device::frameCounter = 0;
    Device::DrainDeletions();

    // What Load needs in descriptor-set mode; the heap build reads nothing from it.
    DescriptorHeapVulkan heap;
    heap.Initialize(1, 1, 0);
    const fs::path compute = ExeDir() / "noop.comp.spv";

    // ── ComputeShader ────────────────────────────────────────────────────────
    VkPipeline pipeline = VK_NULL_HANDLE;
    {
        ComputeShader never;
    }
    Check(Device::deletionQueue.empty(), "a ComputeShader never loaded releases nothing");
    {
        ComputeShader shader;
        shader.Load(compute, heap);
        pipeline = shader.pipeline;
        Check(pipeline != VK_NULL_HANDLE, "the compute pipeline was created");
        // Load's VulkanShader has queued the shader module, nothing more.
        Check(QueuedPipelines() == 0, "a live ComputeShader has queued no pipeline");
    }
    Check(QueuedCount(pipeline) == 1, "destroying a ComputeShader queues its pipeline");
    Device::frameCounter = Window::numFrames;
    Device::ReclaimDeletions();
    Check(Device::deletionQueue.empty(), "the pipeline is freed on its retire frame");
    {
        ComputeShader source;
        source.Load(compute, heap);
        pipeline = source.pipeline;
        ComputeShader moved(std::move(source));
        Check(source.pipeline == VK_NULL_HANDLE && moved.pipeline == pipeline, "a move constructor hands the compute pipeline over");

        ComputeShader target;
        target.Load(compute, heap);
        const VkPipeline replaced = target.pipeline;
        target = std::move(moved);
        Check(moved.pipeline == VK_NULL_HANDLE && target.pipeline == pipeline, "a move assignment hands it over");
        Check(QueuedCount(replaced) == 1, "and queues the compute pipeline it replaced");
        Check(QueuedCount(pipeline) == 0, "the pipeline moved around is not queued while owned");
    }
    Check(QueuedCount(pipeline) == 1 && QueuedPipelines() == 2, "the compute pipeline is queued once, by its last owner");
    Device::DrainDeletions();
    {
        ComputeShader shader;
        shader.Load(compute, heap);
        pipeline = shader.pipeline;
        shader.Load(compute, heap);
        Check(shader.pipeline != pipeline && QueuedCount(pipeline) == 1 && QueuedPipelines() == 1, "loading one again queues the pipeline it had");
    }
    Device::DrainDeletions();

    // ── PipelineRTVulkan ─────────────────────────────────────────────────────
    if (!Device::rayTracingSupported) {
        std::println("SKIP  no Vulkan ray-tracing device: PipelineRTVulkan not checked");
    } else {
        {
            PipelineRTVulkan never;
        }
        Check(Device::deletionQueue.empty(), "a PipelineRTVulkan never initialised releases nothing");

        std::array<VulkanShader, 1> shaders {{ { ExeDir() / "lifetime-raygen.spv", "main", VK_SHADER_STAGE_RAYGEN_BIT_KHR, nullptr } }};
        ShaderBindingTableVulkan table;
        table.Init(shaders);
        std::array<VkRayTracingShaderGroupCreateInfoKHR, 1> raygen {{ {
            .sType = VK_STRUCTURE_TYPE_RAY_TRACING_SHADER_GROUP_CREATE_INFO_KHR, .type = VK_RAY_TRACING_SHADER_GROUP_TYPE_GENERAL_KHR,
            .generalShader = 0, .closestHitShader = VK_SHADER_UNUSED_KHR, .anyHitShader = VK_SHADER_UNUSED_KHR, .intersectionShader = VK_SHADER_UNUSED_KHR,
        } }};
        auto build = [&](PipelineRTVulkan& rt) {
            Run([&](VkCommandBuffer cmd) { rt.Init(cmd, raygen, {}, {}, table); });
        };

        {
            PipelineRTVulkan rt;
            build(rt);
            pipeline = rt.pipeline;
            Check(pipeline != VK_NULL_HANDLE && QueuedPipelines() == 0, "the ray-tracing pipeline was created, and nothing queued while it lives");
        }
        Check(QueuedCount(pipeline) == 1, "destroying a PipelineRTVulkan queues its pipeline");
        Device::DrainDeletions();
        {
            PipelineRTVulkan source;
            build(source);
            pipeline = source.pipeline;
            const VkBuffer sourceTable = source.sbtBuffer.buffer;
            const VkDeviceAddress raygenAddress = source.raygenRegion.deviceAddress;
            PipelineRTVulkan moved(std::move(source));
            Check(source.pipeline == VK_NULL_HANDLE && moved.pipeline == pipeline && moved.sbtBuffer.buffer == sourceTable
                      && moved.raygenRegion.deviceAddress == raygenAddress,
                  "a move constructor hands the ray-tracing pipeline and its binding table over");

            PipelineRTVulkan target;
            build(target);
            const VkPipeline replaced = target.pipeline;
            const VkBuffer replacedTable = target.sbtBuffer.buffer;
            target = std::move(moved);
            Check(moved.pipeline == VK_NULL_HANDLE && target.pipeline == pipeline && target.sbtBuffer.buffer == sourceTable,
                  "a move assignment hands them over");
            Check(QueuedCount(replaced) == 1 && QueuedCount(replacedTable) == 1, "and queues the pipeline and binding table it replaced");
            Check(QueuedCount(pipeline) == 0 && QueuedCount(sourceTable) == 0, "what was moved around is not queued while owned");

            build(target);
            Check(target.pipeline != pipeline && QueuedCount(pipeline) == 1 && QueuedCount(sourceTable) == 1,
                  "initialising one again queues the pipeline and binding table it had");
        }
        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;
}