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

// Issue #64: TLAS host-visible input buffers (instanceBuffer / metadataBuffer)
// grow on a high-water mark instead of being reallocated to the exact instance
// count on every topology change. These buffers only ever need to hold *at
// least* primitiveCount entries — the AS build reads exactly primitiveCount of
// them — so an allocation left over from a larger earlier frame is reused.
//
// This drives the real hardware path: a headless Vulkan RT device (no swapchain
// needed — a TLAS build only touches the queue + command pool), a real cube
// BLAS, and RenderingElement3D::BuildTLAS recorded into one-time command
// buffers at a sequence of instance counts.
//
// What is asserted:
//   - First build at a given count ALLOCATES the host inputs (non-null handle,
//     non-zero device address, size == count·sizeof(entry)).
//   - Growing PAST the current capacity REALLOCATES (new VkBuffer handle, new
//     address, larger size) — the only case that still pays the realloc.
//   - SHRINKING reuses the existing allocation unchanged (same handle, same
//     address, same size) — the core win of this issue, and the case the
//     pre-fix code reallocated on.
//   - Growing back up but still WITHIN the high-water capacity also reuses it.
//   - Growing to EXACTLY the capacity reuses it (boundary: count > capacity,
//     not >=).
//   - instanceBuffer and metadataBuffer grow in lockstep (one capacity check
//     governs both).
//   - A same-count rebuild takes the refit (UPDATE) path: builtInstanceCount
//     is unchanged and the inputs are obviously not reallocated.
//   - builtInstanceCount tracks the live count across every build (the AS
//     itself is still rebuilt on a count change — the fix only spares the two
//     host buffers, not the AS storage/scratch).
//   - The Vulkan validation layer reports ZERO errors across all of the above
//     — the strongest check that feeding an oversized instance buffer to the
//     AS build (with tlasRangeInfo.primitiveCount < capacity) is spec-correct.
//
// Issue #65: this test also guards the instance buffer's memory-type request —
// HOST_VISIBLE (required) with a best-effort DEVICE_LOCAL preference and no
// mandatory HOST_COHERENT, so the per-frame GPU accesses stay in local VRAM
// (BAR) rather than crossing PCIe. See the assertion block after the first
// build. The zero-validation-errors check below also covers the dropped
// HOST_COHERENT: the FlushDevice(cmd, ...) in BuildTLAS gates on the chosen
// type's coherency flag, so a non-coherent BAR type still flushes correctly.
//
// Issue #75: the metadata buffer gets the identical upgrade — it is CPU-written
// every frame and read by the ray shaders as a STORAGE_BUFFER every frame, so
// the same HOST_VISIBLE | prefer-DEVICE_LOCAL, no-mandatory-HOST_COHERENT
// request keeps those per-frame shader reads in local VRAM. A second assertion
// block after the first build checks its chosen memory type the same way. The
// dropped HOST_COHERENT is covered by the zero-validation-errors check: the
// FlushDevice() added after the copy loop in BuildTLAS makes the host writes
// available on a non-coherent type and self-gates to a no-op on a coherent one.
//
// Validation layers are required for the last check to be meaningful; the build
// marks this test as needing the SDK layers.

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

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

// A unit cube (8 verts, 12 triangles) — enough topology for a real BLAS.
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,
    };
}

// One TLAS instance referencing `blasAddr`, identity transform, visible mask.
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;
}

// Register exactly `n` elements from a stable (reserved, never-reallocated)
// pool. Clears the current registration first so the live count is exactly n.
void SetCount(std::vector<RenderingElement3D>& pool, std::uint32_t n) {
    while (!RenderingElement3D::elements.empty()) {
        RenderingElement3D::Remove(RenderingElement3D::elements.back());
    }
    for (std::uint32_t i = 0; i < n; ++i) {
        RenderingElement3D::Add(&pool[i]);
    }
}

// Build the frame-0 TLAS for the currently-registered elements.
void BuildOnce() {
    VkCommandBuffer cmd = BeginCmd();
    RenderingElement3D::BuildTLAS(cmd, 0);
    SubmitWait(cmd);
}

// Snapshot of the host-input buffer identity, for before/after comparison.
struct InputSnapshot {
    VkBuffer       instBuf;
    VkDeviceAddress instAddr;
    std::uint32_t  instSize;
    VkBuffer       metaBuf;
    VkDeviceAddress metaAddr;
    std::uint32_t  metaSize;
};
InputSnapshot Snapshot() {
    auto& tlas = RenderingElement3D::tlases[0];
    return {
        tlas.instanceBuffer.buffer, tlas.instanceBuffer.address, tlas.instanceBuffer.size,
        tlas.metadataBuffer.buffer, tlas.metadataBuffer.address, tlas.metadataBuffer.size,
    };
}

