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

module;
#include "vulkan/vulkan.h"
module Crafter.Graphics:ComputeShader_impl;
import :ComputeShader;
import :ShaderVulkan;
import :Device;
import :DescriptorHeapVulkan;
import std;

using namespace Crafter;

ComputeShader::ComputeShader(ComputeShader&& other) noexcept
    : pipeline(other.pipeline),
      layout(other.layout),
      workaroundNeedsTlas(other.workaroundNeedsTlas),
      workaroundTlasPushOffset(other.workaroundTlasPushOffset) {
    other.pipeline = VK_NULL_HANDLE;
}

ComputeShader& ComputeShader::operator=(ComputeShader&& other) noexcept {
    if (this != &other) {
        if (pipeline != VK_NULL_HANDLE) {
            vkDestroyPipeline(Device::device, pipeline, nullptr);
        }
        pipeline = other.pipeline;
        layout = other.layout;
        workaroundNeedsTlas = other.workaroundNeedsTlas;
        workaroundTlasPushOffset = other.workaroundTlasPushOffset;
        other.pipeline = VK_NULL_HANDLE;
    }
    return *this;
}

ComputeShader::~ComputeShader() {
    if (pipeline != VK_NULL_HANDLE) {
        vkDestroyPipeline(Device::device, pipeline, nullptr);
        pipeline = VK_NULL_HANDLE;
    }
}

void ComputeShader::Load(const std::filesystem::path& spvPath, DescriptorHeapVulkan& heap) {
    if (!Device::descriptorHeapSupported) {
        // Sets mode. The heap build of the shader uses SPV_EXT_descriptor_heap
        // and untyped pointers, which this device cannot consume, so load the
        // descriptor-set build compiled from the same source instead.
        std::filesystem::path setsPath = spvPath;
        setsPath.replace_extension("sets.spv");
        // VulkanShader only prints when a file is missing and then creates an
        // empty module, which surfaces as an unrelated driver error. A missing
        // variant is the likely mistake when a bundle is copied to a device
        // without the heap, so name it.
        if (!std::filesystem::exists(setsPath)) {
            throw std::runtime_error(std::format(
                "ComputeShader::Load: this device has no VK_EXT_descriptor_heap, so the descriptor-set build "
                "of the shader is required, but {} does not exist. Build the shader a second time with "
                "CRAFTER_GRAPHICS_DESCRIPTOR_SETS defined and ship it next to {}.",
                setsPath.string(), spvPath.string()));
        }
        if (heap.pipelineLayout == VK_NULL_HANDLE) {
            throw std::runtime_error("ComputeShader::Load: the descriptor heap must be Initialize()d before a shader is loaded against it");
        }
        VulkanShader shader(setsPath, "main", VK_SHADER_STAGE_COMPUTE_BIT, nullptr);
        layout = heap.pipelineLayout;
        VkComputePipelineCreateInfo info {
            .sType  = VK_STRUCTURE_TYPE_COMPUTE_PIPELINE_CREATE_INFO,
            .stage  = {
                .sType  = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO,
                .stage  = VK_SHADER_STAGE_COMPUTE_BIT,
                .module = shader.shader,
                .pName  = "main",
            },
            .layout = layout,
        };
        const VkResult created = vkCreateComputePipelines(
            Device::device, Device::pipelineCache, 1, &info, nullptr, &pipeline);
        // The module is only an input to pipeline creation; VulkanShader does
        // not destroy it, so it is released here either way.
        vkDestroyShaderModule(Device::device, shader.shader, nullptr);
        Device::CheckVkResult(created);
        return;
    }

    VulkanShader shader(spvPath, "main", VK_SHADER_STAGE_COMPUTE_BIT, nullptr);

    // NVIDIA descriptor-heap AS-read workaround (issue #15 / #7): remember
    // whether VulkanShader rewrote a heap acceleration-structure read in this
    // module, and where it expects the TLAS address pushed, so Dispatch can
    // feed it the per-frame TLAS. Per-shader, not a global — see ComputeShader.
    workaroundNeedsTlas = shader.patchedAS;
    workaroundTlasPushOffset = shader.tlasPushOffset;

    // Spec: with VK_PIPELINE_CREATE_2_DESCRIPTOR_HEAP_BIT_EXT, layout MUST be
    // VK_NULL_HANDLE — bindings come from the bound descriptor heap and push
    // constants are pushed via vkCmdPushDataEXT instead of vkCmdPushConstants.
    VkPipelineCreateFlags2CreateInfo flags2 {
        .sType = VK_STRUCTURE_TYPE_PIPELINE_CREATE_FLAGS_2_CREATE_INFO,
        .flags = VK_PIPELINE_CREATE_2_DESCRIPTOR_HEAP_BIT_EXT,
    };
    VkComputePipelineCreateInfo info {
        .sType  = VK_STRUCTURE_TYPE_COMPUTE_PIPELINE_CREATE_INFO,
        .pNext  = &flags2,
        .stage  = {
            .sType  = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO,
            .stage  = VK_SHADER_STAGE_COMPUTE_BIT,
            .module = shader.shader,
            .pName  = "main",
        },
        .layout = VK_NULL_HANDLE,
    };
    const VkResult created = vkCreateComputePipelines(
        Device::device, Device::pipelineCache, 1, &info, nullptr, &pipeline);
    vkDestroyShaderModule(Device::device, shader.shader, nullptr);
    Device::CheckVkResult(created);
}

