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

// Regression test for reading an acceleration structure out of a
// VK_EXT_descriptor_heap heap (#7, worked around on NVIDIA in #15). It renders
// examples/VulkanTriangle's triangle, tracing the left half of the frame with
// traceRayEXT and ray querying the right half, then reads the frame back and
// checks both halves hit the triangle where it is and miss around it. As run
// by the test runner it goes through the engine's normal path, so on NVIDIA it
// covers the SPIR-V rewrite too; a CRAFTER_GRAPHICS_NO_AS_WORKAROUND inherited
// from the environment is cleared so the rewrite can't be switched off.
//
// `HeapASRead --no-workaround` instead probes whether the driver still has the
// fault: it renders each path on its own in a child process with the rewrite
// off and validation off (GPU-assisted validation instruments the shaders),
// and reports per path whether it works, renders wrong or loses the device.
// It exits 0 only when both work. On 610.43.02 both paths lost the device; on
// 615.78.08 traceRayEXT still does, and the ray query renders nothing.
//
// The validation check only means something with CRAFTER_GRAPHICS_VALIDATION
// 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 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, float fx, float fy, std::array<std::uint8_t, 3> expected, std::string_view what) {
    const int x = static_cast<int>(fx * static_cast<float>(frame.width));
    const int y = static_cast<int>(fy * static_cast<float>(frame.height));
    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]));
}

// Matches `path` in heapas-raygen.glsl.
enum class Path : std::uint32_t { Both = 0, Trace = 1, Query = 2 };

// Spec constants of the raygen, matching heapas-raygen.glsl.
struct RaygenConstants {
    std::uint16_t bufferStart;
    std::uint16_t pad;
    std::uint32_t path;
};

constexpr std::array<std::uint8_t, 3> traceHit = {0, 255, 0};
constexpr std::array<std::uint8_t, 3> queryHit = {255, 255, 255};
constexpr std::array<std::uint8_t, 3> miss     = {128, 128, 128};

// Renders the triangle through `path` into `out` and checks the pixels.
// Returns the exit code.
int Render(Path path, const fs::path& out) {
    Device::Initialize();
    if (!Device::rayTracingSupported) {
        std::println("SKIP  no Vulkan ray-tracing device");
        return 77;
    }
    Device::validationErrorCount = 0;
    if (path != Path::Both) {
        Check(!Device::workaroundDescriptorHeapAS, "the AS workaround is off for the probe");
    }

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

    DescriptorHeapVulkan descriptorHeap;
    descriptorHeap.Initialize(1, 1, 0);

    RaygenConstants raygenConstants = {
        .bufferStart = descriptorHeap.bufferStartElement,
        .pad = 0,
        .path = static_cast<std::uint32_t>(path),
    };
    std::array<VkSpecializationMapEntry, 2> raygenEntries {{
        { .constantID = 0, .offset = offsetof(RaygenConstants, bufferStart), .size = sizeof(std::uint16_t) },
        { .constantID = 1, .offset = offsetof(RaygenConstants, path), .size = sizeof(std::uint32_t) },
    }};
    VkSpecializationInfo raygenSpec = {
        .mapEntryCount = static_cast<std::uint32_t>(raygenEntries.size()),
        .pMapEntries = raygenEntries.data(),
        .dataSize = sizeof(RaygenConstants),
        .pData = &raygenConstants,
    };

    const fs::path dir = ExeDir();
    std::array<VulkanShader, 3> shaders{{
        {dir / "heapas-raygen.spv", "main", VK_SHADER_STAGE_RAYGEN_BIT_KHR, &raygenSpec},
        {dir / "heapas-miss.spv", "main", VK_SHADER_STAGE_MISS_BIT_KHR, nullptr},
        {dir / "heapas-closesthit.spv", "main", VK_SHADER_STAGE_CLOSEST_HIT_BIT_KHR, nullptr},
    }};
    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);

    Mesh triangleMesh;
    std::array<Vector<float, 3, 3>, 3> verts {{{-150, -150, 100}, {0, 150, 100}, {150, -150, 100}}};
    std::array<std::uint32_t, 3> index {{2, 1, 0}};
    triangleMesh.Build(verts, index, cmd);

    RenderingElement3D renderer;
    renderer.instance = {
        .instanceCustomIndex                    = 0,
        .mask                                   = 0xFF,
        .instanceShaderBindingTableRecordOffset = 0,
        .flags                                  = VK_GEOMETRY_INSTANCE_FORCE_OPAQUE_BIT_KHR,
        .accelerationStructureReference         = triangleMesh.blasAddr,
    };
    MatrixRowMajor<float, 4, 3, 1>::Identity().Store(reinterpret_cast<float*>(renderer.instance.transform.matrix));
    RenderingElement3D::Add(&renderer);
    for (std::uint32_t f = 0; f < Window::numFrames; ++f) RenderingElement3D::BuildTLAS(cmd, f);

    window.FinishInit();

    // Per frame: the TLAS at buffer slot 0 and that frame's swapchain image at
    // image slot 0, each in its own copy of the heap. Written again whenever
    // the swapchain is recreated, since the image descriptors point at its
    // images.
    auto imgSlots = descriptorHeap.AllocateImageSlots(1);
    auto bufSlots = descriptorHeap.AllocateBufferSlots(1);
    auto writeDescriptors = [&] {
        std::array<VkDeviceAddressRangeKHR, Window::numFrames> tlasRanges{};
        std::array<VkImageDescriptorInfoEXT, Window::numFrames> imageInfos{};
        std::vector<VkResourceDescriptorInfoEXT> resources;
        std::vector<VkHostAddressRangeEXT> destinations;
        for (std::uint32_t f = 0; f < Window::numFrames; ++f) {
            tlasRanges[f] = { .address = RenderingElement3D::tlases[f].address };
            imageInfos[f] = {
                .sType = VK_STRUCTURE_TYPE_IMAGE_DESCRIPTOR_INFO_EXT,
                .pView = &window.imageViews[f],
                .layout = VK_IMAGE_LAYOUT_GENERAL,
            };
            resources.push_back({
                .sType = VK_STRUCTURE_TYPE_RESOURCE_DESCRIPTOR_INFO_EXT,
                .type = VK_DESCRIPTOR_TYPE_ACCELERATION_STRUCTURE_KHR,
                .data = { .pAddressRange = &tlasRanges[f] },
            });
            destinations.push_back({
                .address = descriptorHeap.resourceHeap[f].value + descriptorHeap.BufferByteOffset(bufSlots.firstElement),
                .size = Device::descriptorHeapProperties.bufferDescriptorSize,
            });
            resources.push_back({
                .sType = VK_STRUCTURE_TYPE_RESOURCE_DESCRIPTOR_INFO_EXT,
                .type = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
                .data = { .pImage = &imageInfos[f] },
            });
            destinations.push_back({
                .address = descriptorHeap.resourceHeap[f].value + descriptorHeap.ImageByteOffset(imgSlots.firstElement),
                .size = Device::descriptorHeapProperties.imageDescriptorSize,
            });
        }
        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);

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

    // More frames than are in flight, so every per-frame heap copy and TLAS
    // 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()) {
        switch (path) {
            case Path::Both:
                // The triangle covers the middle of the frame at any aspect
                // ratio; 5% either side of the centre line is inside it.
                CheckPixel(frame, 0.45f, 0.5f, traceHit, "traceRayEXT hit");
                CheckPixel(frame, 0.55f, 0.5f, queryHit, "rayQueryEXT hit");
                CheckPixel(frame, 0.05f, 0.1f, miss, "traceRayEXT miss");
                CheckPixel(frame, 0.95f, 0.1f, miss, "rayQueryEXT miss");
                break;
            case Path::Trace:
                CheckPixel(frame, 0.5f, 0.5f, traceHit, "traceRayEXT hit");
                CheckPixel(frame, 0.05f, 0.1f, miss, "traceRayEXT miss");
                break;
            case Path::Query:
                CheckPixel(frame, 0.5f, 0.5f, queryHit, "rayQueryEXT hit");
                CheckPixel(frame, 0.05f, 0.1f, miss, "rayQueryEXT miss");
                break;
        }
    }
    RenderingElement3D::Remove(&renderer);
    return failures == 0 ? EXIT_SUCCESS : EXIT_FAILURE;
}

