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

// Issue #118: the per-frame TLAS instance+metadata host rebuild no longer copies
// every entry unconditionally — it copies (and flushes) only the slots whose
// element's host-authored data changed since this frame's buffers last saw it,
// tracked by RenderingElement3D::hostDataVersion against the per-frame
// TlasWithBuffer::uploadedVersion record. Elements left at version 0 stay
// "untracked" and are copied every frame (the pre-#118 behaviour), so callers
// that don't opt in keep working; MarkHostDataDirty() opts an element in.
//
// This drives the real hardware path (a headless Vulkan RT device, a real cube
// BLAS, RenderingElement3D::BuildTLAS into one-time command buffers) exactly as
// TLASHighWaterMark does, and asserts behaviour by reading back the host-mapped
// instanceBuffer/metadataBuffer (host-visible, and we only inspect host-written
// fields, so the mapping reflects what the copy loop wrote — a skipped slot
// retains its previous bytes, which is the whole point).
//
// What is asserted:
//   - A tracked (MarkHostDataDirty'd) element is uploaded on first build.
//   - Mutating a tracked element's host data WITHOUT marking it dirty and
//     rebuilding leaves the buffer holding the OLD value — the slot was skipped.
//   - Marking it dirty and rebuilding uploads the NEW value.
//   - An untracked element (version 0) is uploaded on every build even without
//     a mark — the backward-compatible always-copy path.
//   - Relocation on the refit path: a remove+add that nets the same instance
//     count (so BuildTLAS refits rather than rebuilding) re-uploads exactly the
//     slots whose occupying element changed, and leaves the unchanged slots
//     intact — proving globally-unique versions handle swap-and-pop without
//     tracking element identity.
//   - The Vulkan validation layer reports ZERO errors across all of the above —
//     covering the ranged FlushDevice(offset, bytes) the dirty span feeds.

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

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

std::vector<Vector<float, 3, 3>> CubeVerts(float s) {
    return {
        {-s,-s,-s}, { s,-s,-s}, { s, s,-s}, {-s, s,-s},
        {-s,-s, s}, { s,-s, s}, { s, s, s}, {-s, s, s},
    };
}
std::vector<std::uint32_t> CubeIndices() {
    return {
        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,
    };
}

VkAccelerationStructureInstanceKHR MakeInstance(VkDeviceAddress blasAddr) {
    VkAccelerationStructureInstanceKHR inst{};
    inst.transform = VkTransformMatrixKHR{{
        {1.0f, 0.0f, 0.0f, 0.0f},
        {0.0f, 1.0f, 0.0f, 0.0f},
        {0.0f, 0.0f, 1.0f, 0.0f},
    }};
    inst.mask                            = 0xFF;
    inst.accelerationStructureReference  = blasAddr;
    return inst;
}

void BuildOnce() {
    VkCommandBuffer cmd = BeginCmd();
    RenderingElement3D::BuildTLAS(cmd, 0);
    SubmitWait(cmd);
}

