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

// The UI path has two binding models. Where the driver has
// VK_EXT_descriptor_heap, shaders index the heap directly; where it does not
// (Mesa Turnip on Adreno), the same slots live in a bindless descriptor set
// and the shaders are a second build of the same source (<name>.sets.spv,
// CRAFTER_GRAPHICS_DESCRIPTOR_SETS). Nothing but the binding model may differ
// between the two, so this test renders one fixed scene through each and
// compares the frames pixel by pixel.
//
// The scene is chosen to reach everything the descriptor-set path implements
// and HelloUI does not exercise:
//
//   - the four standalone kernels (DispatchQuads / Circles / Images / Text)
//     and the fused kernel with all four categories non-empty;
//   - a sampled image and two samplers indexed per item (bindings 1 and 2,
//     nonuniform indices), next to the font atlas;
//   - slot reuse: an image slot and a sampler slot are each freed and
//     registered again before the first frame, so the second write has to
//     replace the first one's view / sampler. A stale descriptor shows the
//     wrong texture and fails the comparison;
//   - a resize, which rewrites the per-frame swapchain descriptors;
//   - an image created with ImageVulkan::Create's default consumer stage,
//     which has to be resolved to a stage the device actually has;
//   - a user-style ui.Dispatch whose push payload is not a whole number of
//     words, which the descriptor-set path pads (the heap run pushes the
//     whole header and must draw the same quad).
//
// How it runs: with no arguments the test is the driver. It starts itself
// twice with `--render <png>`, once as is and once with
// CRAFTER_GRAPHICS_NO_DESCRIPTOR_HEAP set, then decodes and compares the two
// PNGs. A child also fails on any Vulkan validation error raised while
// rendering, which only means something when CRAFTER_GRAPHICS_VALIDATION is
// set for the run.
//
// On a driver that has no descriptor heap at all, both children take the
// descriptor-set path. There is then nothing to compare against, and the test
// passes on the strength of the scene rendering cleanly — which is the run
// that matters on such a device. `--render <png>` can also be used directly to
// produce a frame for inspection.
//
// Needs a live Wayland compositor and a Vulkan device, like FrameLoopSync.

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

// PNG decode for the comparison. 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 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;
}

// The test binary's own directory: the build places the compiled shaders
// (both variants) and font.ttf next to it, 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);
}

// RGBA8 texel as a 4-byte word, matching VK_FORMAT_R8G8B8A8_UNORM.
using Texel = std::uint32_t;
constexpr std::uint16_t kTexSize = 8;

constexpr std::uint32_t kWidth  = 640;
constexpr std::uint32_t kHeight = 480;
// The size after the mid-run resize. Smaller than the first one so the scene,
// laid out in absolute pixels, still fits.
constexpr std::uint32_t kResizedWidth  = 600;
constexpr std::uint32_t kResizedHeight = 440;

template<typename T>
void CreateItemBuffer(GraphicsBuffer<T, true>& buffer, std::uint32_t count) {
    buffer.Create(
        VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT,
        VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, count);
}

void CreateTexture(ImageVulkan<Texel>& image, VkCommandBuffer cmd, auto&& texelAt) {
    image.Create(kTexSize, kTexSize, /*mipLevels*/ 1, cmd,
                 VK_FORMAT_R8G8B8A8_UNORM,
                 VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_SAMPLED_BIT,
                 VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
                 VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT);
    for (std::uint32_t y = 0; y < kTexSize; ++y) {
        for (std::uint32_t x = 0; x < kTexSize; ++x) {
            image.buffer.value[y * kTexSize + x] = texelAt(x, y);
        }
    }
    image.Update(cmd, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);
}

VkSamplerCreateInfo SamplerInfo(VkFilter filter) {
    return {
        .sType        = VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO,
        .magFilter    = filter,
        .minFilter    = filter,
        .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,
        .maxAnisotropy = 1.0f,
        .maxLod       = VK_LOD_CLAMP_NONE,
    };
}

