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

// Device::Initialize enables what the library's own shaders use plus what the
// application passes it, skips a GPU that lacks any of that, and VulkanShader
// refuses a module that declares a capability nobody enabled. The decisions
// behind those are pure logic: Device::ResolveFeatures (the features to
// enable, and the gaps a GPU has), Device::ChooseDevice (the GPU to use) and
// Device::CheckShaderCapabilities (whether a module's capabilities are
// enabled). This test drives them with synthetic GPUs and hand-built SPIR-V,
// so it needs no GPU, mirroring PipelineCacheValidation. tests/FeatureRequests
// checks the same wiring end to end on a real GPU.

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

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

namespace {

int failures = 0;

void Check(bool ok, std::string_view what) {
    std::println("{}  {}", ok ? "PASS" : "FAIL", what);
    if (!ok) ++failures;
}

// The five core feature structs as DeviceFeatures lays them out: where each
// one's VkBool32 members start, and how many it has (vk.xml; a core feature
// struct never changes once its version ships).
struct Block {
    std::string_view version;
    std::size_t first;
    std::size_t count;
};
const std::array<Block, 5> blocks {{
    { "vulkan10", offsetof(DeviceFeatures, vulkan10) + offsetof(VkPhysicalDeviceFeatures, robustBufferAccess), 55 },
    { "vulkan11", offsetof(DeviceFeatures, vulkan11) + offsetof(VkPhysicalDeviceVulkan11Features, storageBuffer16BitAccess), 12 },
    { "vulkan12", offsetof(DeviceFeatures, vulkan12) + offsetof(VkPhysicalDeviceVulkan12Features, samplerMirrorClampToEdge), 47 },
    { "vulkan13", offsetof(DeviceFeatures, vulkan13) + offsetof(VkPhysicalDeviceVulkan13Features, robustImageAccess), 15 },
    { "vulkan14", offsetof(DeviceFeatures, vulkan14) + offsetof(VkPhysicalDeviceVulkan14Features, globalPriorityQuery), 21 },
}};

void SetSlot(DeviceFeatures& features, std::size_t offset) {
    const VkBool32 on = VK_TRUE;
    std::memcpy(reinterpret_cast<std::byte*>(&features) + offset, &on, sizeof(on));
}

// Whether a and b have the same VkBool32s on, ignoring sType, pNext and
// padding.
bool SameFeatures(const DeviceFeatures& a, const DeviceFeatures& b) {
    for (const Block& block : blocks) {
        for (std::size_t i = 0; i < block.count; ++i) {
            const std::size_t offset = block.first + i * sizeof(VkBool32);
            VkBool32 x;
            VkBool32 y;
            std::memcpy(&x, reinterpret_cast<const std::byte*>(&a) + offset, sizeof(x));
            std::memcpy(&y, reinterpret_cast<const std::byte*>(&b) + offset, sizeof(y));
            if ((x != VK_FALSE) != (y != VK_FALSE)) {
                return false;
            }
        }
    }
    return true;
}

// Every core feature on.
DeviceFeatures AllOn() {
    DeviceFeatures features;
    for (const Block& block : blocks) {
        for (std::size_t i = 0; i < block.count; ++i) {
            SetSlot(features, block.first + i * sizeof(VkBool32));
        }
    }
    return features;
}

// A module that declares Shader and then `capabilities`, ending the
// capability section with an OpMemoryModel.
std::vector<std::uint32_t> Module(std::initializer_list<std::uint32_t> capabilities) {
    constexpr std::uint32_t opCapability = 17;
    constexpr std::uint32_t opMemoryModel = 14;
    std::vector<std::uint32_t> words { 0x07230203u, 0x00010400u, 0, 1, 0 };
    words.insert(words.end(), { opCapability | (2u << 16), 1 /* Shader */ });
    for (std::uint32_t capability : capabilities) {
        words.insert(words.end(), { opCapability | (2u << 16), capability });
    }
    words.insert(words.end(), { opMemoryModel | (3u << 16), 0, 1 });
    return words;
}

// The message CheckShaderCapabilities throws for `spirv`, or "" when it lets
// the module through.
std::string Refusal(const std::vector<std::uint32_t>& spirv) {
    try {
        Device::CheckShaderCapabilities(spirv, "shaders/test.comp.spv");
    } catch (const std::runtime_error& error) {
        return error.what();
    }
    return "";
}

// A 1.4 GPU with every core feature, enabled nowhere, and no optional group.
void ResetDevice() {
    Device::enabledFeatures = {};
    Device::supportedFeatures = AllOn();
    Device::apiVersion = VK_API_VERSION_1_4;
    Device::vulkan11Properties = { .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_1_PROPERTIES };
    Device::vulkan11Properties.subgroupSupportedOperations = VK_SUBGROUP_FEATURE_BASIC_BIT;
    Device::shaderInt64ForWorkaroundOnly = false;
    Device::descriptorHeapSupported = false;
    Device::rayTracingSupported = false;
}

} // namespace

