/*
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 as published by the Free Software Foundation; either
version 3.0 of the License, or (at your option) any later version.

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_core.h>
module Crafter.Graphics:RenderingElement3DVulkan_impl;
import :RenderingElement3DVulkan;
import std;

using namespace Crafter;

std::vector<RenderingElement3DVulkan> RenderingElement3DVulkan::elements;

RenderingElement3DVulkan::RenderingElement3DVulkan(Mesh& mesh, std::uint32_t customIndex, std::uint32_t mask, std::uint32_t shaderOffset) {
    VkAccelerationStructureDeviceAddressInfoKHR addrInfo {
        .sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_DEVICE_ADDRESS_INFO_KHR,
        .accelerationStructure = mesh.accelerationStructure
    };
    VkDeviceAddress blasDeviceAddr = VulkanDevice::vkGetAccelerationStructureDeviceAddressKHR(VulkanDevice::device, &addrInfo);

    instance = {
        .instanceCustomIndex                    = customIndex,
        .mask                                   = mask,
        .instanceShaderBindingTableRecordOffset = shaderOffset,
        .flags                                  = VK_GEOMETRY_INSTANCE_FORCE_OPAQUE_BIT_KHR,
        .accelerationStructureReference         = blasDeviceAddr
    };
}

void RenderingElement3DVulkan::BuildTLAS(VkCommandBuffer cmd, std::uint32_t index) {
    tlases[index].instanceBuffer.Resize(VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_ACCELERATION_STRUCTURE_BUILD_INPUT_READ_ONLY_BIT_KHR, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT, elements.size());

    for(std::uint32_t i = 0; i < elements.size(); i++) {
        tlases[index].instanceBuffer.value[i] = elements[i].instance;
    }

    tlases[index].instanceBuffer.FlushDevice(cmd, VK_ACCESS_MEMORY_READ_BIT, VK_PIPELINE_STAGE_ACCELERATION_STRUCTURE_BUILD_BIT_KHR);

    VkAccelerationStructureGeometryInstancesDataKHR instancesData = VkAccelerationStructureGeometryInstancesDataKHR {
        .sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_INSTANCES_DATA_KHR,
        .arrayOfPointers = VK_FALSE,
        .data            = {tlases[index].instanceBuffer.address}
    };

    VkAccelerationStructureGeometryDataKHR geometryData;
    geometryData.instances = instancesData;

    VkAccelerationStructureGeometryKHR tlasGeometry {
        .sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_KHR,
        .geometryType = VK_GEOMETRY_TYPE_INSTANCES_KHR,
        .geometry     = geometryData
    };

    VkAccelerationStructureBuildGeometryInfoKHR tlasBuildGeometryInfo{
        .sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_GEOMETRY_INFO_KHR,
        .type          = VK_ACCELERATION_STRUCTURE_TYPE_TOP_LEVEL_KHR,
        .mode          = VK_BUILD_ACCELERATION_STRUCTURE_MODE_BUILD_KHR,
        .geometryCount = 1,
        .pGeometries   = &tlasGeometry
    };

    // Query the memory sizes that will be needed for this TLAS
    auto primitiveCount = static_cast<uint32_t>(elements.size());

    VkAccelerationStructureBuildSizesInfoKHR tlasBuildSizes {
        .sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_SIZES_INFO_KHR
    };
    VulkanDevice::vkGetAccelerationStructureBuildSizesKHR(
        VulkanDevice::device,
        VK_ACCELERATION_STRUCTURE_BUILD_TYPE_DEVICE_KHR,
        &tlasBuildGeometryInfo,
        &primitiveCount,
        &tlasBuildSizes
    );

    scratchBuffer.Resize(VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, tlasBuildSizes.buildScratchSize);
    tlasBuildGeometryInfo.scratchData.deviceAddress = scratchBuffer.address;

    tlases[index].buffer.Resize(VK_BUFFER_USAGE_ACCELERATION_STRUCTURE_STORAGE_BIT_KHR | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT | VK_BUFFER_USAGE_ACCELERATION_STRUCTURE_BUILD_INPUT_READ_ONLY_BIT_KHR, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, tlasBuildSizes.accelerationStructureSize);

    // Create and store the TLAS handle
    VkAccelerationStructureCreateInfoKHR tlasCreateInfo {
        .sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_CREATE_INFO_KHR,
        .buffer = tlases[index].buffer.buffer,
        .offset = 0,
        .size   = tlasBuildSizes.accelerationStructureSize,
        .type   = VK_ACCELERATION_STRUCTURE_TYPE_TOP_LEVEL_KHR,
    };

    VulkanDevice::CheckVkResult(VulkanDevice::vkCreateAccelerationStructureKHR(VulkanDevice::device, &tlasCreateInfo, nullptr, &tlases[index].accelerationStructure));
    tlasBuildGeometryInfo.dstAccelerationStructure = tlases[index].accelerationStructure;

    // Prepare the build range for the TLAS
    VkAccelerationStructureBuildRangeInfoKHR tlasRangeInfo {
        .primitiveCount  = primitiveCount,
        .primitiveOffset = 0,
        .firstVertex     = 0,
        .transformOffset = 0
    };

    VkAccelerationStructureBuildRangeInfoKHR* tlasRangeInfoPP = &tlasRangeInfo;
    VulkanDevice::vkCmdBuildAccelerationStructuresKHR(cmd, 1, &tlasBuildGeometryInfo, &tlasRangeInfoPP);

    vkCmdPipelineBarrier(
        cmd,
        VK_PIPELINE_STAGE_ACCELERATION_STRUCTURE_BUILD_BIT_KHR,
        VK_PIPELINE_STAGE_RAY_TRACING_SHADER_BIT_KHR,
        0,
        0, nullptr,
        0, nullptr,
        0, nullptr
    );
}