/*
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 :Mesh;
import :VulkanDevice;
import std;

using namespace Crafter;

constexpr std::size_t alignUp(std::size_t value, std::size_t alignment) {
    return (value + alignment - 1) & ~(alignment - 1);
}

void Mesh::Build(std::span<Vertex> verticies, std::span<std::uint32_t> indicies, VkCommandBuffer cmd) {
    new (&vertexStaging) VulkanBuffer<Vertex>(VK_BUFFER_USAGE_TRANSFER_SRC_BIT | VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT, verticies.size());
    new (&indexStaging) VulkanBuffer<std::uint32_t>(VK_BUFFER_USAGE_TRANSFER_SRC_BIT | VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT, indicies.size());

    std::memcpy(vertexStaging.value, verticies.data(), verticies.size() * sizeof(Vertex));
    std::memcpy(indexStaging.value, indicies.data(), indicies.size() * sizeof(std::uint32_t));

    new (&vertexBuffer) VulkanBuffer<Vertex>(VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT | VK_BUFFER_USAGE_2_ACCELERATION_STRUCTURE_BUILD_INPUT_READ_ONLY_BIT_KHR, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, verticies.size());
    new (&indexBuffer) VulkanBuffer<std::uint32_t>(VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT | VK_BUFFER_USAGE_2_ACCELERATION_STRUCTURE_BUILD_INPUT_READ_ONLY_BIT_KHR, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, indicies.size());

    VkBufferCopy copyRegion = {
        .srcOffset = 0,
        .dstOffset = 0,
        .size = verticies.size() * sizeof(Vertex)
    };

    vkCmdCopyBuffer(
        cmd,
        vertexStaging.buffer,
        vertexBuffer.buffer,
        1,
        &copyRegion
    );

    VkBufferCopy copyRegion2 = {
        .srcOffset = 0,
        .dstOffset = 0,
        .size = indicies.size() * sizeof(std::uint32_t)
    };

    vkCmdCopyBuffer(
        cmd,
        indexStaging.buffer,
        indexBuffer.buffer,
        1,
        &copyRegion2
    );

    VkBufferMemoryBarrier barrier = {
        .sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER,
        .srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT,
        .dstAccessMask = VK_ACCESS_MEMORY_READ_BIT,
        .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
        .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
        .buffer =  vertexBuffer.buffer,
        .offset = 0,
        .size = VK_WHOLE_SIZE
    };

    vkCmdPipelineBarrier(
        cmd,
        VK_PIPELINE_STAGE_TRANSFER_BIT,
        VK_PIPELINE_STAGE_ACCELERATION_STRUCTURE_BUILD_BIT_KHR,
        0,
        0, NULL,
        1, &barrier,
        0, NULL
    );

    VkBufferMemoryBarrier barrier2 = {
        .sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER,
        .srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT,
        .dstAccessMask = VK_ACCESS_MEMORY_READ_BIT,
        .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
        .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
        .buffer =  indexStaging.buffer,
        .offset = 0,
        .size = VK_WHOLE_SIZE
    };

    vkCmdPipelineBarrier(
        cmd,
        VK_PIPELINE_STAGE_TRANSFER_BIT,
        VK_PIPELINE_STAGE_ACCELERATION_STRUCTURE_BUILD_BIT_KHR,
        0,
        0, NULL,
        1, &barrier2,
        0, NULL
    );


    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(Vertex),
        .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
    );

    new (&scratchBuffer) VulkanBuffer<char>(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;

    new (&blasBuffer) VulkanBuffer<char>(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,
    };

    VkAccelerationStructureKHR aclereationStructure {};

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

    // 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);
}