void ComputeShader::Dispatch(VkCommandBuffer cmd,
                             const void* push, std::uint32_t pushBytes,
                             std::uint32_t gx,
                             std::uint32_t gy,
                             std::uint32_t gz,
                             VkDeviceAddress tlasAddress) const {
    vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, pipeline);
    if (push != nullptr && pushBytes > 0) {
        if (!Device::descriptorHeapSupported) {
            // Sets mode: the same bytes go through the shared layout's
            // push-constant range, which DescriptorHeapVulkan sized to the
            // device limit rounded down to whole words.
            const std::uint32_t limit = Device::deviceProperties.limits.maxPushConstantsSize & ~3u;
            const std::uint32_t wholeWords = (pushBytes + 3u) & ~3u;
            if (wholeWords > limit) {
                throw std::length_error(std::format(
                    "ComputeShader::Dispatch: {} bytes of push constants exceed this device's limit of {}",
                    pushBytes, limit));
            }
            if (wholeWords == pushBytes) {
                vkCmdPushConstants(cmd, layout, VK_SHADER_STAGE_COMPUTE_BIT, 0, pushBytes, push);
            } else {
                // vkCmdPushConstants only takes whole words, and the caller's
                // buffer cannot be read past its end, so push a padded copy.
                // vkCmdPushDataEXT has the same rule, but nothing enforces it
                // on the heap path and drivers accept a ragged size there;
                // padding keeps such a payload from failing only on a device
                // without the heap.
                std::vector<std::byte> padded(wholeWords);
                std::memcpy(padded.data(), push, pushBytes);
                vkCmdPushConstants(cmd, layout, VK_SHADER_STAGE_COMPUTE_BIT, 0, wholeWords, padded.data());
            }
        } else {
            VkPushDataInfoEXT pushInfo {
                .sType  = VK_STRUCTURE_TYPE_PUSH_DATA_INFO_EXT,
                .offset = 0,
                .data   = { .address = const_cast<void*>(push), .size = pushBytes },
            };
            Device::vkCmdPushDataEXT(cmd, &pushInfo);
        }
    }
    // NVIDIA descriptor-heap AS-read workaround (issue #15 / #7): if this shader
    // ray-queries the TLAS through the heap it was rewritten to read the TLAS
    // device address from a push constant; push the caller-supplied address
    // where the rewrite reads it (after any user payload, or offset 0 if none).
    // Mirrors RTPass::Record for the RT pipeline. Inert on every other driver.
    if (Device::workaroundDescriptorHeapAS && workaroundNeedsTlas) {
        VkPushDataInfoEXT tlasPush {
            .sType  = VK_STRUCTURE_TYPE_PUSH_DATA_INFO_EXT,
            .offset = workaroundTlasPushOffset,
            .data   = { .address = &tlasAddress, .size = sizeof(tlasAddress) },
        };
        Device::vkCmdPushDataEXT(cmd, &tlasPush);
    }
    vkCmdDispatch(cmd, gx, gy, gz);
}