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

// Issue #114: a static ImageVulkan used to keep its host-visible staging
// `buffer` (sized w*h, persistently mapped) alive for the whole life of the
// image, pinning HOST_VISIBLE / small-BAR memory long after the one upload that
// reads it has retired — and Destroy() never freed it. Update now releases it
// via VulkanBuffer::DeferredClear() right after recording the buffer→image
// copy, so the fence-keyed deletion queue (#101/#102) frees it once that
// submit's frame has cleared. A `streamed` image (the FontAtlas) keeps its
// staging, since it re-fills the persistent map and re-uploads every frame.
//
// This drives the real upload path on a headless device — no swapchain/window
// needed, a buffer→image copy only touches the queue + command pool — and
// asserts:
//   - After a static Update, image.buffer owns no handle (released), and
//     exactly that one allocation was handed to the deletion queue tagged with
//     the current frame (so it retires later, not now).
//   - The upload + readback still complete with ZERO validation errors and the
//     image reads back byte-equal to the staged source — proof the staging
//     genuinely outlived the submit (the queue had not yet retired it), so
//     releasing it is no use-after-free.
//   - The deferred entry retires (and frees) only once framesInFlight frames
//     elapse, and not before.
//   - A `streamed` image keeps its staging across Update and UpdateRegion (the
//     FontAtlas contract), and Destroy() then frees it validation-clean.

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

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

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

// RGBA8 texel as a 4-byte word, matching VK_FORMAT_R8G8B8A8_UNORM.
using Texel = std::uint32_t;

constexpr std::uint16_t kW = 4;
constexpr std::uint16_t kH = 4;
constexpr std::uint32_t kCount = static_cast<std::uint32_t>(kW) * kH;

} // namespace

