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

// Regression test for dynamic descriptor-heap indexing in ray-tracing hit
// shaders. On NVIDIA 610.43.02, selecting a VK_EXT_descriptor_heap descriptor
// with a runtime index inside a hit shader lost the device, for storage
// buffers and sampled images alike (README, "Native RT limitation"), so the
// examples select hit-shader descriptors through spec constants. This renders
// two instances of one triangle whose closest-hit picks a storage buffer and a
// texture by a runtime index: instance 0 the same set everywhere, instance 1 a
// different set in every other pixel column, so neighbouring lanes of a wave
// select different descriptors. It reads the frame back and checks every
// pixel got its own set's buffer red and texture green.
//
// The fault fixed in NVIDIA 615.78.08, so on older NVIDIA drivers the test
// skips instead of aborting the suite. GPU-assisted validation hides the
// fault (examples/Sponza/README.md), so with CRAFTER_GRAPHICS_VALIDATION set
// the scene is also rendered in a child process with validation off; the
// validation check only means something with it set. Needs a live Wayland
// compositor and a Vulkan RT device; without RT it skips.

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

// PNG decode for the read-back frame. The library only compiles the writer
// (Window::SaveFrame); STB_IMAGE_STATIC keeps the reader private to this file.
#define STB_IMAGE_IMPLEMENTATION
#define STB_IMAGE_STATIC
#define STBI_ONLY_PNG
#include "../../lib/stb_image.h"

import Crafter.Graphics;
import Crafter.Event;
import Crafter.Math;
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);
    // Flushed per line: a driver with the fault aborts the process from
    // inside Render(), and the checks before it should still show.
    std::fflush(stdout);
    if (!ok) ++failures;
}

// The build places the compiled shaders next to the test binary, and the
// test runner's working directory is not that directory.
fs::path ExeDir() {
    return fs::read_symlink("/proc/self/exe").parent_path();
}

int RunCommand(const std::string& cmd) {
    int status = std::system(cmd.c_str());
    if (status == -1) return -1;
    return (status & 0x7f) == 0 ? ((status >> 8) & 0xff) : 128 + (status & 0x7f);
}

struct RGBA8 { std::uint8_t r, g, b, a; };

struct Frame {
    int width = 0;
    int height = 0;
    std::vector<std::uint8_t> rgba;

    std::array<std::uint8_t, 3> At(int x, int y) const {
        const std::size_t i = (static_cast<std::size_t>(y) * width + x) * 4;
        return { rgba[i], rgba[i + 1], rgba[i + 2] };
    }
};

Frame LoadFrame(const fs::path& path) {
    Frame frame;
    int channels = 0;
    stbi_uc* pixels = stbi_load(path.string().c_str(), &frame.width, &frame.height, &channels, STBI_rgb_alpha);
    if (pixels != nullptr) {
        frame.rgba.assign(pixels, pixels + static_cast<std::size_t>(frame.width) * frame.height * 4);
        stbi_image_free(pixels);
    }
    return frame;
}

bool Near(std::array<std::uint8_t, 3> pixel, std::array<std::uint8_t, 3> expected) {
    for (std::size_t c = 0; c < 3; ++c) {
        if (std::abs(static_cast<int>(pixel[c]) - static_cast<int>(expected[c])) > 8) return false;
    }
    return true;
}

void CheckPixel(const Frame& frame, int x, int y, std::array<std::uint8_t, 3> expected, std::string_view what) {
    const std::array<std::uint8_t, 3> pixel = frame.At(x, y);
    Check(Near(pixel, expected), std::format("{} at ({}, {}) is ({}, {}, {}), expected ({}, {}, {})",
          what, x, y, pixel[0], pixel[1], pixel[2], expected[0], expected[1], expected[2]));
}

// The pixel column world x on the triangles' plane (z = 100, 400 in front of
// the camera) lands in, for the raygen's 60° vertical field of view.
int FrameX(float worldX, const Frame& frame) {
    const float aspect = static_cast<float>(frame.width) / static_cast<float>(frame.height);
    const float ndc = worldX / (400.0f * aspect * std::tan(std::numbers::pi_v<float> / 6.0f));
    return static_cast<int>((0.5f + ndc * 0.5f) * static_cast<float>(frame.width));
}

