/*
Crafter®.Graphics
Copyright (C) 2026 Catcrafts®
catcrafts.net

This library is free software; you can redistribute it and/or
modify it under the terms of the GNU Lesser General Public
License version 3.0 as published by the Free Software Foundation;

This library is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
Lesser General Public License for more details.

You should have received a copy of the GNU Lesser General Public
License along with this library; if not, write to the Free Software
Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301 USA
*/

// Regression test for issue #18: the NVIDIA descriptor-heap AS-read workaround
// (WorkaroundNvidiaAS::Patch) used to bolt a brand-new push-constant block onto
// every patched ray-tracing shader. SPIR-V allows at most one push-constant
// block statically used per entry point, so any shader that already declared
// one ended up with two and failed spirv-val:
//
//   Entry point id '4' uses more than one PushConstant interface.
//
// This test compiles representative ray-generation shaders with glslang, runs
// them through the real Patch(), and asserts with spirv-val that the result is
// valid and contains exactly one push-constant variable — both for shaders
// that already have a push constant (merge path) and for those that don't
// (synthesize path). It also checks the published TLAS push-constant offset.
//
// Delete this test together with the rest of the workaround once a fixed NVIDIA
// driver ships.

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

import Crafter.Graphics;
import std;
using namespace Crafter;

namespace {

namespace fs = std::filesystem;

int RunCommand(const std::string& cmd) {
    int status = std::system(cmd.c_str());
    if (status == -1) return -1;
    // Mirror WEXITSTATUS without pulling in <sys/wait.h>: glibc encodes the
    // exit code in bits 8..15 of the wait status when the low byte is zero.
    return (status & 0x7f) == 0 ? ((status >> 8) & 0xff) : 128 + (status & 0x7f);
}

std::vector<std::uint32_t> ReadSpirv(const fs::path& p) {
    std::ifstream f(p, std::ios::binary | std::ios::ate);
    if (!f) return {};
    std::streamsize size = f.tellg();
    f.seekg(0);
    std::vector<std::uint32_t> words(static_cast<std::size_t>(size) / sizeof(std::uint32_t));
    f.read(reinterpret_cast<char*>(words.data()), size);
    return words;
}

void WriteSpirv(const fs::path& p, const std::vector<std::uint32_t>& words) {
    std::ofstream f(p, std::ios::binary);
    f.write(reinterpret_cast<const char*>(words.data()),
            static_cast<std::streamsize>(words.size() * sizeof(std::uint32_t)));
}

// Count OpVariable instructions in the PushConstant storage class (SC == 9).
int CountPushConstantVariables(const std::vector<std::uint32_t>& words) {
    constexpr std::uint32_t OpVariable = 59;
    constexpr std::uint32_t StorageClassPushConstant = 9;
    int count = 0;
    for (std::size_t i = 5; i < words.size();) {
        std::uint32_t len = words[i] >> 16;
        if (len == 0 || i + len > words.size()) break;
        if ((words[i] & 0xFFFFu) == OpVariable && len >= 4 && words[i + 3] == StorageClassPushConstant)
            ++count;
        i += len;
    }
    return count;
}

struct Case {
    std::string_view name;
    std::string_view glsl;
    bool readsAccelStruct;     // whether Patch should rewrite anything
    bool hasExistingPushConst; // whether the source already declares a push block
    std::uint32_t expectedOffset; // expected Device::workaroundTlasPushOffset (only checked when readsAccelStruct)
};

// Shared raygen scaffolding: a heap AS + heap image, traced and stored to.
constexpr std::string_view kHeader =
    "#version 460\n"
    "#extension GL_EXT_ray_tracing : enable\n"
    "#extension GL_EXT_shader_image_load_formatted : enable\n"
    "#extension GL_EXT_descriptor_heap : enable\n"
    "#extension GL_EXT_nonuniform_qualifier : enable\n"
    "layout(descriptor_heap) uniform accelerationStructureEXT topLevelAS[];\n"
    "layout(descriptor_heap) uniform writeonly image2D image[];\n"
    "layout(location = 0) rayPayloadEXT vec3 hitValue;\n";

const std::array<Case, 5> kCases = {{
    // No push constant at all → Patch synthesizes a fresh single-member block at offset 0.
    { "no-push-constant", std::string_view{
        ""
      }, true, false, 0 },

    // Existing block {mat4 @0, vec3 @64, uint @76}; ends at 80, already 8-aligned.
    { "merge-mat4-vec3-uint", std::string_view{
        "layout(push_constant) uniform PC { mat4 m; vec3 l; uint f; } pc;\n"
      }, true, true, 80 },

    // Existing block {uint @0}; ends at 4, TLAS rounds up to the next 8.
    { "merge-uint", std::string_view{
        "layout(push_constant) uniform PC { uint f; } pc;\n"
      }, true, true, 8 },

    // Existing block {vec4 v[2] @0 (32 bytes), uint @32}; ends at 36, rounds to 40.
    { "merge-array", std::string_view{
        "layout(push_constant) uniform PC { vec4 v[2]; uint f; } pc;\n"
      }, true, true, 40 },

    // Push constant but NO acceleration-structure read → Patch is a no-op; the
    // single user block must survive untouched and still validate.
    { "push-constant-no-as", std::string_view{
        "layout(push_constant) uniform PC { vec4 tint; } pc;\n"
      }, false, true, 0 },
}};

std::string BuildSource(const Case& c) {
    std::string s(kHeader);
    s += c.glsl;
    s += "void main() {\n";
    s += "  uvec2 pixel = gl_LaunchIDEXT.xy;\n";
    s += "  vec3 origin = vec3(0.0, 0.0, -300.0);\n";
    s += "  vec3 dir = normalize(vec3(0.0, 0.0, 1.0));\n";
    if (c.readsAccelStruct)
        s += "  traceRayEXT(topLevelAS[0], gl_RayFlagsNoneEXT, 0xff, 0,0,0, origin, 0.001, dir, 10000.0, 0);\n";
    // Reference the push constant so glslang keeps the block in the module.
    std::string_view g = c.glsl;
    std::string extra = "vec4(hitValue, 1.0)";
    if (g.find("mat4 m;") != std::string_view::npos)
        extra = "pc.m * vec4(hitValue, 1.0) + vec4(pc.l, float(pc.f))";
    else if (g.find("uint f; } pc;") != std::string_view::npos && g.find("vec4 v[2]") != std::string_view::npos)
        extra = "vec4(hitValue, 1.0) + pc.v[0] + pc.v[1] + vec4(float(pc.f))";
    else if (g.find("uint f; } pc;") != std::string_view::npos)
        extra = "vec4(hitValue, float(pc.f))";
    else if (g.find("vec4 tint;") != std::string_view::npos)
        extra = "vec4(hitValue, 1.0) + pc.tint";
    s += "  imageStore(image[0], ivec2(pixel), " + extra + ");\n";
    s += "}\n";
    return s;
}

} // namespace