int main() {
    Device::Initialize();
    Device::validationErrorCount = 0;

    // Control the deferred-deletion clock: enqueue at frame 0, retire after 2.
    Device::framesInFlight = 2;
    Device::frameCounter   = 0;
    Device::deletionQueue.clear();

    // ── Static texture: staging is released right after the upload copy. ──────
    ImageVulkan<Texel> tex;
    {
        VkCommandBuffer cmd = BeginCmd();
        tex.Create(kW, kH, /*mipLevels*/ 1, cmd,
                   VK_FORMAT_R8G8B8A8_UNORM,
                   VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_TRANSFER_SRC_BIT
                       | VK_IMAGE_USAGE_SAMPLED_BIT,
                   VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);

        Check(tex.buffer.buffer != VK_NULL_HANDLE,
              "Create allocated the host-visible staging buffer");
        Check(!tex.streamed, "a default ImageVulkan is not streamed");

        // Fill the staging map with a known pattern, then upload.
        for (std::uint32_t i = 0; i < kCount; ++i) {
            tex.buffer.value[i] = 0xA0B0C0D0u + i;
        }
        tex.Update(cmd, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);

        // ── The fix: staging released to the queue at record time, not pinned. ─
        Check(tex.buffer.buffer == VK_NULL_HANDLE,
              "static Update releases the staging buffer (no handle pinned)");
        Check(Device::deletionQueue.size() == 1,
              "exactly one allocation (the staging) handed to the deletion queue");
        const bool taggedNow = !Device::deletionQueue.empty()
            && Device::deletionQueue.front().retireAfter == 0
            && Device::deletionQueue.front().buffer != VK_NULL_HANDLE;
        Check(taggedNow,
              "deferred staging is tagged with the enqueue frame and still a live handle");

        // Read the image back to prove the copy saw valid staging. The upload's
        // recorded copy still references the now-released staging; it must
        // outlive THIS submit (queue retires at frame 2, we're at 0). Move the
        // image to TRANSFER_SRC and copy into a host-visible readback buffer in
        // the same command buffer.
        VkImageMemoryBarrier toSrc {
            .sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER,
            .srcAccessMask = VK_ACCESS_SHADER_READ_BIT,
            .dstAccessMask = VK_ACCESS_TRANSFER_READ_BIT,
            .oldLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
            .newLayout = VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL,
            .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
            .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
            .image = tex.image,
            .subresourceRange = { VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1 },
        };
        vkCmdPipelineBarrier(cmd, VK_PIPELINE_STAGE_RAY_TRACING_SHADER_BIT_KHR,
                             VK_PIPELINE_STAGE_TRANSFER_BIT, 0, 0, nullptr,
                             0, nullptr, 1, &toSrc);

        VulkanBuffer<Texel, true> readback;
        readback.Create(VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT,
                        VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT, kCount);
        VkBufferImageCopy region {};
        region.imageSubresource = { VK_IMAGE_ASPECT_COLOR_BIT, 0, 0, 1 };
        region.imageExtent = { kW, kH, 1 };
        vkCmdCopyImageToBuffer(cmd, tex.image, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL,
                               readback.buffer, 1, &region);

        SubmitWait(cmd);
        Check(Device::validationErrorCount == 0,
              "static upload + readback raised no validation errors");

        readback.FlushHost();
        bool match = true;
        for (std::uint32_t i = 0; i < kCount; ++i) {
            if (readback.value[i] != 0xA0B0C0D0u + i) { match = false; break; }
        }
        Check(match,
              "image reads back byte-equal to the staged source "
              "(released staging outlived the submit)");
    }

    // ── Retire timing: the staging frees once framesInFlight frames elapse,
    //    and not a frame before. We've already wait-idled, so freeing is safe. ─
    Device::ReclaimDeletions();
    Check(Device::deletionQueue.size() == 1,
          "staging is NOT freed before framesInFlight frames have elapsed");
    Device::frameCounter = 2;  // retireAfter(0) + framesInFlight(2) <= 2
    Device::ReclaimDeletions();
    Check(Device::deletionQueue.empty(),
          "staging is freed once its retire frame is reached");

    // Destroying the image after its staging was released must not double-free
    // (buffer handle is already null) and must stay validation-clean.
    Device::validationErrorCount = 0;
    tex.Destroy();
    Check(Device::validationErrorCount == 0,
          "Destroy after a released-staging upload is validation-clean");

    // ── Streamed image: the persistent staging is kept across uploads. ────────
    Device::frameCounter = 0;
    Device::deletionQueue.clear();
    Device::validationErrorCount = 0;
    {
        ImageVulkan<Texel> atlas;
        VkCommandBuffer cmd = BeginCmd();
        atlas.Create(kW, kH, /*mipLevels*/ 1, cmd,
                     VK_FORMAT_R8G8B8A8_UNORM,
                     VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_SAMPLED_BIT,
                     VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
                     VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
                     /*streamed*/ true);
        Check(atlas.streamed, "streamed=true is recorded on the image");

        for (std::uint32_t i = 0; i < kCount; ++i) atlas.buffer.value[i] = i;
        atlas.Update(cmd, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);
        Check(atlas.buffer.buffer != VK_NULL_HANDLE,
              "streamed Update keeps the persistent staging buffer");
        // UpdateRegion (the FontAtlas re-upload path) must also keep it.
        atlas.UpdateRegion(cmd, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL, 0, 0, kW, kH);
        Check(atlas.buffer.buffer != VK_NULL_HANDLE,
              "streamed UpdateRegion keeps the persistent staging buffer");
        Check(Device::deletionQueue.empty(),
              "a streamed image enqueues nothing for deferred deletion");
        SubmitWait(cmd);
        Check(Device::validationErrorCount == 0,
              "streamed upload raised no validation errors");

        // Destroy frees the still-live persistent staging (the leak #114 names).
        atlas.Destroy();
        Check(atlas.buffer.buffer == VK_NULL_HANDLE,
              "Destroy frees the streamed image's persistent staging buffer");
        Check(Device::validationErrorCount == 0,
              "streamed Destroy (freeing live staging) is validation-clean");
    }

    std::println("{}", failures == 0 ? "ALL PASS" : "FAILURES PRESENT");
    return failures == 0 ? 0 : 1;
}