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

module;
#ifndef CRAFTER_GRAPHICS_WINDOW_DOM
#include "vulkan/vulkan.h"
#endif // !CRAFTER_GRAPHICS_WINDOW_DOM
export module Crafter.Graphics:PipelineRTVulkan;
#ifndef CRAFTER_GRAPHICS_WINDOW_DOM
import std;
import :Device;
import :VulkanBuffer;
import :ShaderBindingTableVulkan;
import :Types;

export namespace Crafter {
    struct PipelineRTVulkan {
        VkPipeline pipeline;
        std::vector<std::uint8_t> shaderHandles;  
        VulkanBuffer<std::uint8_t, true> sbtBuffer;
        VkStridedDeviceAddressRegionKHR raygenRegion;
        VkStridedDeviceAddressRegionKHR missRegion;
        VkStridedDeviceAddressRegionKHR hitRegion;
        VkStridedDeviceAddressRegionKHR callableRegion;

        // NVIDIA descriptor-heap AS-read workaround (issue #15 / #7): copied
        // from the shader table at Init so RTPass can push the active TLAS
        // device address into the patched shaders' push constant. Inert on
        // every other driver.
        bool workaroundNeedsTlas = false;
        std::uint32_t workaroundTlasPushOffset = 0;

        // maxRecursionDepth: the maximum ray-recursion depth the pipeline must
        // support — i.e. the deepest chain of nested traceRayEXT calls. The
        // raygen counts as depth 1, so a closest-hit shader that traces a shadow
        // ray needs 2. Tracing beyond the value the pipeline was created with is
        // undefined behaviour and faults the device, so a consumer with any
        // recursion past the raygen must raise this. Defaults to 1 (raygen-only,
        // matching the simple examples) and is clamped to the device's
        // maxRayRecursionDepth.
        void Init(VkCommandBuffer cmd, std::span<VkRayTracingShaderGroupCreateInfoKHR> raygenGroups, std::span<VkRayTracingShaderGroupCreateInfoKHR> missGroups, std::span<VkRayTracingShaderGroupCreateInfoKHR> hitGroups, ShaderBindingTableVulkan& shaderTable, std::uint32_t maxRecursionDepth = 1) {
            workaroundNeedsTlas = shaderTable.workaroundNeedsTlas;
            workaroundTlasPushOffset = shaderTable.workaroundTlasPushOffset;

            std::vector<VkRayTracingShaderGroupCreateInfoKHR> groups;
            groups.reserve(raygenGroups.size() + missGroups.size() + hitGroups.size());

            groups.insert(groups.end(), raygenGroups.begin(), raygenGroups.end());
            groups.insert(groups.end(), missGroups.begin(), missGroups.end());
            groups.insert(groups.end(), hitGroups.begin(), hitGroups.end());

            VkPipelineCreateFlags2CreateInfo flags2 = {
                .sType = VK_STRUCTURE_TYPE_PIPELINE_CREATE_FLAGS_2_CREATE_INFO,
                .flags = VK_PIPELINE_CREATE_2_DESCRIPTOR_HEAP_BIT_EXT
            };

            VkRayTracingPipelineCreateInfoKHR rtPipelineInfo {
                .sType = VK_STRUCTURE_TYPE_RAY_TRACING_PIPELINE_CREATE_INFO_KHR,
                .pNext = &flags2,
                .flags = 0,
                .stageCount = static_cast<std::uint32_t>(shaderTable.shaderStages.size()),
                .pStages = shaderTable.shaderStages.data(),
                .groupCount = static_cast<std::uint32_t>(groups.size()),
                .pGroups = groups.data(),
                .maxPipelineRayRecursionDepth = std::min(maxRecursionDepth, Device::rayTracingProperties.maxRayRecursionDepth),
                .layout = VK_NULL_HANDLE
            };

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

            std::size_t dataSize = Device::rayTracingProperties.shaderGroupHandleSize * rtPipelineInfo.groupCount;
            shaderHandles.resize(dataSize);
            Device::CheckVkResult(Device::vkGetRayTracingShaderGroupHandlesKHR(Device::device, pipeline, 0, rtPipelineInfo.groupCount, dataSize, shaderHandles.data()));

            std::uint32_t sbtStride = AlignUp(Device::rayTracingProperties.shaderGroupHandleSize, Device::rayTracingProperties.shaderGroupHandleAlignment);

            raygenRegion.stride = sbtStride;
            raygenRegion.deviceAddress = 0;
            raygenRegion.size = raygenGroups.size() * sbtStride;

            missRegion.stride = sbtStride;
            missRegion.deviceAddress = AlignUp(raygenRegion.size, Device::rayTracingProperties.shaderGroupBaseAlignment);
            missRegion.size = missGroups.size() * sbtStride;

            hitRegion.stride = sbtStride;
            hitRegion.deviceAddress = AlignUp(missRegion.deviceAddress + missRegion.size, Device::rayTracingProperties.shaderGroupBaseAlignment);
            hitRegion.size = hitGroups.size() * sbtStride;

            std::size_t bufferSize =  hitRegion.deviceAddress + hitRegion.size;
            // The SBT is written once here (the memcpys below) and read by the
            // GPU on every vkCmdTraceRaysKHR for the pipeline's lifetime — the
            // textbook write-once/read-many buffer (issue #72). Get it into
            // device-local memory so trace dispatches read raygen/miss/hit
            // records out of VRAM instead of over PCIe from system RAM.
            //
            // Route the placement through #89: when a DEVICE_LOCAL|HOST_VISIBLE
            // type exists (ReBAR/UMA, or a BAR window — the SBT is tiny so the
            // window budget is never the constraint), prefer DEVICE_LOCAL on top
            // of the required HOST_VISIBLE so the allocation lands in device
            // memory we can still map. The one-time memcpy goes write-combined
            // over PCIe (write-only, sequential — ideal); GPU reads then hit
            // local VRAM. When no combined type exists at all (no spec
            // guarantee), PreferDirectDeviceWrite returns false and the preferred
            // hint is dropped, so GetMemoryType falls back to plain HOST_VISIBLE
            // (current behaviour, the per-trace PCIe read) — the cheaper fallback
            // the issue blesses while staging isn't wired here.
            //
            // Either way the buffer stays mapped, and the FlushDevice below gates
            // its flush on the *chosen* memory type's flags (issue #60), so a
            // direct-write type lacking HOST_COHERENT is still flushed correctly.
            VkMemoryPropertyFlags sbtPreferred = Device::PreferDirectDeviceWrite(bufferSize)
                ? VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT : 0;
            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, sbtPreferred);

