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

// A re-upload into a buffer the GPU may still be reading lands on the GPU
// timeline, after the reads already submitted, instead of being written from
// the host at once. On ReBAR hardware UploadDeviceLocal (same size) and
// UploadDeviceLocalRange mapped the buffer and wrote it in place while a frame
// in flight could still read it: a mesh refitted every frame had the previous
// frame's BLAS update read this frame's vertices, or half of them.
//
// The test records a re-upload, then, before submitting it, submits and waits
// for a copy of the buffer: a frame in flight that runs after the upload was
// recorded. That copy must still see the old contents, and a copy recorded
// after the upload the new ones. Both upload calls are checked, and two
// ranged uploads to one buffer in one frame, which share that frame's staging.
//
// An empty re-upload records nothing (a copy or barrier of no bytes is
// invalid usage). And without a window no frame passes, so nothing ever
// reclaimed the staging: each re-upload added to it. DrainDeletions, after
// waiting for the queue, now lets the staging start over.
//
// Needs a Vulkan device and no window: the test steps Device::frameCounter
// itself. Run with CRAFTER_GRAPHICS_VALIDATION=1 and
// VK_KHRONOS_VALIDATION_VALIDATE_SYNC=1 to check the uploads' barriers too.

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

import Crafter.Graphics;
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;
}

constexpr std::uint32_t kCount = 256;
constexpr VkBufferUsageFlags2 kUsage = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT
                                     | VK_BUFFER_USAGE_TRANSFER_SRC_BIT;

VkCommandBuffer BeginCmd() {
    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));
    return cmd;
}

void Submit(VkCommandBuffer 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));
}

// Copies `source` into the host-visible `out`, after the uploads' consumers
// (the transfer stage here) and readable by the host once the queue is idle.
void CopyOut(VkCommandBuffer cmd, VulkanBuffer<std::uint32_t, false>& source, VulkanBuffer<std::uint32_t, true>& out) {
    const VkBufferCopy region { .srcOffset = 0, .dstOffset = 0, .size = kCount * sizeof(std::uint32_t) };
    vkCmdCopyBuffer(cmd, source.buffer, out.buffer, 1, &region);
    const VkBufferMemoryBarrier toHost {
        .sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER,
        .srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT,
        .dstAccessMask = VK_ACCESS_HOST_READ_BIT,
        .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
        .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
        .buffer = out.buffer,
        .offset = 0,
        .size = VK_WHOLE_SIZE,
    };
    vkCmdPipelineBarrier(cmd, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_HOST_BIT, 0, 0, nullptr, 1, &toHost, 0, nullptr);
}

std::vector<std::uint32_t> Contents(VulkanBuffer<std::uint32_t, true>& out) {
    out.FlushHost();
    return std::vector<std::uint32_t>(out.value, out.value + kCount);
}

std::vector<std::uint32_t> Filled(std::uint32_t base) {
    std::vector<std::uint32_t> values(kCount);
    for (std::uint32_t i = 0; i < kCount; ++i) values[i] = base + i;
    return values;
}

