/*
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
*/
module;
#include "vulkan/vulkan.h"
module Crafter.Graphics:ComputeShader_impl;
import :ComputeShader;
import :ShaderVulkan;
import :Device;
import std;

using namespace Crafter;

ComputeShader::ComputeShader(ComputeShader&& other) noexcept
    : pipeline(other.pipeline),
      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;
        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) {
    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,
    };
    Device::CheckVkResult(vkCreateComputePipelines(
        Device::device, VK_NULL_HANDLE, 1, &info, nullptr, &pipeline));
}

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