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

// RTPass pushes the application's own push data (pushPtr/pushBytes, as the
// DOM backend's RTPass takes it) and traces at an explicit size instead of
// always the window's. Before, an app with a camera to push had to copy
// RTPass's Record, NVIDIA TLAS push included, into its own pass. The frame
// is cleared to magenta, then the pass traces only the top-left quadrant,
// colouring hits from a push block that changes every frame; the read-back
// frame shows the last frame's colour on the triangle, the miss colour
// around it inside the quadrant, and the clear colour outside it. On NVIDIA
// the TLAS address is pushed too, after the app's block, so this also covers
// the two pushes together.
//
// Needs a live Wayland compositor and a Vulkan RT device; without RT it
// skips. The validation check only means something with
// CRAFTER_GRAPHICS_VALIDATION set.

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

// PNG decode for the read-back frame. The library only compiles the writer
// (Window::CaptureNextFrame); 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);
    if (!ok) ++failures;
}

fs::path ExeDir() {
    return fs::read_symlink("/proc/self/exe").parent_path();
}

struct ClearPass : RenderPass {
    void Record(VkCommandBuffer cmd, std::uint32_t, Window& window) override {
        const VkClearColorValue magenta = {{1.0f, 0.0f, 1.0f, 1.0f}};
        const VkImageSubresourceRange range = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1};
        vkCmdClearColorImage(cmd, window.images[window.currentBuffer], VK_IMAGE_LAYOUT_GENERAL, &magenta, 1, &range);
    }
    VkPipelineStageFlags SwapchainStage() const override { return VK_PIPELINE_STAGE_TRANSFER_BIT; }
};

// Matches the raygen's push block.
struct Push {
    float hitColour[4];
};

} // namespace

int main() {
    Device::Initialize();
    if (!Device::rayTracingSupported) {
        std::println("SKIP  no Vulkan ray-tracing device");
        return 77;
    }
    Device::validationErrorCount = 0;

    Window window(640, 360, "RTPassPushData test");
    if ((window.swapchainUsage & VK_IMAGE_USAGE_TRANSFER_DST_BIT) == 0) {
        std::println("SKIP  the surface can't clear the swapchain image");
        return 77;
    }
    VkCommandBuffer cmd = window.StartInit();

    DescriptorHeapVulkan descriptorHeap;
    descriptorHeap.Initialize(1, 1, 0);
    VkSpecializationMapEntry entry = { .constantID = 0, .offset = 0, .size = sizeof(std::uint16_t) };
    VkSpecializationInfo spec = { .mapEntryCount = 1, .pMapEntries = &entry, .dataSize = sizeof(std::uint16_t),
                                  .pData = &descriptorHeap.bufferStartElement };
    const fs::path dir = ExeDir();
    std::array<VulkanShader, 3> shaders{{
        {dir / "rtpass-raygen.spv", "main", VK_SHADER_STAGE_RAYGEN_BIT_KHR, &spec},
        {dir / "rtpass-miss.spv", "main", VK_SHADER_STAGE_MISS_BIT_KHR, nullptr},
        {dir / "rtpass-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();

    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;

    ClearPass clear;
    Push push{};
    RTPass rtPass(&pipeline);
    rtPass.pushPtr = &push;
    rtPass.pushBytes = sizeof(push);
    window.passes.push_back(&clear);
    window.passes.push_back(&rtPass);

    // A different colour every frame; the frame read back has the last one.
    const int frames = Window::numFrames * 2;
    for (int i = 0; i < frames; ++i) {
        rtPass.width = window.width / 2;
        rtPass.height = window.height / 2;
        push = { { 0.0f, static_cast<float>(i) / static_cast<float>(frames - 1), 1.0f, 1.0f } };
        window.Render();
    }
    Device::CheckVkResult(vkQueueWaitIdle(Device::queue));

    const fs::path out = fs::temp_directory_path() / "crafter-rtpass-push-data.png";
    std::error_code ec;
    fs::remove(out, ec);
    window.CaptureNextFrame(out);
    window.Render();
    if (std::getenv("CRAFTER_GRAPHICS_VALIDATION") != nullptr) {
        Check(Device::validationErrorCount == 0,
              std::format("no Vulkan validation errors while rendering and capturing ({} seen)", Device::validationErrorCount));
    }
    Device::CheckVkResult(vkQueueWaitIdle(Device::queue));
    int width = 0, height = 0, channels = 0;
    stbi_uc* pixels = stbi_load(out.string().c_str(), &width, &height, &channels, STBI_rgb_alpha);
    Check(pixels != nullptr, "the frame was read back");
    if (pixels != nullptr) {
        auto at = [&](int x, int y) {
            const stbi_uc* p = pixels + (static_cast<std::size_t>(y) * width + x) * 4;
            return std::array<int, 3>{p[0], p[1], p[2]};
        };
        auto near = [](std::array<int, 3> got, std::array<int, 3> want) {
            for (int c = 0; c < 3; ++c) if (std::abs(got[c] - want[c]) > 3) return false;
            return true;
        };
        const auto inside = at(width / 4, height / 4);
        const auto quadrantCorner = at(width / 40, height / 40);
        const auto outside = at(3 * width / 4, 3 * height / 4);
        Check(near(inside, {0, 255, 255}), std::format("the triangle has the last frame's pushed colour (got {}, {}, {})", inside[0], inside[1], inside[2]));
        Check(near(quadrantCorner, {0, 0, 64}), std::format("misses inside the traced quadrant (got {}, {}, {})", quadrantCorner[0], quadrantCorner[1], quadrantCorner[2]));
        Check(near(outside, {255, 0, 255}), std::format("outside the trace size the clear stays (got {}, {}, {})", outside[0], outside[1], outside[2]));
        stbi_image_free(pixels);
    }

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