            std::uint8_t* offset = sbtBuffer.value;
            std::uint8_t* handleOffset = shaderHandles.data();

            std::memcpy(offset, handleOffset, raygenGroups.size() * Device::rayTracingProperties.shaderGroupHandleSize);
            offset += AlignUp(raygenRegion.size, Device::rayTracingProperties.shaderGroupBaseAlignment); 
            handleOffset += raygenGroups.size() * Device::rayTracingProperties.shaderGroupHandleSize;

            std::memcpy(offset, handleOffset, missGroups.size() * Device::rayTracingProperties.shaderGroupHandleSize);
            offset += AlignUp(missRegion.size, Device::rayTracingProperties.shaderGroupBaseAlignment); 
            handleOffset += missGroups.size() * Device::rayTracingProperties.shaderGroupHandleSize;

            std::memcpy(offset, handleOffset, hitGroups.size() * Device::rayTracingProperties.shaderGroupHandleSize);

            sbtBuffer.FlushDevice(cmd, VK_ACCESS_MEMORY_READ_BIT, VK_PIPELINE_STAGE_RAY_TRACING_SHADER_BIT_KHR);

            raygenRegion.deviceAddress += sbtBuffer.address;
            missRegion.deviceAddress += sbtBuffer.address;
            hitRegion.deviceAddress += sbtBuffer.address;

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

        ~PipelineRTVulkan() {
            vkDestroyPipeline(Device::device, pipeline, nullptr);
        }
    };
}
#endif // !CRAFTER_GRAPHICS_WINDOW_DOM