// Read back what the host copy loop last wrote into frame-0's metadata buffer
// slot. Host-visible and host-written, so the mapping reflects the last copy
// (or the slot's prior contents if it was skipped) without an invalidate.
std::uint32_t MetaSlot(std::uint32_t i) {
    return RenderingElement3D::tlases[0].metadataBuffer.value[i];
}
std::uint32_t CustomIndexSlot(std::uint32_t i) {
    return RenderingElement3D::tlases[0].instanceBuffer.value[i].instanceCustomIndex;
}

} // namespace

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

    Mesh cube;
    {
        auto verts = CubeVerts(1.0f);
        auto idx   = CubeIndices();
        VkCommandBuffer cmd = BeginCmd();
        cube.Build(verts, idx, cmd, RTBuildOptions{ .allowUpdate = true });
        SubmitWait(cmd);
    }
    Check(cube.blasAddr != 0, "cube BLAS produced a non-zero blasAddr");

    constexpr std::uint32_t kMaxElems = 8;
    std::vector<RenderingElement3D> pool(kMaxElems);
    for (auto& e : pool) e.instance = MakeInstance(cube.blasAddr);

    // ── 1. A tracked element is uploaded on first build. ────────────────────
    pool[0].userMetadata = 100;
    pool[0].instance.instanceCustomIndex = 100;
    pool[0].MarkHostDataDirty();
    RenderingElement3D::Add(&pool[0]);
    BuildOnce();
    Check(MetaSlot(0) == 100, "tracked element uploaded on first build (metadata)");
    Check(CustomIndexSlot(0) == 100, "tracked element uploaded on first build (customIndex)");

    // ── 2. Mutating a tracked element WITHOUT marking it dirty and rebuilding
    //       leaves the old value — the slot is skipped. The instance count is
    //       unchanged, so this is the refit path (no topology reset). ─────────
    pool[0].userMetadata = 200;
    pool[0].instance.instanceCustomIndex = 200;
    BuildOnce();
    Check(MetaSlot(0) == 100,
          "unmarked mutation is NOT uploaded — clean slot skipped (metadata stays 100)");
    Check(CustomIndexSlot(0) == 100,
          "unmarked mutation is NOT uploaded — clean slot skipped (customIndex stays 100)");

    // ── 3. Marking it dirty and rebuilding uploads the new value. ───────────
    pool[0].MarkHostDataDirty();
    BuildOnce();
    Check(MetaSlot(0) == 200, "marking dirty re-uploads the new metadata (200)");
    Check(CustomIndexSlot(0) == 200, "marking dirty re-uploads the new customIndex (200)");

    // ── 4. An untracked element (version 0) is uploaded every build even
    //       without a mark — the backward-compatible always-copy path. Add a
    //       second element to grow the count (topology change), then mutate it
    //       in place across same-count rebuilds. ─────────────────────────────
    pool[1].userMetadata = 300;          // left untracked: hostDataVersion == 0
    RenderingElement3D::Add(&pool[1]);
    BuildOnce();
    Check(MetaSlot(1) == 300, "untracked element uploaded on first build (300)");
    pool[1].userMetadata = 400;          // mutate in place, no MarkHostDataDirty
    BuildOnce();
    Check(MetaSlot(1) == 400,
          "untracked element re-uploaded every frame without a mark (400) — legacy path");
    // The tracked element [0] must NOT have been disturbed by [1]'s churn.
    Check(MetaSlot(0) == 200, "tracked element [0] untouched while [1] churns (still 200)");

    // ── 5. Relocation on the refit path. Both live elements are tracked; a
    //       remove+add that nets the same instance count takes the refit path
    //       (no topology reset), so only the slots whose occupying element
    //       changed must be re-uploaded — proving globally-unique versions
    //       handle swap-and-pop without tracking element identity. ───────────
    // Mark both current elements so neither is the always-copy legacy path.
    pool[0].userMetadata = 10; pool[0].MarkHostDataDirty();
    pool[1].userMetadata = 11; pool[1].MarkHostDataDirty();
    BuildOnce();                                   // slots: [0]=10, [1]=11
    Check(MetaSlot(0) == 10 && MetaSlot(1) == 11, "both tracked elements uploaded (10, 11)");

    // Remove pool[0] (swap-and-pop moves the last element, pool[1], into slot 0),
    // then add pool[2] (appended at slot 1). Net count is unchanged at 2, so
    // BuildTLAS refits. New mapping: slot 0 -> pool[1] (was pool[0]),
    // slot 1 -> pool[2] (was pool[1]).
    pool[2].userMetadata = 22; pool[2].MarkHostDataDirty();
    RenderingElement3D::Remove(&pool[0]);
    RenderingElement3D::Add(&pool[2]);
    BuildOnce();
    Check(MetaSlot(0) == 11,
          "relocation: slot 0 re-uploaded to the swapped-in element (pool[1] -> 11)");
    Check(MetaSlot(1) == 22,
          "relocation: slot 1 re-uploaded to the appended element (pool[2] -> 22)");

    // Unregister everything.
    while (!RenderingElement3D::elements.empty()) {
        RenderingElement3D::Remove(RenderingElement3D::elements.back());
    }

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

    // Tear down the frame-0 TLAS while the device is still alive (same rationale
    // as TLASHighWaterMark — the static tlases[] outlive the static device).
    {
        auto& t0 = RenderingElement3D::tlases[0];
        if (t0.accelerationStructure != VK_NULL_HANDLE) {
            Device::vkDestroyAccelerationStructureKHR(Device::device, t0.accelerationStructure, nullptr);
            t0.accelerationStructure = VK_NULL_HANDLE;
        }
        if (t0.instanceBuffer.buffer != VK_NULL_HANDLE) t0.instanceBuffer.Clear();
        if (t0.metadataBuffer.buffer != VK_NULL_HANDLE) t0.metadataBuffer.Clear();
        if (t0.scratchBuffer.buffer  != VK_NULL_HANDLE) t0.scratchBuffer.Clear();
        if (t0.buffer.buffer         != VK_NULL_HANDLE) t0.buffer.Clear();
    }

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