// 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-as-read";
    std::error_code ec;
    fs::remove_all(dir, ec);
    fs::create_directories(dir);
    return dir;
}

// --no-workaround: renders each path in its own child process, so a device
// loss on one doesn't hide the other.
int Probe() {
    const fs::path exe = fs::read_symlink("/proc/self/exe");
    const fs::path dir = OutputDir();
    bool allWork = true;
    for (const auto& [name, arg] : {std::pair{"traceRayEXT", "trace"}, std::pair{"rayQueryEXT", "query"}}) {
        const int exit = RunCommand(std::format(
            "env -u CRAFTER_GRAPHICS_VALIDATION CRAFTER_GRAPHICS_NO_AS_WORKAROUND=1 '{}' --probe {} '{}'",
            exe.string(), arg, (dir / std::format("{}.png", arg)).string()));
        const std::string verdict =
            exit == 0  ? std::string("works") :
            exit == 77 ? std::string("skipped (no RT device)") :
            exit == 1  ? std::string("renders wrong") :
            std::format("loses the device or crashes (exit {})", exit);
        std::println("PROBE  {} reading the TLAS from the heap: {}", name, verdict);
        if (exit != 0) allWork = false;
    }
    std::println("{}", allWork ? "the driver reads the TLAS from the heap on both paths"
                               : "the driver still needs the workaround");
    return allWork ? EXIT_SUCCESS : EXIT_FAILURE;
}

} // namespace

int main(int argc, char** argv) {
    if (argc == 2 && std::string_view(argv[1]) == "--no-workaround") {
        return Probe();
    }
    if (argc == 4 && std::string_view(argv[1]) == "--probe") {
        const Path path = std::string_view(argv[2]) == "trace" ? Path::Trace : Path::Query;
        return Render(path, argv[3]);
    }

    unsetenv("CRAFTER_GRAPHICS_NO_AS_WORKAROUND");
    const int result = Render(Path::Both, OutputDir() / "both.png");
    if (result == 77) return result;
    Check(Device::workaroundDescriptorHeapAS
              == (Device::descriptorHeapSupported && Device::driverProperties.driverID == VK_DRIVER_ID_NVIDIA_PROPRIETARY),
          "the AS workaround is on exactly where the driver needs it");
    std::println("{}", failures == 0 ? "all checks passed" : "check(s) failed");
    return failures == 0 ? EXIT_SUCCESS : EXIT_FAILURE;
}