// Both host inputs untouched (same VkBuffer handle, address and size).
bool Reused(const InputSnapshot& a, const InputSnapshot& b) {
    return a.instBuf == b.instBuf && a.instAddr == b.instAddr && a.instSize == b.instSize
        && a.metaBuf == b.metaBuf && a.metaAddr == b.metaAddr && a.metaSize == b.metaSize;
}

constexpr std::uint32_t kEntrySize = sizeof(VkAccelerationStructureInstanceKHR);

} // namespace

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

    // One real cube BLAS that every TLAS instance references.
    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");

    // Stable backing store for the elements. Reserve to the max count this
    // test ever registers so the vector never reallocates (the static
    // `elements` array holds raw pointers into it).
    constexpr std::uint32_t kMaxElems = 32;
    std::vector<RenderingElement3D> pool(kMaxElems);
    for (auto& e : pool) e.instance = MakeInstance(cube.blasAddr);

    // ── 1. First build at 4 instances → allocates the host inputs. ──────────
    SetCount(pool, 4);
    BuildOnce();
    InputSnapshot s4 = Snapshot();
    Check(s4.instBuf != VK_NULL_HANDLE, "first build allocated the instance buffer");
    Check(s4.metaBuf != VK_NULL_HANDLE, "first build allocated the metadata buffer");
    Check(s4.instAddr != 0, "instance buffer has a non-zero device address");
    Check(s4.instSize == 4 * kEntrySize, "instance buffer sized for 4 entries");
    Check(RenderingElement3D::tlases[0].builtInstanceCount == 4,
          "builtInstanceCount == 4 after first build");

    // ── Issue #65: the instance buffer is allocated HOST_VISIBLE (required,
    // mappable) with a best-effort DEVICE_LOCAL preference (BAR/VRAM) and the
    // mandatory HOST_COHERENT dropped — so the per-frame compute-write of the
    // transform field and the AS build's read of the instances both stay in
    // local VRAM instead of crossing PCIe. The chosen type must match exactly
    // what GetMemoryType resolves for this buffer's memoryTypeBits with that
    // prefer-but-don't-require request: DEVICE_LOCAL when a host-visible
    // device-local type is reachable (resizable BAR), plain HOST_VISIBLE
    // otherwise. Asserting equality with the recomputed selection makes the
    // check device-independent — it tracks the call-site request, not the
    // particular GPU this runs on. ────────────────────────────────────────────
    {
        auto& ib = RenderingElement3D::tlases[0].instanceBuffer;
        VkMemoryRequirements req{};
        vkGetBufferMemoryRequirements(Device::device, ib.buffer, &req);
        std::uint32_t expectIdx = Device::GetMemoryType(
            req.memoryTypeBits,
            VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT,
            VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);
        VkMemoryPropertyFlags expectFlags =
            Device::memoryProperties.memoryTypes[expectIdx].propertyFlags;
        Check((ib.memoryPropertyFlagsChosen & VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT) != 0,
              "instance buffer memory is HOST_VISIBLE (mappable — the required flag)");
        Check(ib.memoryPropertyFlagsChosen == expectFlags,
              "instance buffer memory matches the HOST_VISIBLE | prefer-DEVICE_LOCAL request (#65)");

        bool barReachable = false;
        constexpr VkMemoryPropertyFlags kBar =
            VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT;
        for (std::uint32_t i = 0; i < Device::memoryProperties.memoryTypeCount; ++i) {
            if ((req.memoryTypeBits & (1u << i))
                && (Device::memoryProperties.memoryTypes[i].propertyFlags & kBar) == kBar) {
                barReachable = true;
                break;
            }
        }
        if (barReachable) {
            Check((ib.memoryPropertyFlagsChosen & VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT) != 0,
                  "instance buffer landed on a DEVICE_LOCAL (BAR/VRAM) type where one is reachable");
        } else {
            std::println("INFO  no host-visible device-local type reachable — "
                         "fell back to plain HOST_VISIBLE (expected on non-BAR GPUs)");
        }
    }

    // ── Issue #75: the metadata buffer carries the same memory-type request as
    // the instance buffer — HOST_VISIBLE (required, mappable) with a best-effort
    // DEVICE_LOCAL preference (BAR/VRAM) and the mandatory HOST_COHERENT dropped.
    // It is CPU-written every frame and read by the ray shaders every frame, so
    // landing it in local VRAM keeps those per-frame storage-buffer reads off
    // PCIe. Same device-independent check: assert the recorded chosen flags equal
    // what GetMemoryType resolves for this buffer's memoryTypeBits with that
    // request, and that DEVICE_LOCAL is chosen wherever a host-visible
    // device-local type is reachable. ───────────────────────────────────────────
    {
        auto& mb = RenderingElement3D::tlases[0].metadataBuffer;
        VkMemoryRequirements req{};
        vkGetBufferMemoryRequirements(Device::device, mb.buffer, &req);
        std::uint32_t expectIdx = Device::GetMemoryType(
            req.memoryTypeBits,
            VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT,
            VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);
        VkMemoryPropertyFlags expectFlags =
            Device::memoryProperties.memoryTypes[expectIdx].propertyFlags;
        Check((mb.memoryPropertyFlagsChosen & VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT) != 0,
              "metadata buffer memory is HOST_VISIBLE (mappable — the required flag)");
        Check(mb.memoryPropertyFlagsChosen == expectFlags,
              "metadata buffer memory matches the HOST_VISIBLE | prefer-DEVICE_LOCAL request (#75)");

        bool barReachable = false;
        constexpr VkMemoryPropertyFlags kBar =
            VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT;
        for (std::uint32_t i = 0; i < Device::memoryProperties.memoryTypeCount; ++i) {
            if ((req.memoryTypeBits & (1u << i))
                && (Device::memoryProperties.memoryTypes[i].propertyFlags & kBar) == kBar) {
                barReachable = true;
                break;
            }
        }
        if (barReachable) {
            Check((mb.memoryPropertyFlagsChosen & VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT) != 0,
                  "metadata buffer landed on a DEVICE_LOCAL (BAR/VRAM) type where one is reachable");
        } else {
            std::println("INFO  no host-visible device-local type reachable — "
                         "metadata fell back to plain HOST_VISIBLE (expected on non-BAR GPUs)");
        }
    }

    // ── 2. Grow to 16 (past capacity) → REALLOCATES. ────────────────────────
    SetCount(pool, 16);
    BuildOnce();
    InputSnapshot s16 = Snapshot();
    Check(s16.instBuf != s4.instBuf, "growing past capacity reallocated the instance buffer");
    Check(s16.metaBuf != s4.metaBuf, "growing past capacity reallocated the metadata buffer");
    Check(s16.instSize == 16 * kEntrySize, "instance buffer grew to 16 entries");
    Check(RenderingElement3D::tlases[0].builtInstanceCount == 16,
          "builtInstanceCount == 16 after growth");

    // ── 3. Shrink to 2 → REUSES the 16-entry allocation (the core win). ─────
    SetCount(pool, 2);
    BuildOnce();
    InputSnapshot s2 = Snapshot();
    Check(Reused(s16, s2),
          "shrinking to 2 reused the existing buffers (no realloc — high-water mark)");
    Check(s2.instSize == 16 * kEntrySize,
          "instance buffer kept its 16-entry high-water size after shrink");
    Check(RenderingElement3D::tlases[0].builtInstanceCount == 2,
          "builtInstanceCount tracks the live count (2) even on the reuse path");

    // ── 4. Grow to 10 (still within the high-water capacity) → REUSES. ──────
    SetCount(pool, 10);
    BuildOnce();
    InputSnapshot s10 = Snapshot();
    Check(Reused(s16, s10),
          "growing to 10 (≤ capacity 16) reused the existing buffers");

    // ── 5. Grow to exactly 16 (== capacity) → REUSES (boundary: > not >=). ──
    SetCount(pool, 16);
    BuildOnce();
    InputSnapshot s16b = Snapshot();
    Check(Reused(s16, s16b),
          "growing to exactly the capacity (16) reused the existing buffers");

    // ── 6. Same count again → refit (UPDATE) path, inputs untouched. ────────
    BuildOnce();
    InputSnapshot s16c = Snapshot();
    Check(Reused(s16, s16c), "same-count rebuild left the inputs untouched");
    Check(RenderingElement3D::tlases[0].builtInstanceCount == 16,
          "same-count rebuild kept builtInstanceCount == 16 (took the refit path)");

    // ── 7. Grow to 32 (past the high-water) → REALLOCATES again. ────────────
    SetCount(pool, 32);
    BuildOnce();
    InputSnapshot s32 = Snapshot();
    Check(s32.instBuf != s16.instBuf, "growing past the high-water (32) reallocated again");
    Check(s32.instSize == 32 * kEntrySize, "instance buffer grew to 32 entries");

    // Unregister everything so nothing dangles past the pool's lifetime.
    SetCount(pool, 0);

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

    // Tear down the frame-0 TLAS while the Vulkan device is still alive. The
    // static tlases[] array is destroyed at process exit, by which point the
    // device (also static) may already be gone — so its VulkanBuffer
    // destructors would fault. Releasing here keeps the exit path clean.
    {
        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;
}