/*
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;
#ifdef CRAFTER_GRAPHICS_VULKAN
#include <vulkan/vulkan.h>
#endif
export module Crafter.Graphics:PipelineRTVulkan;
#ifdef CRAFTER_GRAPHICS_VULKAN
import std;
import :VulkanDevice;
import :VulkanBuffer;
import :Types;

export namespace Crafter {
    template <typename Raygen, typename ClosestHit, typename Miss, typename ShadowClosestHit, typename ShadowMiss, typename... Shaders>
    class PipelineRTVulkan {
        public:
        inline static VkPipeline pipeline;
        inline static VkPipelineLayout pipelineLayout;
        inline static std::vector<std::uint8_t> shaderHandles;  
        inline static VulkanBuffer<std::uint8_t, true, true, false> sbtBuffer;
        inline static VkStridedDeviceAddressRegionKHR raygenRegion;
        inline static VkStridedDeviceAddressRegionKHR missRegion;
        inline static VkStridedDeviceAddressRegionKHR hitRegion;
        inline static VkStridedDeviceAddressRegionKHR callableRegion;

        static void Init(std::span<VkDescriptorSetLayout> setLayouts) {
            VkPipelineLayoutCreateInfo pipelineLayoutInfo {
                .sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO,
                .setLayoutCount = static_cast<std::uint32_t>(setLayouts.size()),
                .pSetLayouts = setLayouts.data()
            };

			VulkanDevice::CheckVkResult(vkCreatePipelineLayout(VulkanDevice::device, &pipelineLayoutInfo, nullptr, &pipelineLayout));

			std::array<VkPipelineShaderStageCreateInfo, 5> shaderStages;

            shaderStages[0].sType  = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
            shaderStages[0].stage  = Raygen::_stage;
            shaderStages[0].module = Raygen::shader;
            shaderStages[0].pName  = Raygen::_entrypoint.value;
            shaderStages[0].flags = 0;
            shaderStages[0].pSpecializationInfo = nullptr;
            shaderStages[0].pNext = nullptr;

            shaderStages[1].sType  = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
            shaderStages[1].stage  = Miss::_stage;
            shaderStages[1].module = Miss::shader;
            shaderStages[1].pName  = Miss::_entrypoint.value;
            shaderStages[1].flags = 0;
            shaderStages[1].pSpecializationInfo = nullptr;
            shaderStages[1].pNext = nullptr;

            shaderStages[2].sType  = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
            shaderStages[2].stage  = ClosestHit::_stage;
            shaderStages[2].module = ClosestHit::shader;
            shaderStages[2].pName  = ClosestHit::_entrypoint.value;
            shaderStages[2].flags = 0;
            shaderStages[2].pSpecializationInfo = nullptr;
            shaderStages[2].pNext = nullptr;

            shaderStages[3].sType  = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
            shaderStages[3].stage  = VK_SHADER_STAGE_CLOSEST_HIT_BIT_KHR;
            shaderStages[3].module = ShadowClosestHit::shader;
            shaderStages[3].pName  = ShadowClosestHit::_entrypoint.value;
            shaderStages[3].flags = 0;
            shaderStages[3].pSpecializationInfo = nullptr;
            shaderStages[3].pNext = nullptr;

            shaderStages[4].sType  = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
            shaderStages[4].stage  = ShadowMiss::_stage;
            shaderStages[4].module = ShadowMiss::shader;
            shaderStages[4].pName  = ShadowMiss::_entrypoint.value;
            shaderStages[4].flags = 0;
            shaderStages[4].pSpecializationInfo = nullptr;
            shaderStages[4].pNext = nullptr;

            std::array<VkRayTracingShaderGroupCreateInfoKHR, 5> groups {{
                {
                    .sType = VK_STRUCTURE_TYPE_RAY_TRACING_SHADER_GROUP_CREATE_INFO_KHR,
                    .type = VK_RAY_TRACING_SHADER_GROUP_TYPE_GENERAL_KHR,
                    .generalShader = 0,
                    .closestHitShader = VK_SHADER_UNUSED_KHR,
                    .anyHitShader = VK_SHADER_UNUSED_KHR,
                    .intersectionShader = VK_SHADER_UNUSED_KHR
                },
                {
                    .sType = VK_STRUCTURE_TYPE_RAY_TRACING_SHADER_GROUP_CREATE_INFO_KHR,
                    .type = VK_RAY_TRACING_SHADER_GROUP_TYPE_GENERAL_KHR,
                    .generalShader = 1,
                    .closestHitShader = VK_SHADER_UNUSED_KHR,
                    .anyHitShader = VK_SHADER_UNUSED_KHR,
                    .intersectionShader = VK_SHADER_UNUSED_KHR
                },
                {
                    .sType = VK_STRUCTURE_TYPE_RAY_TRACING_SHADER_GROUP_CREATE_INFO_KHR,
                    .type = VK_RAY_TRACING_SHADER_GROUP_TYPE_TRIANGLES_HIT_GROUP_KHR,
                    .generalShader = VK_SHADER_UNUSED_KHR,
                    .closestHitShader = 2,
                    .anyHitShader = VK_SHADER_UNUSED_KHR,
                    .intersectionShader = VK_SHADER_UNUSED_KHR
                },
                {
                    .sType = VK_STRUCTURE_TYPE_RAY_TRACING_SHADER_GROUP_CREATE_INFO_KHR,
                    .type = VK_RAY_TRACING_SHADER_GROUP_TYPE_TRIANGLES_HIT_GROUP_KHR,
                    .generalShader = VK_SHADER_UNUSED_KHR,
                    .closestHitShader = 3,
                    .anyHitShader = VK_SHADER_UNUSED_KHR,
                    .intersectionShader = VK_SHADER_UNUSED_KHR
                },
                {
                    .sType = VK_STRUCTURE_TYPE_RAY_TRACING_SHADER_GROUP_CREATE_INFO_KHR,
                    .type = VK_RAY_TRACING_SHADER_GROUP_TYPE_GENERAL_KHR,
                    .generalShader = 4,
                    .closestHitShader = VK_SHADER_UNUSED_KHR,
                    .anyHitShader = VK_SHADER_UNUSED_KHR,
                    .intersectionShader = VK_SHADER_UNUSED_KHR
                }
            }};

            VkRayTracingPipelineCreateInfoKHR rtPipelineInfo{
                .sType = VK_STRUCTURE_TYPE_RAY_TRACING_PIPELINE_CREATE_INFO_KHR,
                .stageCount = static_cast<std::uint32_t>(shaderStages.size()),
                .pStages = shaderStages.data(),
                .groupCount = static_cast<std::uint32_t>(groups.size()),
                .pGroups = groups.data(),
                .maxPipelineRayRecursionDepth = 2,
                .layout = pipelineLayout
            };

            VulkanDevice::CheckVkResult(VulkanDevice::vkCreateRayTracingPipelinesKHR(VulkanDevice::device, {}, {}, 1, &rtPipelineInfo, nullptr, &pipeline));

            std::uint32_t handleSize      = VulkanDevice::rayTracingProperties.shaderGroupHandleSize;
            std::uint32_t handleAlignment = VulkanDevice::rayTracingProperties.shaderGroupHandleAlignment;
            std::uint32_t baseAlignment   = VulkanDevice::rayTracingProperties.shaderGroupBaseAlignment;
            std::uint32_t groupCount      = rtPipelineInfo.groupCount;

            std::size_t dataSize = handleSize * groupCount;
            shaderHandles.resize(dataSize);
            VulkanDevice::CheckVkResult(VulkanDevice::vkGetRayTracingShaderGroupHandlesKHR(VulkanDevice::device, pipeline, 0, groupCount, dataSize, shaderHandles.data()));

            std::uint32_t sbtStride = AlignUp(handleSize, handleAlignment);
            std::uint32_t raygenOffset = 0;
            std::uint32_t missOffset     = AlignUp(raygenOffset + sbtStride, baseAlignment);
            std::uint32_t hitOffset      = AlignUp(missOffset + sbtStride * 2, baseAlignment);

            std::uint32_t hitGroupCount = 2;
            std::size_t bufferSize = hitOffset + sbtStride * hitGroupCount;

            sbtBuffer.Create(VK_BUFFER_USAGE_2_SHADER_BINDING_TABLE_BIT_KHR | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT, bufferSize);

            // Ray generation shader (group 0)
            std::memcpy(sbtBuffer.value + raygenOffset, shaderHandles.data() + 0 * handleSize, handleSize);
            raygenRegion.deviceAddress = sbtBuffer.address + raygenOffset;
            raygenRegion.stride        = sbtStride;
            raygenRegion.size          = sbtStride;

            std::memcpy(sbtBuffer.value + missOffset, shaderHandles.data() + 1 * handleSize, handleSize);
            std::memcpy(sbtBuffer.value + missOffset + 1 * sbtStride, shaderHandles.data() + 4 * handleSize, handleSize);
            missRegion.deviceAddress = sbtBuffer.address + missOffset;
            missRegion.stride        = sbtStride;
            missRegion.size          = sbtStride * 2;

            std::memcpy(sbtBuffer.value + hitOffset + 0 * sbtStride, shaderHandles.data() + 2 * handleSize, handleSize);
            std::memcpy(sbtBuffer.value + hitOffset + 1 * sbtStride, shaderHandles.data() + 3 * handleSize, handleSize);
            hitRegion.deviceAddress = sbtBuffer.address + hitOffset;
            hitRegion.stride = sbtStride;
            hitRegion.size   = sbtStride * hitGroupCount;

            callableRegion.deviceAddress = 0;
            callableRegion.stride        = 0;
            callableRegion.size          = 0;
        }
    };
}

#endif