/*
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) {
    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;
        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);

    // 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) 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);
    }
    vkCmdDispatch(cmd, gx, gy, gz);
}