int main() {
    // --- what gets enabled --------------------------------------------------
    {
        DeviceFeatures library;
        library.vulkan10.samplerAnisotropy = VK_TRUE;
        // Designated initializers leave sType zero, as an application's will.
        const DeviceFeatures required { .vulkan12 = { .shaderFloat16 = VK_TRUE } };
        const DeviceFeatures optional {
            .vulkan11 = { .storageBuffer16BitAccess = VK_TRUE },
            .vulkan13 = { .synchronization2 = VK_TRUE },
        };
        DeviceFeatures supported = AllOn();
        supported.vulkan13.synchronization2 = VK_FALSE;

        FeatureGaps gaps;
        const DeviceFeatures enabled = Device::ResolveFeatures(library, required, optional, supported, VK_API_VERSION_1_4, gaps);
        Check(gaps.library.empty() && gaps.application.empty(), "a GPU with everything needed has no gaps");
        Check(enabled.vulkan10.samplerAnisotropy && enabled.vulkan12.shaderFloat16,
              "the library's and the application's required features are enabled");
        Check(enabled.vulkan11.storageBuffer16BitAccess, "an optional feature the GPU has is enabled");
        Check(!enabled.vulkan13.synchronization2, "an optional feature the GPU lacks stays off, silently");
        Check(!enabled.vulkan10.shaderInt16 && !enabled.vulkan14.pushDescriptor,
              "nothing nobody asked for is enabled");
        Check(enabled.vulkan12.sType == VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_2_FEATURES && enabled.vulkan12.pNext == nullptr,
              "the result carries the sTypes and no pNext, ready to chain");
    }

    // --- gaps, by who needs the feature -------------------------------------
    {
        DeviceFeatures library;
        library.vulkan10.samplerAnisotropy = VK_TRUE;
        const DeviceFeatures required { .vulkan12 = { .shaderFloat16 = VK_TRUE } };
        DeviceFeatures supported = AllOn();
        supported.vulkan10.samplerAnisotropy = VK_FALSE;
        supported.vulkan12.shaderFloat16 = VK_FALSE;

        FeatureGaps gaps;
        const DeviceFeatures enabled = Device::ResolveFeatures(library, required, {}, supported, VK_API_VERSION_1_4, gaps);
        Check(gaps.library == std::vector<std::string>{ "vulkan10.samplerAnisotropy" },
              "a library need the GPU lacks is the library's gap");
        Check(gaps.application == std::vector<std::string>{ "vulkan12.shaderFloat16" },
              "a required feature the GPU lacks is the application's gap, named as the application writes it");
        Check(!enabled.vulkan12.shaderFloat16, "a feature the GPU lacks is never enabled");
    }

    // --- Vulkan 1.3 GPUs ----------------------------------------------------
    {
        DeviceFeatures library;
        library.vulkan14.maintenance5 = VK_TRUE;
        const DeviceFeatures required { .vulkan14 = { .pushDescriptor = VK_TRUE } };
        // Below 1.4 a GPU supports no vulkan14 feature but maintenance5, which
        // Initialize copies there from VK_KHR_maintenance5's struct (that copy
        // is tests/FeatureRequests' to check, on a real GPU).
        DeviceFeatures supported = AllOn();
        supported.vulkan14 = { .maintenance5 = VK_TRUE };

        FeatureGaps gaps;
        const DeviceFeatures enabled = Device::ResolveFeatures(library, required, {}, supported, VK_API_VERSION_1_3, gaps);
        Check(gaps.library.empty() && enabled.vulkan14.maintenance5,
              "at 1.3 a supported vulkan14.maintenance5 meets the library's need, with no version note");
        Check(gaps.application == std::vector<std::string>{ "vulkan14.pushDescriptor (Vulkan 1.4 only)" },
              "at 1.3 a required vulkan14 feature is a gap that names the version");

        FeatureGaps noMaintenance5;
        DeviceFeatures without = AllOn();
        without.vulkan14 = {};
        Device::ResolveFeatures(library, {}, {}, without, VK_API_VERSION_1_3, noMaintenance5);
        Check(noMaintenance5.library == std::vector<std::string>{ "vulkan14.maintenance5 (VK_KHR_maintenance5 reports it unsupported)" },
              "at 1.3 a missing maintenance5 blames VK_KHR_maintenance5's feature bit, not the version");

        FeatureGaps gaps14;
        DeviceFeatures supported14 = AllOn();
        supported14.vulkan14.pushDescriptor = VK_FALSE;
        Device::ResolveFeatures(library, required, {}, supported14, VK_API_VERSION_1_4, gaps14);
        Check(gaps14.application == std::vector<std::string>{ "vulkan14.pushDescriptor" },
              "at 1.4 the same gap carries no version note");
    }

    // --- every core feature, by name ----------------------------------------
    {
        std::set<std::string> names;
        bool eachNamedOnce = true;
        bool eachEnabledAlone = true;
        bool eachUnderItsVersion = true;
        for (const Block& block : blocks) {
            for (std::size_t i = 0; i < block.count; ++i) {
                DeviceFeatures required;
                SetSlot(required, block.first + i * sizeof(VkBool32));
                FeatureGaps gaps;
                Device::ResolveFeatures({}, required, {}, {}, VK_API_VERSION_1_4, gaps);
                if (gaps.application.size() != 1 || !names.insert(gaps.application[0]).second) {
                    eachNamedOnce = false;
                    continue;
                }
                eachUnderItsVersion = eachUnderItsVersion && gaps.application[0].starts_with(std::format("{}.", block.version));
                // On a GPU with everything, only that one gets enabled.
                FeatureGaps none;
                const DeviceFeatures enabled = Device::ResolveFeatures({}, required, {}, AllOn(), VK_API_VERSION_1_4, none);
                eachEnabledAlone = eachEnabledAlone && SameFeatures(enabled, required);
            }
        }
        Check(eachNamedOnce && names.size() == 150, "each of the 150 core features has a gap name of its own");
        Check(eachUnderItsVersion, "each name starts with its struct's member in DeviceFeatures");
        Check(eachEnabledAlone, "requiring one feature enables that one alone");
        Check(names.contains("vulkan10.robustBufferAccess") && names.contains("vulkan10.shaderInt16")
                  && names.contains("vulkan11.storageBuffer16BitAccess") && names.contains("vulkan12.scalarBlockLayout")
                  && names.contains("vulkan14.pushDescriptor"),
              "names match the members (spot checks)");
    }

    // --- choosing the GPU ---------------------------------------------------
    {
        const std::vector<std::string> noGaps;
        const std::vector<std::string> gap { "lacks vulkan12.shaderFloat16, which the application requires" };
        const DeviceChoice discrete { VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU, "dGPU", noGaps };
        const DeviceChoice discreteLacking { VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU, "old dGPU", gap };
        const DeviceChoice integrated { VK_PHYSICAL_DEVICE_TYPE_INTEGRATED_GPU, "iGPU", noGaps };
        const DeviceChoice integratedLacking { VK_PHYSICAL_DEVICE_TYPE_INTEGRATED_GPU, "old iGPU", gap };
        const DeviceChoice cpu { VK_PHYSICAL_DEVICE_TYPE_CPU, "llvmpipe", noGaps };

        Check(Device::ChooseDevice(std::vector{ integrated, discrete }) == 1, "a discrete GPU wins over an integrated one");
        Check(Device::ChooseDevice(std::vector{ discrete, discrete }) == 0, "of two equal GPUs the first is used");
        Check(Device::ChooseDevice(std::vector{ discreteLacking, integrated, cpu }) == 1,
              "a GPU that lacks something is skipped for the next one that can run the application");
        Check(Device::ChooseDevice(std::vector{ cpu }) == 0, "a CPU device runs the application when nothing else is listed");

        std::string message;
        try {
            Device::ChooseDevice(std::vector{ discreteLacking, integratedLacking, cpu });
        } catch (const std::runtime_error& error) {
            message = error.what();
        }
        Check(!message.empty(), "GPUs that lack something don't fall back to a CPU device");
        Check(message.contains("old dGPU: lacks vulkan12.shaderFloat16") && message.contains("old iGPU: lacks")
                  && message.contains("llvmpipe: a CPU device"),
              "the error names every GPU and why it can't be used");

        std::string none;
        try {
            Device::ChooseDevice(std::vector<DeviceChoice>{});
        } catch (const std::runtime_error& error) {
            none = error.what();
        }
        Check(none.contains("lists none"), "no GPU at all is an error that says so");
    }

    // --- SPIR-V capabilities against the enabled features -------------------
    {
        constexpr std::uint32_t int16 = 22;
        constexpr std::uint32_t int64 = 11;
        constexpr std::uint32_t int64Atomics = 12;
        constexpr std::uint32_t storageImageReadWithoutFormat = 55;
        constexpr std::uint32_t groupNonUniform = 61;
        constexpr std::uint32_t groupNonUniformArithmetic = 63;
        constexpr std::uint32_t shaderNonUniform = 5301;
        constexpr std::uint32_t rayQueryKHR = 4472;
        constexpr std::uint32_t descriptorHeapEXT = 5128;
        constexpr std::uint32_t groupNonUniformRotateKHR = 6026;
        constexpr std::uint32_t floatControls2 = 6029;

        ResetDevice();
        const std::string refusal = Refusal(Module({ int16 }));
        Check(refusal.contains("shaders/test.comp.spv") && refusal.contains("Int16")
                  && refusal.contains("vulkan10.shaderInt16: pass it to Device::Initialize"),
              "Int16 without shaderInt16, on a GPU that has it, is refused, naming the file and the field to pass");
        Device::enabledFeatures.vulkan10.shaderInt16 = VK_TRUE;
        Check(Refusal(Module({ int16 })).empty(), "...and accepted once shaderInt16 is enabled");

        ResetDevice();
        Device::supportedFeatures.vulkan10.shaderInt16 = VK_FALSE;
        Check(Refusal(Module({ int16 })).contains("this GPU lacks it"),
              "on a GPU without shaderInt16 the refusal says the GPU lacks it, not to pass it");

        ResetDevice();
        Device::apiVersion = VK_API_VERSION_1_3;
        Device::supportedFeatures.vulkan14 = { .maintenance5 = VK_TRUE };
        Check(Refusal(Module({ floatControls2 })).contains("a Vulkan 1.4 feature, and this GPU runs Vulkan 1.3"),
              "a vulkan14 capability on a 1.3 GPU names the version");
        ResetDevice();
        Check(Refusal(Module({ floatControls2 })).contains("vulkan14.shaderFloatControls2: pass it"),
              "FloatControls2 is refused without vulkan14.shaderFloatControls2");
        Device::enabledFeatures.vulkan14.shaderFloatControls2 = VK_TRUE;
        Check(Refusal(Module({ floatControls2 })).empty(), "...and accepted with it");

        ResetDevice();
        Check(Refusal(Module({ storageImageReadWithoutFormat, shaderNonUniform })).empty(),
              "capabilities a core version enables are left to the driver");

        ResetDevice();
        Device::enabledFeatures.vulkan10.shaderInt64 = VK_TRUE;
        Device::shaderInt64ForWorkaroundOnly = true;
        Check(Refusal(Module({ int64 })).contains("NVIDIA workaround"),
              "Int64 on only for the NVIDIA workaround doesn't let an application's Int64 through");
        Device::shaderInt64ForWorkaroundOnly = false;
        Check(Refusal(Module({ int64 })).empty(), "...but does once the application asked for shaderInt64");

        ResetDevice();
        Device::enabledFeatures.vulkan12.shaderSharedInt64Atomics = VK_TRUE;
        Check(Refusal(Module({ int64Atomics })).contains("Int64Atomics, which implicitly declares Int64"),
              "Int64Atomics alone still needs shaderInt64: it implicitly declares Int64");
        Device::enabledFeatures.vulkan10.shaderInt64 = VK_TRUE;
        Check(Refusal(Module({ int64Atomics })).empty(), "a capability with two enabling features takes either");
        Device::shaderInt64ForWorkaroundOnly = true;
        Check(Refusal(Module({ int64Atomics })).contains("NVIDIA workaround"),
              "an implied Int64 doesn't pass on the NVIDIA workaround's shaderInt64 either");

        ResetDevice();
        Check(Refusal(Module({ groupNonUniformArithmetic })).contains("VK_SUBGROUP_FEATURE_ARITHMETIC_BIT"),
              "subgroup arithmetic the GPU lacks is refused");
        Device::vulkan11Properties.subgroupSupportedOperations |= VK_SUBGROUP_FEATURE_ARITHMETIC_BIT;
        Check(Refusal(Module({ groupNonUniformArithmetic })).empty(),
              "...and accepted where the GPU has it, whatever the stage: Vulkan gates running subgroup operations by stage, not declaring them");
        Device::enabledFeatures.vulkan14.shaderSubgroupRotate = VK_TRUE;
        Device::vulkan11Properties.subgroupSupportedOperations = 0;
        Check(Refusal(Module({ groupNonUniformRotateKHR })).contains("implicitly declares GroupNonUniform, but this GPU has no"),
              "GroupNonUniformRotateKHR implicitly declares GroupNonUniform, and is held to it");

        ResetDevice();
        Check(Refusal(Module({ rayQueryKHR })).contains("Device::rayTracingSupported"),
              "a ray-query module is refused where ray tracing is off");
        Device::rayTracingSupported = true;
        Check(Refusal(Module({ rayQueryKHR })).empty(), "...and accepted where it is on");
        Check(Refusal(Module({ descriptorHeapEXT })).contains("Device::descriptorHeapSupported"),
              "a descriptor-heap module is refused where the heap is off");
        Device::descriptorHeapSupported = true;
        Check(Refusal(Module({ descriptorHeapEXT })).empty(), "...and accepted where it is on");

        ResetDevice();
        std::vector<std::uint32_t> afterLayout = Module({});
        afterLayout.insert(afterLayout.end(), { 17u | (2u << 16), int16 });
        Check(Refusal(afterLayout).empty(), "only the capability section at the start of the module is read");
        std::vector<std::uint32_t> malformed { 0x07230203u, 0x00010400u, 0, 1, 0, 17u /* length 0 */ };
        Check(Refusal(malformed).empty(), "a malformed instruction ends the scan instead of running past the module");
        Check(Refusal(Module({ groupNonUniform })).empty(), "GroupNonUniform passes on the BASIC operations every GPU has");
    }
    ResetDevice();

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