/*
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:Mesh_impl;
import Crafter.Math;
import :Mesh;
import :VulkanDevice;
import :Types;
import std;

using namespace Crafter;

void Mesh::Build(std::span<Vector<float, 3, 3>> verticies, std::span<std::uint32_t> indicies, VkCommandBuffer cmd) {
    vertexBuffer.Resize(VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT | VK_BUFFER_USAGE_2_ACCELERATION_STRUCTURE_BUILD_INPUT_READ_ONLY_BIT_KHR, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT, verticies.size());
    indexBuffer.Resize(VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT | VK_BUFFER_USAGE_2_ACCELERATION_STRUCTURE_BUILD_INPUT_READ_ONLY_BIT_KHR | VK_BUFFER_USAGE_STORAGE_BUFFER_BIT, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT, indicies.size());

    std::memcpy(vertexBuffer.value, verticies.data(), verticies.size() * sizeof(Vector<float, 3, 3>));
    std::memcpy(indexBuffer.value, indicies.data(), indicies.size() * sizeof(std::uint32_t));

    vertexBuffer.FlushDevice(cmd, VK_ACCESS_MEMORY_READ_BIT, VK_PIPELINE_STAGE_ACCELERATION_STRUCTURE_BUILD_BIT_KHR);
    vertexBuffer.FlushDevice(cmd, VK_ACCESS_MEMORY_READ_BIT, VK_PIPELINE_STAGE_ACCELERATION_STRUCTURE_BUILD_BIT_KHR);

    VkDeviceOrHostAddressConstKHR vertexAddr;
    vertexAddr.deviceAddress = vertexBuffer.address;

    VkDeviceOrHostAddressConstKHR indexAddr;
    indexAddr.deviceAddress = indexBuffer.address;

    auto trianglesData = VkAccelerationStructureGeometryTrianglesDataKHR {
        .sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_TRIANGLES_DATA_KHR,
        .vertexFormat = VK_FORMAT_R32G32B32_SFLOAT,
        .vertexData = vertexAddr,
        .vertexStride = sizeof(Vector<float, 3, 3>),
        .maxVertex = static_cast<std::uint32_t>(verticies.size())-1,
        .indexType = VK_INDEX_TYPE_UINT32,
        .indexData = indexAddr,
        .transformData = {.deviceAddress = 0}
    };
    VkAccelerationStructureGeometryDataKHR geometryData(trianglesData);
    VkAccelerationStructureGeometryKHR blasGeometry {
        .sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_KHR,
        .geometryType = VK_GEOMETRY_TYPE_TRIANGLES_KHR,
        .geometry     = geometryData,
        .flags        = VK_GEOMETRY_OPAQUE_BIT_KHR
    };

    VkAccelerationStructureBuildGeometryInfoKHR blasBuildGeometryInfo{
        .sType         = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_GEOMETRY_INFO_KHR,
        .type          = VK_ACCELERATION_STRUCTURE_TYPE_BOTTOM_LEVEL_KHR,
        .mode          = VK_BUILD_ACCELERATION_STRUCTURE_MODE_BUILD_KHR,
        .geometryCount = 1,
        .pGeometries   = &blasGeometry,
    };

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

    VkAccelerationStructureBuildSizesInfoKHR blasBuildSizes = {
        .sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_SIZES_INFO_KHR
    };

    VulkanDevice::vkGetAccelerationStructureBuildSizesKHR(
        VulkanDevice::device,
        VK_ACCELERATION_STRUCTURE_BUILD_TYPE_DEVICE_KHR,
        &blasBuildGeometryInfo,
        &primitiveCount,
        &blasBuildSizes
    );

    scratchBuffer.Resize(VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, blasBuildSizes.buildScratchSize);

    VkPhysicalDeviceAccelerationStructurePropertiesKHR asProps{
        .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_ACCELERATION_STRUCTURE_PROPERTIES_KHR
    };

    VkDeviceAddress scratchAddr = scratchBuffer.address;
    blasBuildGeometryInfo.scratchData.deviceAddress = scratchAddr;

    blasBuffer.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, blasBuildSizes.accelerationStructureSize);

    // Create and store the BLAS handle
    VkAccelerationStructureCreateInfoKHR blasCreateInfo{
        .sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_CREATE_INFO_KHR,
        .buffer = blasBuffer.buffer,
        .offset = 0,
        .size   = blasBuildSizes.accelerationStructureSize,
        .type   = VK_ACCELERATION_STRUCTURE_TYPE_BOTTOM_LEVEL_KHR,
    };


    VulkanDevice::CheckVkResult(VulkanDevice::vkCreateAccelerationStructureKHR(VulkanDevice::device, &blasCreateInfo, nullptr, &accelerationStructure));
    blasBuildGeometryInfo.dstAccelerationStructure = accelerationStructure;

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

    VkAccelerationStructureBuildRangeInfoKHR* blasRangeInfoPP = &blasRangeInfo;
    VulkanDevice::vkCmdBuildAccelerationStructuresKHR(cmd, 1, &blasBuildGeometryInfo, &blasRangeInfoPP);

    VkAccelerationStructureDeviceAddressInfoKHR addrInfo {
        .sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_DEVICE_ADDRESS_INFO_KHR,
        .accelerationStructure = accelerationStructure
    };
    blasAddr = VulkanDevice::vkGetAccelerationStructureDeviceAddressKHR(VulkanDevice::device, &addrInfo);
}