// Spec constants of the closest-hit, matching dynidx-closesthit.glsl.
struct HitConstants {
    std::uint16_t bufferStart;
    std::uint16_t pad;
    std::uint32_t textureBase;
    std::uint32_t colourBase;
};

// What each descriptor set holds: the red its buffer gives and the green its
// texture gives.
constexpr std::array<float, 2> bufferRed = {1.0f, 0.5f};
constexpr std::array<std::uint8_t, 2> textureGreen = {64, 191};
constexpr std::array<std::array<std::uint8_t, 3>, 2> setColour = {{{255, 64, 0}, {128, 191, 0}}};

// NVIDIA's proprietary driver packs its version into 10.8.8.6 bits.
std::tuple<std::uint32_t, std::uint32_t, std::uint32_t> NvidiaVersion() {
    const std::uint32_t v = Device::deviceProperties.driverVersion;
    return {(v >> 22) & 0x3ff, (v >> 14) & 0xff, (v >> 6) & 0xff};
}

// Renders the scene into `out` and checks the pixels. Device::Initialize has
// already run. Returns the number of failed checks.
int RenderScene(const fs::path& out) {
    const int failuresBefore = failures;
    Device::validationErrorCount = 0;

    Window window(640, 360, "HeapDynamicIndexHit test");
    VkCommandBuffer cmd = window.StartInit();

    // Image slots: 0 = the frame's swapchain image, 1-2 = the textures.
    // Buffer slots: 0 = the TLAS, 1-2 = the colour buffers. Sampler slot 0.
    DescriptorHeapVulkan descriptorHeap;
    descriptorHeap.Initialize(3, 3, 1);

    VkSpecializationMapEntry raygenEntry = { .constantID = 0, .offset = 0, .size = sizeof(std::uint16_t) };
    VkSpecializationInfo raygenSpec = {
        .mapEntryCount = 1,
        .pMapEntries = &raygenEntry,
        .dataSize = sizeof(std::uint16_t),
        .pData = &descriptorHeap.bufferStartElement,
    };
    HitConstants hitConstants = {
        .bufferStart = descriptorHeap.bufferStartElement,
        .pad = 0,
        .textureBase = 1,
        .colourBase = 1,
    };
    std::array<VkSpecializationMapEntry, 3> hitEntries {{
        { .constantID = 0, .offset = offsetof(HitConstants, bufferStart), .size = sizeof(std::uint16_t) },
        { .constantID = 1, .offset = offsetof(HitConstants, textureBase), .size = sizeof(std::uint32_t) },
        { .constantID = 2, .offset = offsetof(HitConstants, colourBase), .size = sizeof(std::uint32_t) },
    }};
    VkSpecializationInfo hitSpec = {
        .mapEntryCount = static_cast<std::uint32_t>(hitEntries.size()),
        .pMapEntries = hitEntries.data(),
        .dataSize = sizeof(HitConstants),
        .pData = &hitConstants,
    };

    const fs::path dir = ExeDir();
    std::array<VulkanShader, 3> shaders{{
        {dir / "dynidx-raygen.spv", "main", VK_SHADER_STAGE_RAYGEN_BIT_KHR, &raygenSpec},
        {dir / "dynidx-miss.spv", "main", VK_SHADER_STAGE_MISS_BIT_KHR, nullptr},
        {dir / "dynidx-closesthit.spv", "main", VK_SHADER_STAGE_CLOSEST_HIT_BIT_KHR, &hitSpec},
    }};
    ShaderBindingTableVulkan shaderTable;
    shaderTable.Init(shaders);

    std::array<VkRayTracingShaderGroupCreateInfoKHR, 1> raygenGroups {{ {
        .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,
    } }};
    std::array<VkRayTracingShaderGroupCreateInfoKHR, 1> missGroups {{ {
        .sType = VK_STRUCTURE_TYPE_RAY_TRACING_SHADER_GROUP_CREATE_INFO_KHR,
        .type = VK_RAY_TRACING_SHADER_GROUP_TYPE_GENERAL_KHR,
        .generalShader = 1, .closestHitShader = VK_SHADER_UNUSED_KHR,
        .anyHitShader = VK_SHADER_UNUSED_KHR, .intersectionShader = VK_SHADER_UNUSED_KHR,
    } }};
    std::array<VkRayTracingShaderGroupCreateInfoKHR, 1> hitGroups {{ {
        .sType = VK_STRUCTURE_TYPE_RAY_TRACING_SHADER_GROUP_CREATE_INFO_KHR,
        .type = VK_RAY_TRACING_SHADER_GROUP_TYPE_TRIANGLES_HIT_GROUP_KHR,
        .generalShader = VK_SHADER_UNUSED_KHR, .closestHitShader = 2,
        .anyHitShader = VK_SHADER_UNUSED_KHR, .intersectionShader = VK_SHADER_UNUSED_KHR,
    } }};
    PipelineRTVulkan pipeline;
    pipeline.Init(cmd, raygenGroups, missGroups, hitGroups, shaderTable);

    // One triangle, two instances 100 units either side of the centre line.
    Mesh triangleMesh;
    std::array<Vector<float, 3, 3>, 3> verts {{{-75, -75, 100}, {0, 75, 100}, {75, -75, 100}}};
    std::array<std::uint32_t, 3> index {{2, 1, 0}};
    triangleMesh.Build(verts, index, cmd);

    std::array<RenderingElement3D, 2> renderers{};
    for (std::uint32_t i = 0; i < renderers.size(); ++i) {
        renderers[i].instance = {
            .instanceCustomIndex                    = i,
            .mask                                   = 0xFF,
            .instanceShaderBindingTableRecordOffset = 0,
            .flags                                  = VK_GEOMETRY_INSTANCE_FORCE_OPAQUE_BIT_KHR,
            .accelerationStructureReference         = triangleMesh.blasAddr,
        };
        MatrixRowMajor<float, 4, 3, 1>::Translation(i == 0 ? -100.0f : 100.0f, 0.0f, 0.0f)
            .Store(reinterpret_cast<float*>(renderers[i].instance.transform.matrix));
        RenderingElement3D::Add(&renderers[i]);
    }
    for (std::uint32_t f = 0; f < Window::numFrames; ++f) RenderingElement3D::BuildTLAS(cmd, f);

    // Per descriptor set: a 2x2 texture whose green the hit shader reads, and
    // a one-vec4 buffer whose red it reads.
    std::array<Image2D<RGBA8>, 2> textures;
    std::array<VulkanBuffer<float, true>, 2> colourBuffers;
    for (std::size_t i = 0; i < textures.size(); ++i) {
        textures[i].Create(2, 2, 1, cmd, VK_FORMAT_R8G8B8A8_UNORM,
                           VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_SAMPLED_BIT,
                           VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);
        for (std::size_t p = 0; p < 4; ++p) textures[i].buffer.value[p] = {0, textureGreen[i], 0, 255};
        textures[i].Update(cmd, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);

        colourBuffers[i].Create(VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT,
                                VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT, 4);
        colourBuffers[i].value[0] = bufferRed[i];
        colourBuffers[i].value[1] = 0.0f;
        colourBuffers[i].value[2] = 0.0f;
        colourBuffers[i].value[3] = 1.0f;
        colourBuffers[i].FlushDevice();
    }

    window.FinishInit();

    auto imgSlots = descriptorHeap.AllocateImageSlots(3);
    auto bufSlots = descriptorHeap.AllocateBufferSlots(3);
    auto samplerSlots = descriptorHeap.AllocateSamplerSlots(1);

    std::array<VkImageViewCreateInfo, 2> textureViews{};
    std::array<VkImageDescriptorInfoEXT, 2> textureInfos{};
    std::array<VkDeviceAddressRangeKHR, 2> colourRanges{};
    for (std::size_t i = 0; i < 2; ++i) {
        textureViews[i] = {
            .sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO,
            .image = textures[i].image,
            .viewType = VK_IMAGE_VIEW_TYPE_2D,
            .format = VK_FORMAT_R8G8B8A8_UNORM,
            .subresourceRange = { .aspectMask = VK_IMAGE_ASPECT_COLOR_BIT, .baseMipLevel = 0,
                                  .levelCount = 1, .baseArrayLayer = 0, .layerCount = 1 },
        };
        textureInfos[i] = {
            .sType = VK_STRUCTURE_TYPE_IMAGE_DESCRIPTOR_INFO_EXT,
            .pView = &textureViews[i],
            .layout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
        };
        colourRanges[i] = { .address = colourBuffers[i].address, .size = 4 * sizeof(float) };
    }

    // Written again whenever the swapchain is recreated, since the output
    // image descriptors point at its images.
    auto writeDescriptors = [&] {
        std::array<VkDeviceAddressRangeKHR, Window::numFrames> tlasRanges{};
        std::array<VkImageDescriptorInfoEXT, Window::numFrames> outInfos{};
        std::vector<VkResourceDescriptorInfoEXT> resources;
        std::vector<VkHostAddressRangeEXT> destinations;
        auto add = [&](std::uint32_t f, VkResourceDescriptorInfoEXT resource, std::uint32_t offset, std::uint32_t size) {
            resources.push_back(resource);
            destinations.push_back({ .address = descriptorHeap.resourceHeap[f].value + offset, .size = size });
        };
        const std::uint32_t imageSize = Device::descriptorHeapProperties.imageDescriptorSize;
        const std::uint32_t bufferSize = Device::descriptorHeapProperties.bufferDescriptorSize;
        for (std::uint32_t f = 0; f < Window::numFrames; ++f) {
            tlasRanges[f] = { .address = RenderingElement3D::tlases[f].address };
            outInfos[f] = {
                .sType = VK_STRUCTURE_TYPE_IMAGE_DESCRIPTOR_INFO_EXT,
                .pView = &window.imageViews[f],
                .layout = VK_IMAGE_LAYOUT_GENERAL,
            };
            add(f, { .sType = VK_STRUCTURE_TYPE_RESOURCE_DESCRIPTOR_INFO_EXT,
                     .type = VK_DESCRIPTOR_TYPE_ACCELERATION_STRUCTURE_KHR,
                     .data = { .pAddressRange = &tlasRanges[f] } },
                descriptorHeap.BufferByteOffset(bufSlots.firstElement), bufferSize);
            add(f, { .sType = VK_STRUCTURE_TYPE_RESOURCE_DESCRIPTOR_INFO_EXT,
                     .type = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
                     .data = { .pImage = &outInfos[f] } },
                descriptorHeap.ImageByteOffset(imgSlots.firstElement), imageSize);
            for (std::uint16_t i = 0; i < 2; ++i) {
                add(f, { .sType = VK_STRUCTURE_TYPE_RESOURCE_DESCRIPTOR_INFO_EXT,
                         .type = VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE,
                         .data = { .pImage = &textureInfos[i] } },
                    descriptorHeap.ImageByteOffset(imgSlots.firstElement + 1 + i), imageSize);
                add(f, { .sType = VK_STRUCTURE_TYPE_RESOURCE_DESCRIPTOR_INFO_EXT,
                         .type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
                         .data = { .pAddressRange = &colourRanges[i] } },
                    descriptorHeap.BufferByteOffset(bufSlots.firstElement + 1 + i), bufferSize);
            }
        }
        Device::vkWriteResourceDescriptorsEXT(Device::device, static_cast<std::uint32_t>(resources.size()),
                                              resources.data(), destinations.data());
        for (auto& heap : descriptorHeap.resourceHeap) heap.FlushDevice();
    };
    writeDescriptors();
    EventListener<void> resizeListener(&window.onResize, writeDescriptors);

    VkSamplerCreateInfo samplerInfo = {
        .sType = VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO,
        .magFilter = VK_FILTER_NEAREST,
        .minFilter = VK_FILTER_NEAREST,
        .mipmapMode = VK_SAMPLER_MIPMAP_MODE_NEAREST,
        .addressModeU = VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE,
        .addressModeV = VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE,
        .addressModeW = VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE,
        .maxLod = VK_LOD_CLAMP_NONE,
    };
    std::array<VkSamplerCreateInfo, Window::numFrames> samplerInfos{};
    std::array<VkHostAddressRangeEXT, Window::numFrames> samplerDestinations{};
    for (std::uint32_t f = 0; f < Window::numFrames; ++f) {
        samplerInfos[f] = samplerInfo;
        samplerDestinations[f] = {
            .address = descriptorHeap.samplerHeap[f].value + descriptorHeap.SamplerByteOffset(samplerSlots.firstElement),
            .size = Device::descriptorHeapProperties.samplerDescriptorSize,
        };
    }
    Device::vkWriteSamplerDescriptorsEXT(Device::device, Window::numFrames, samplerInfos.data(), samplerDestinations.data());
    for (auto& heap : descriptorHeap.samplerHeap) heap.FlushDevice();

    window.descriptorHeap = &descriptorHeap;
    RTPass rtPass(&pipeline);
    window.passes.push_back(&rtPass);

    // More frames than are in flight, so every per-frame heap copy is read at
    // least once. A device loss here throws out of Render().
    for (int i = 0; i < Window::numFrames * 2; ++i) window.Render();
    Device::CheckVkResult(vkQueueWaitIdle(Device::queue));

    // Counted before the capture: SaveFrame reads back an image that has
    // already been handed to the presentation engine, which the validation
    // layer reports on its own.
    if (std::getenv("CRAFTER_GRAPHICS_VALIDATION") != nullptr) {
        Check(Device::validationErrorCount == 0,
              std::format("no Vulkan validation errors while rendering ({} seen)", Device::validationErrorCount));
    }

    window.SaveFrame(out);
    Device::CheckVkResult(vkQueueWaitIdle(Device::queue));
    const Frame frame = LoadFrame(out);
    Check(!frame.rgba.empty(), "the frame was read back");
    if (!frame.rgba.empty()) {
        const int y = frame.height / 2;
        // An even column then the odd one next to it, in each triangle.
        const int left = FrameX(-100.0f, frame) & ~1;
        const int right = FrameX(100.0f, frame) & ~1;
        CheckPixel(frame, left, y, setColour[0], "instance 0, even column (set 0)");
        CheckPixel(frame, left + 1, y, setColour[0], "instance 0, odd column (set 0)");
        CheckPixel(frame, right, y, setColour[1], "instance 1, even column (set 1)");
        CheckPixel(frame, right + 1, y, setColour[0], "instance 1, odd column (set 0)");
        CheckPixel(frame, frame.width / 2, y, {0, 0, 255}, "miss between the triangles");
        CheckPixel(frame, frame.width / 20, frame.height / 10, {0, 0, 255}, "miss in the corner");
    }

    for (RenderingElement3D& renderer : renderers) RenderingElement3D::Remove(&renderer);
    return failures - failuresBefore;
}