// One round: upload `before`, record a re-upload with `reupload`, and run a
// copy of the buffer before that re-upload's submission.
void Round(std::string_view name, const std::function<void(VulkanBuffer<std::uint32_t, false>&, VkCommandBuffer)>& reupload,
           const std::vector<std::uint32_t>& after) {
    const std::vector<std::uint32_t> before = Filled(1000);
    VulkanBuffer<std::uint32_t, false> data;
    VulkanBuffer<std::uint32_t, true> seenInFlight;
    VulkanBuffer<std::uint32_t, true> seenAfter;
    seenInFlight.Create(VK_BUFFER_USAGE_TRANSFER_DST_BIT, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT, kCount);
    seenAfter.Create(VK_BUFFER_USAGE_TRANSFER_DST_BIT, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT, kCount);

    VkCommandBuffer setup = BeginCmd();
    data.UploadDeviceLocal(kUsage, before.data(), kCount, setup, VK_ACCESS_TRANSFER_READ_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT);
    Submit(setup);
    Device::CheckVkResult(vkQueueWaitIdle(Device::queue));
    ++Device::frameCounter;

    // The next frame records its re-upload...
    VkCommandBuffer next = BeginCmd();
    reupload(data, next);
    CopyOut(next, data, seenAfter);

    // ...while a frame submitted before it still reads the buffer.
    VkCommandBuffer inFlight = BeginCmd();
    CopyOut(inFlight, data, seenInFlight);
    Submit(inFlight);
    Device::CheckVkResult(vkQueueWaitIdle(Device::queue));

    Submit(next);
    Device::CheckVkResult(vkQueueWaitIdle(Device::queue));

    Check(Contents(seenInFlight) == before, std::format("{}: a frame submitted before the upload still reads the old contents", name));
    Check(Contents(seenAfter) == after, std::format("{}: a read recorded after the upload sees the new contents", name));

    data.Clear();
    seenInFlight.Clear();
    seenAfter.Clear();
    Device::DrainDeletions();
}

} // namespace

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

    const std::vector<std::uint32_t> fresh = Filled(5000);
    Round("UploadDeviceLocal of the same size", [&](VulkanBuffer<std::uint32_t, false>& data, VkCommandBuffer cmd) {
        data.UploadDeviceLocal(kUsage, fresh.data(), kCount, cmd, VK_ACCESS_TRANSFER_READ_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT);
    }, fresh);

    Round("UploadDeviceLocalRange", [&](VulkanBuffer<std::uint32_t, false>& data, VkCommandBuffer cmd) {
        data.UploadDeviceLocalRange(fresh.data(), 0, kCount, cmd, VK_ACCESS_TRANSFER_READ_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT);
    }, fresh);

    // Two ranged uploads in one frame: the second must not overwrite the
    // first's staged bytes before the first copy has run.
    std::vector<std::uint32_t> halves = Filled(1000);
    for (std::uint32_t i = 0; i < kCount / 2; ++i) halves[i] = 7000 + i;
    for (std::uint32_t i = kCount / 2; i < kCount; ++i) halves[i] = 9000 + i;
    Round("two UploadDeviceLocalRange calls in one frame", [&](VulkanBuffer<std::uint32_t, false>& data, VkCommandBuffer cmd) {
        data.UploadDeviceLocalRange(halves.data(), 0, kCount / 2, cmd, VK_ACCESS_TRANSFER_READ_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT);
        data.UploadDeviceLocalRange(halves.data() + kCount / 2, kCount / 2, kCount / 2, cmd, VK_ACCESS_TRANSFER_READ_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT);
    }, halves);

    {
        VulkanBuffer<std::uint32_t, false> data;
        VkCommandBuffer cmd = BeginCmd();
        data.UploadDeviceLocal(kUsage, fresh.data(), kCount, cmd, VK_ACCESS_TRANSFER_READ_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT);
        Submit(cmd);
        Device::CheckVkResult(vkQueueWaitIdle(Device::queue));
        ++Device::frameCounter;
        const std::uint32_t errorsBefore = Device::validationErrorCount;
        cmd = BeginCmd();
        data.UploadDeviceLocal(kUsage, fresh.data(), 0, cmd, VK_ACCESS_TRANSFER_READ_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT);
        data.UploadDeviceLocalRange(fresh.data(), 0, 0, cmd, VK_ACCESS_TRANSFER_READ_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT);
        Submit(cmd);
        Device::CheckVkResult(vkQueueWaitIdle(Device::queue));
        Check(Device::validationErrorCount == errorsBefore, "an empty re-upload records no copy or barrier (validation stays quiet)");
        data.Clear();
        Device::DrainDeletions();
    }

    {
        // As before any window: no frames in flight, and the frame count
        // stands still.
        Device::framesInFlight = 0;
        VulkanBuffer<std::uint32_t, false> data;
        VkCommandBuffer cmd = BeginCmd();
        data.UploadDeviceLocal(kUsage, fresh.data(), kCount, cmd, VK_ACCESS_TRANSFER_READ_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT);
        Submit(cmd);
        Device::CheckVkResult(vkQueueWaitIdle(Device::queue));
        for (int round = 0; round < 8; ++round) {
            cmd = BeginCmd();
            data.UploadDeviceLocalRange(fresh.data(), 0, kCount, cmd, VK_ACCESS_TRANSFER_READ_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT);
            Submit(cmd);
            Device::CheckVkResult(vkQueueWaitIdle(Device::queue));
            Device::DrainDeletions();
        }
        const VkDeviceSize staged = data.stagingRing.empty() ? 0 : data.stagingRing[0].capacity;
        Check(staged <= kCount * sizeof(std::uint32_t),
              std::format("without a window, re-uploads between drains hold one upload's staging ({} bytes for {})", staged, kCount * sizeof(std::uint32_t)));
        data.Clear();
        Device::DrainDeletions();
        Device::framesInFlight = Window::numFrames;
    }

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