/* Copyright (c) 2025-2026 Valve Corporation
 * Copyright (c) 2025-2026 LunarG, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "ray_tracing_utils.h"

#include "generated/dispatch_functions.h"
#include "containers/span.h"

#include <cassert>
#include <vector>
#include <vulkan/utility/vk_struct_helper.hpp>

namespace rt {

static VkAccelerationStructureBuildSizesInfoKHR ComputeBuildSizes(const VkDevice device,
                                                                  const VkAccelerationStructureBuildTypeKHR build_type,
                                                                  const VkAccelerationStructureBuildGeometryInfoKHR& build_info,
                                                                  const VkAccelerationStructureBuildRangeInfoKHR* range_infos) {
    std::vector<uint32_t> primitive_counts(build_info.geometryCount);
    for (const auto [i, build_range] : vvl::enumerate(range_infos, build_info.geometryCount)) {
        primitive_counts[i] = build_range.primitiveCount;
    }
    VkAccelerationStructureBuildSizesInfoKHR size_info = vku::InitStructHelper();
    DispatchGetAccelerationStructureBuildSizesKHR(device, build_type, &build_info, primitive_counts.data(), &size_info);

    return size_info;
}

VkDeviceSize ComputeScratchSize(BuildType build_type, const VkDevice device,
                                const VkAccelerationStructureBuildGeometryInfoKHR& build_info,
                                const VkAccelerationStructureBuildRangeInfoKHR* range_infos) {
    if (!range_infos) {
        // range_infos is null for indirect builds (vkCmdBuildAccelerationStructuresIndirectKHR) or VK_GEOMETRY_TYPE_MICROMAP_KHR
        // but because build ranges are stored in device memory and cannot be accessed during CPU-side validation. In this case, we
        // cannot compute the actual acceleration structure size, so return 0.
        return 0;
    }
    const VkAccelerationStructureBuildSizesInfoKHR size_info =
        ComputeBuildSizes(device,
                          build_type == BuildType::Device ? VK_ACCELERATION_STRUCTURE_BUILD_TYPE_DEVICE_KHR
                                                          : VK_ACCELERATION_STRUCTURE_BUILD_TYPE_HOST_KHR,
                          build_info, range_infos);
    switch (build_info.mode) {
        case VK_BUILD_ACCELERATION_STRUCTURE_MODE_BUILD_KHR:
            return size_info.buildScratchSize;
        case VK_BUILD_ACCELERATION_STRUCTURE_MODE_UPDATE_KHR:
            return size_info.updateScratchSize;
        default:
            assert(false);
            return static_cast<VkDeviceSize>(0);
            break;
    }
}

VkDeviceSize ComputeAccelerationStructureSize(BuildType build_type, const VkDevice device,
                                              const VkAccelerationStructureBuildGeometryInfoKHR& build_info,
                                              const VkAccelerationStructureBuildRangeInfoKHR* range_infos) {
    if (!range_infos) {
        // range_infos is null for indirect builds (vkCmdBuildAccelerationStructuresIndirectKHR) or VK_GEOMETRY_TYPE_MICROMAP_KHR
        // but because build ranges are stored in device memory and cannot be accessed during CPU-side validation. In this case, we
        // cannot compute the actual acceleration structure size, so return 0.
        return 0;
    }

    const VkAccelerationStructureBuildSizesInfoKHR size_info =
        ComputeBuildSizes(device,
                          build_type == BuildType::Device ? VK_ACCELERATION_STRUCTURE_BUILD_TYPE_DEVICE_KHR
                                                          : VK_ACCELERATION_STRUCTURE_BUILD_TYPE_HOST_KHR,
                          build_info, range_infos);
    return size_info.accelerationStructureSize;
}

uint64_t MicromapUsageTotalTriangleCount(const VkAccelerationStructureBuildGeometryInfoKHR& build_info) {
    if (build_info.geometryCount == 0) {
        return 0;
    }

    const VkAccelerationStructureGeometryKHR& geometry = GetGeometry(build_info, 0);
    const auto* micromap_data = vku::FindStructInPNextChain<VkAccelerationStructureGeometryMicromapDataKHR>(geometry.pNext);
    if (!micromap_data) {
        return 0;
    }

    uint64_t total = 0;
    for (uint32_t usage_i = 0; usage_i < micromap_data->usageCountsCount; ++usage_i) {
        total += GetMicroMapUsage(*micromap_data, usage_i).count;
    }
    return total;
}

}  // namespace rt