int main() {
    const fs::path dir = fs::temp_directory_path() / "crafter-pcrewrite-test";
    std::error_code ec;
    fs::create_directories(dir, ec);

    int failures = 0;
    for (const Case& c : kCases) {
        const fs::path glslPath = dir / (std::string(c.name) + ".rgen.glsl");
        const fs::path spvPath  = dir / (std::string(c.name) + ".spv");
        const fs::path patched  = dir / (std::string(c.name) + ".patched.spv");

        { std::ofstream f(glslPath); f << BuildSource(c); }

        std::string compile = "glslang --target-env vulkan1.4 -V -S rgen \""
            + glslPath.string() + "\" -o \"" + spvPath.string() + "\" > /dev/null";
        if (RunCommand(compile) != 0) {
            std::println(std::cerr, "[{}] glslang failed to compile the source shader", c.name);
            ++failures;
            continue;
        }

        std::vector<std::uint32_t> words = ReadSpirv(spvPath);
        if (words.size() < 5) {
            std::println(std::cerr, "[{}] could not read compiled SPIR-V", c.name);
            ++failures;
            continue;
        }

        Device::workaroundTlasPushOffset = 0xDEADBEEFu; // poison so we know Patch set it
        WorkaroundNvidiaAS::Patch(words);
        WriteSpirv(patched, words);

        // 1. The patched module must pass spirv-val under the engine's flags.
        std::string validate = "spirv-val \"" + patched.string()
            + "\" --relax-block-layout --scalar-block-layout --target-env vulkan1.4";
        if (RunCommand(validate) != 0) {
            std::println(std::cerr, "[{}] spirv-val rejected the patched module", c.name);
            ++failures;
            continue;
        }

        // 2. Exactly one push-constant variable — the whole point of issue #18.
        int pcVars = CountPushConstantVariables(words);
        if (pcVars != 1) {
            std::println(std::cerr, "[{}] expected exactly 1 push-constant variable, found {}", c.name, pcVars);
            ++failures;
            continue;
        }

        // 3. The TLAS offset Patch published must match the expected layout end.
        if (c.readsAccelStruct && Device::workaroundTlasPushOffset != c.expectedOffset) {
            std::println(std::cerr, "[{}] expected TLAS push offset {}, got {}",
                         c.name, c.expectedOffset, Device::workaroundTlasPushOffset);
            ++failures;
            continue;
        }

        std::println(std::cout, "[{}] ok (push-constant vars: {}, tlas offset: {})",
                     c.name, pcVars, c.readsAccelStruct ? Device::workaroundTlasPushOffset : 0u);
    }

    if (failures != 0) {
        std::println(std::cerr, "{} case(s) failed", failures);
        return 1;
    }
    std::println(std::cout, "all push-constant rewrite cases passed");
    return 0;
}