// A fresh directory for the read-back frames, so a frame left by an earlier
// run can't stand in for one this run failed to write.
fs::path OutputDir() {
    const fs::path dir = fs::temp_directory_path() / "crafter-heap-dynamic-index-hit";
    std::error_code ec;
    fs::remove_all(dir, ec);
    fs::create_directories(dir);
    return dir;
}

} // namespace

int main(int argc, char** argv) {
    Device::Initialize();
    if (!Device::rayTracingSupported) {
        std::println("SKIP  no Vulkan ray-tracing device");
        return 77;
    }
    if (Device::driverProperties.driverID == VK_DRIVER_ID_NVIDIA_PROPRIETARY
        && NvidiaVersion() < std::tuple(615u, 78u, 8u)) {
        const auto [major, minor, patch] = NvidiaVersion();
        std::println("SKIP  NVIDIA {}.{:02}.{:02} loses the device on this (fixed in 615.78.08)", major, minor, patch);
        return 77;
    }

    // `--render <png>`: the child run, rendering as configured.
    if (argc == 3 && std::string_view(argv[1]) == "--render") {
        RenderScene(argv[2]);
        return failures == 0 ? EXIT_SUCCESS : EXIT_FAILURE;
    }

    const fs::path dir = OutputDir();
    if (std::getenv("CRAFTER_GRAPHICS_VALIDATION") != nullptr) {
        const fs::path exe = fs::read_symlink("/proc/self/exe");
        const int exit = RunCommand(std::format("env -u CRAFTER_GRAPHICS_VALIDATION '{}' --render '{}'",
                                                exe.string(), (dir / "unvalidated.png").string()));
        Check(exit == 0, std::format("the scene renders with validation off (exit {})", exit));
    }
    RenderScene(dir / "frame.png");

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