// Renders the scene and writes the last frame to `out`. Also writes
// `<out>.mode` holding "heap" or "sets", so the driver can tell which binding
// model this process actually ran with.
int RenderScene(const fs::path& out) {
    const fs::path dir = ExeDir();

    Device::Initialize();
    Window window(kWidth, kHeight, "UIFallbackParity test");
    VkCommandBuffer init = window.StartInit();

    GraphicsDescriptorHeap heap;
    heap.Initialize(/*images*/ 8, /*buffers*/ 16, /*samplers*/ 4);
    window.descriptorHeap = &heap;

    Font font(dir / "font.ttf");
    FontAtlas atlas;
    atlas.Initialize(init);

    UIRenderer ui;
    ui.fontAtlas = &atlas;
    ui.Initialize(window, heap, init,
                  dir / "ui-quads.comp.spv", dir / "ui-circles.comp.spv",
                  dir / "ui-images.comp.spv", dir / "ui-text.comp.spv",
                  dir / "ui-fused.comp.spv");
    window.passes.push_back(&ui);

    // Two textures. `red` is registered first and its slot freed again, so
    // `checker` lands on the same slot: its descriptor (in sets mode, its
    // view) must replace red's. If the write were lost the image quads below
    // would come out red in one model and checkered in the other.
    ImageVulkan<Texel> red;
    ImageVulkan<Texel> checker;
    CreateTexture(red, init, [](std::uint32_t, std::uint32_t) { return 0xFF0000FFu; });
    CreateTexture(checker, init, [](std::uint32_t x, std::uint32_t y) {
        return ((x + y) & 1u) != 0 ? 0xFFFFC040u : 0xFF204080u;
    });

    // A third image, created with Create's default consumer stage and never
    // sampled. That default names the ray tracing stage, which is not a valid
    // barrier stage on a device without ray tracing, so Create has to resolve
    // it for both the member and the transition it records itself. Only the
    // validation layer can see a wrong stage here.
    ImageVulkan<Texel> defaultStage;
    defaultStage.Create(kTexSize, kTexSize, /*mipLevels*/ 1, init,
                        VK_FORMAT_R8G8B8A8_UNORM,
                        VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_SAMPLED_BIT,
                        VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);
    std::fill_n(defaultStage.buffer.value, kTexSize * kTexSize, Texel{0});
    defaultStage.Update(init, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);

    ImageSlot imageSlot = ui.RegisterImage(red, VK_FORMAT_R8G8B8A8_UNORM);
    const std::uint16_t firstImageSlot = imageSlot;
    imageSlot.Reset();
    imageSlot = ui.RegisterImage(checker, VK_FORMAT_R8G8B8A8_UNORM);
    Check(static_cast<std::uint16_t>(imageSlot) == firstImageSlot,
          "the freed image slot is handed out again");

    // Same for a sampler slot: linear first, then nearest on the reused slot.
    // The 8x8 checker is magnified ~12x, so the two filters differ in almost
    // every pixel of the quad.
    SamplerSlot nearestSlot = ui.RegisterSampler(SamplerInfo(VK_FILTER_LINEAR));
    const std::uint16_t firstSamplerSlot = nearestSlot;
    nearestSlot.Reset();
    nearestSlot = ui.RegisterSampler(SamplerInfo(VK_FILTER_NEAREST));
    Check(static_cast<std::uint16_t>(nearestSlot) == firstSamplerSlot,
          "the freed sampler slot is handed out again");
    SamplerSlot linearSlot = ui.RegisterLinearClampSampler();

    // One set of item buffers for the standalone dispatches (left column) and
    // one for the fused dispatch (right column).
    GraphicsBuffer<QuadItem,   true> quadsA, quadsB;
    GraphicsBuffer<CircleItem, true> circlesA, circlesB;
    GraphicsBuffer<ImageItem,  true> imagesA, imagesB;
    GraphicsBuffer<GlyphItem,  true> glyphsA, glyphsB;
    CreateItemBuffer(quadsA, 8);     CreateItemBuffer(quadsB, 8);
    CreateItemBuffer(circlesA, 8);   CreateItemBuffer(circlesB, 8);
    CreateItemBuffer(imagesA, 8);    CreateItemBuffer(imagesB, 8);
    CreateItemBuffer(glyphsA, 64);   CreateItemBuffer(glyphsB, 64);
    BufferSlot quadsASlot   = ui.RegisterBuffer(quadsA);
    BufferSlot circlesASlot = ui.RegisterBuffer(circlesA);
    BufferSlot imagesASlot  = ui.RegisterBuffer(imagesA);
    BufferSlot glyphsASlot  = ui.RegisterBuffer(glyphsA);
    BufferSlot quadsBSlot   = ui.RegisterBuffer(quadsB);
    BufferSlot circlesBSlot = ui.RegisterBuffer(circlesB);
    BufferSlot imagesBSlot  = ui.RegisterBuffer(imagesB);
    BufferSlot glyphsBSlot  = ui.RegisterBuffer(glyphsB);

    // A single quad for the odd-sized generic dispatch below.
    GraphicsBuffer<QuadItem, true> quadsOdd;
    CreateItemBuffer(quadsOdd, 1);
    BufferSlot quadsOddSlot = ui.RegisterBuffer(quadsOdd);

    // Everything is laid out in absolute pixels so both runs draw the same
    // thing whatever size the compositor settles on.
    auto fillColumn = [&](float ox, GraphicsBuffer<QuadItem, true>& quads,
                          GraphicsBuffer<CircleItem, true>& circles,
                          GraphicsBuffer<ImageItem, true>& images,
                          GraphicsBuffer<GlyphItem, true>& glyphs,
                          bool withBackground, std::string_view label,
                          std::array<std::uint32_t, 4>& counts) {
        std::uint32_t q = 0;
        if (withBackground) {
            quads.value[q++] = QuadItem{ 0, 0, 4096, 4096,  0.08f, 0.09f, 0.12f, 1.0f,  0, 0, 0, 0,  0, 0, 0, 0 };
        }
        quads.value[q++] = QuadItem{ ox, 20, 220, 70,  0.20f, 0.45f, 0.85f, 1.0f,  4, 12, 20, 28,  3, 1.0f, 0.8f, 0.2f };
        quads.value[q++] = QuadItem{ ox + 40, 60, 140, 50,  0.90f, 0.30f, 0.30f, 0.6f,  0, 0, 0, 0,  0, 0, 0, 0 };

        std::uint32_t c = 0;
        circles.value[c++] = CircleItem{ ox + 50, 150, 34, 0,  0.30f, 0.85f, 0.45f, 1.0f,  0, 0, 0, 0 };
        circles.value[c++] = CircleItem{ ox + 150, 150, 30, 0,  0.0f, 0.0f, 0.0f, 0.0f,  5, 0.95f, 0.95f, 1.0f };

        std::uint32_t i = 0;
        images.value[i++] = ImageItem{ ox, 200, 96, 96,  0, 0, 1, 1,  1, 1, 1, 1,
                                       imageSlot, nearestSlot, 0, 0 };
        images.value[i++] = ImageItem{ ox + 110, 200, 96, 96,  0, 0, 1, 1,  1.0f, 0.8f, 0.8f, 1.0f,
                                       imageSlot, linearSlot, 0, 0 };

        std::uint32_t g = ui.ShapeText(font, 22.0f, ox, 340.0f, label,
                                       {1.0f, 1.0f, 1.0f, 1.0f}, glyphs.value, 64);
        counts = { q, c, i, g };
    };

    std::array<std::uint32_t, 4> countsA{};
    std::array<std::uint32_t, 4> countsB{};
    EventListener<UIBuildArgs> buildSub(&ui.onBuild, [&](UIBuildArgs a) {
        VkCommandBuffer cmd = a.cmd;
        fillColumn(30.0f,  quadsA, circlesA, imagesA, glyphsA, /*withBackground*/ true,  "standalone", countsA);
        fillColumn(330.0f, quadsB, circlesB, imagesB, glyphsB, /*withBackground*/ false, "fused",      countsB);

        constexpr VkAccessFlags        read  = VK_ACCESS_SHADER_READ_BIT;
        constexpr VkPipelineStageFlags stage = VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT;
        quadsA.FlushDevice(cmd, read, stage);   quadsB.FlushDevice(cmd, read, stage);
        circlesA.FlushDevice(cmd, read, stage); circlesB.FlushDevice(cmd, read, stage);
        imagesA.FlushDevice(cmd, read, stage);  imagesB.FlushDevice(cmd, read, stage);
        glyphsA.FlushDevice(cmd, read, stage);  glyphsB.FlushDevice(cmd, read, stage);

        ui.DispatchQuads(cmd,   quadsASlot,   countsA[0]);
        ui.DispatchCircles(cmd, circlesASlot, countsA[1]);
        ui.DispatchImages(cmd,  imagesASlot,  countsA[2]);
        ui.DispatchText(cmd,    glyphsASlot,  countsA[3]);

        ui.DispatchFused(cmd,
                         {quadsBSlot,   countsB[0]},
                         {circlesBSlot, countsB[1]},
                         {imagesBSlot,  countsB[2]},
                         {glyphsBSlot,  countsB[3]});

        // One more quad through the generic Dispatch a custom shader would
        // use. On the descriptor-set path the last three bytes of the
        // header's reserved `_pad` are left off, so ComputeShader::Dispatch
        // has to pad the payload for vkCmdPushConstants. The heap run pushes
        // the whole header (a ragged size is a validation error there and
        // nothing pads it), and both must draw the same quad.
        quadsOdd.value[0] = QuadItem{ 30, 370, 180, 30,  0.85f, 0.65f, 0.15f, 1.0f,  6, 6, 6, 6,  0, 0, 0, 0 };
        quadsOdd.FlushDevice(cmd, read, stage);
        const UIDispatchHeader oddHeader = ui.FillHeader(quadsOddSlot, 1);
        const std::uint32_t oddBytes = Device::descriptorHeapSupported
            ? static_cast<std::uint32_t>(sizeof(oddHeader))
            : static_cast<std::uint32_t>(sizeof(oddHeader) - 3);
        ui.Dispatch(cmd, ui.drawQuads, &oddHeader, oddBytes,
                    (window.width + 7) / 8, (window.height + 7) / 8);
    });

    window.FinishInit();

    // More frames than there are in flight, so every per-frame heap copy /
    // descriptor set is used; then a resize, which recreates the swapchain
    // and rewrites its descriptors; then enough frames for the glyphs shaped
    // on the first frame to have reached the atlas texture.
    for (int i = 0; i < Window::numFrames * 2; ++i) window.Render();
    window.Resize(kResizedWidth, kResizedHeight);
    for (int i = 0; i < Window::numFrames * 3; ++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 in either binding model. What is under test here is the
    // rendering above.
    const std::uint32_t renderErrors = Device::validationErrorCount;
    Check(renderErrors == 0,
          std::format("no Vulkan validation errors while rendering ({} seen)", renderErrors));

    window.SaveFrame(out);
    std::ofstream(fs::path(out).concat(".mode")) << (Device::descriptorHeapSupported ? "heap" : "sets");
    Check(fs::exists(out), "the frame was written");

    // The buffers, images and heap go out of scope here while the last
    // frames may still be presenting; settle the queue first.
    Device::CheckVkResult(vkQueueWaitIdle(Device::queue));
    return failures == 0 ? EXIT_SUCCESS : EXIT_FAILURE;
}

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

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

std::string ReadMode(const fs::path& png) {
    std::ifstream in(fs::path(png).concat(".mode"));
    std::string mode;
    in >> mode;
    return mode;
}

int Drive() {
    const fs::path exe = fs::read_symlink("/proc/self/exe");
    const fs::path dir = fs::temp_directory_path() / "crafter-ui-fallback-parity";
    std::error_code ec;
    fs::remove_all(dir, ec);
    fs::create_directories(dir);
    const fs::path heapPng = dir / "heap.png";
    const fs::path setsPng = dir / "sets.png";

    // `env -u` so a CRAFTER_GRAPHICS_NO_DESCRIPTOR_HEAP inherited from the
    // caller's environment cannot turn the reference run into a second
    // fallback run.
    const int heapExit = RunCommand(std::format(
        "env -u CRAFTER_GRAPHICS_NO_DESCRIPTOR_HEAP '{}' --render '{}'", exe.string(), heapPng.string()));
    const int setsExit = RunCommand(std::format(
        "env CRAFTER_GRAPHICS_NO_DESCRIPTOR_HEAP=1 '{}' --render '{}'", exe.string(), setsPng.string()));

    Check(heapExit == 0, std::format("reference run rendered the scene (exit {})", heapExit));
    Check(setsExit == 0, std::format("descriptor-set run rendered the scene (exit {})", setsExit));
    Check(ReadMode(setsPng) == "sets", "the forced run took the descriptor-set path");
    if (failures != 0) return EXIT_FAILURE;

    const Frame sets = LoadFrame(setsPng);
    Check(!sets.rgba.empty(), "descriptor-set frame decodes");
    if (failures != 0) return EXIT_FAILURE;

    // A frame that is one flat colour means nothing was composited, which
    // two equally broken runs would agree on.
    std::set<std::uint32_t> colours;
    for (std::size_t i = 0; i + 3 < sets.rgba.size() && colours.size() < 16; i += 4) {
        std::uint32_t texel = 0;
        std::memcpy(&texel, &sets.rgba[i], 4);
        colours.insert(texel);
    }
    Check(colours.size() >= 16, std::format("descriptor-set frame has content ({}+ distinct colours)", colours.size()));

    if (ReadMode(heapPng) != "heap") {
        std::println("NOTE  this driver has no VK_EXT_descriptor_heap: both runs used descriptor sets, "
                     "so there is no heap frame to compare against");
        return failures == 0 ? EXIT_SUCCESS : EXIT_FAILURE;
    }

    const Frame heap = LoadFrame(heapPng);
    Check(!heap.rgba.empty(), "descriptor-heap frame decodes");
    Check(heap.width == sets.width && heap.height == sets.height,
          std::format("both frames are the same size ({}x{} vs {}x{})",
                      heap.width, heap.height, sets.width, sets.height));
    if (failures != 0) return EXIT_FAILURE;

    // One code value of slack per channel: both shaders are the same source,
    // but the driver compiles two different SPIR-V modules and may contract
    // or reorder float math differently on an anti-aliased edge.
    std::size_t differing = 0;
    int maxDiff = 0;
    for (std::size_t i = 0; i < heap.rgba.size(); ++i) {
        const int diff = std::abs(static_cast<int>(heap.rgba[i]) - static_cast<int>(sets.rgba[i]));
        if (diff != 0) ++differing;
        maxDiff = std::max(maxDiff, diff);
    }
    Check(maxDiff <= 1,
          std::format("heap and descriptor-set frames match (max channel difference {}, {} of {} channel values differ)",
                      maxDiff, differing, heap.rgba.size()));
    if (maxDiff > 1) {
        std::println("      frames kept for inspection: {} and {}", heapPng.string(), setsPng.string());
    }
    return failures == 0 ? EXIT_SUCCESS : EXIT_FAILURE;
}

} // namespace

int main(int argc, char** argv) {
    if (argc == 3 && std::string_view(argv[1]) == "--render") {
        return RenderScene(argv[2]);
    }
    const int result = Drive();
    std::println("{}", result == EXIT_SUCCESS ? "all checks passed" : "check(s) failed");
    return result;
}