// =============================================
//  Aster: physical_device.cpp
//  Copyright (c) 2020-2024 Anish Bhobe
// =============================================

#include "physical_device.h"

#include "context.h"
#include "window.h"

[[nodiscard]] vk::SurfaceCapabilitiesKHR
GetSurfaceCapabilities(const vk::PhysicalDevice physicalDevice, const vk::SurfaceKHR surface)
{
    vk::SurfaceCapabilitiesKHR surfaceCapabilities;

    vk::Result result = physicalDevice.getSurfaceCapabilitiesKHR(surface, &surfaceCapabilities);
    ERROR_IF(failed(result), "Fetching surface capabilities failed. Cause: {}", result)
    THEN_ABORT(result);

    return surfaceCapabilities;
}

[[nodiscard]] eastl::vector<vk::SurfaceFormatKHR>
GetSurfaceFormats(const vk::PhysicalDevice physicalDevice, const vk::SurfaceKHR surface)
{
    // vk::Result::eIncomplete should not occur in this function. The rest are errors. Thus, abort is allowed.
    u32 count = 0;
    vk::Result result = physicalDevice.getSurfaceFormatsKHR(surface, &count, nullptr);
    ERROR_IF(failed(result), "Could not get surface formats. Cause: {}", result)
    THEN_ABORT(result);

    eastl::vector<vk::SurfaceFormatKHR> surfaceFormats(count);
    result = physicalDevice.getSurfaceFormatsKHR(surface, &count, surfaceFormats.data());
    ERROR_IF(failed(result), "Could not get surface formats. Cause: {}", result)
    THEN_ABORT(result);

    return surfaceFormats;
}

[[nodiscard]] eastl::vector<vk::PresentModeKHR>
GetSurfacePresentModes(const vk::PhysicalDevice physicalDevice, const vk::SurfaceKHR surface)
{
    // vk::Result::eIncomplete should not occur in this function. The rest are errors. Thus, abort is allowed.
    u32 count = 0;
    vk::Result result = physicalDevice.getSurfacePresentModesKHR(surface, &count, nullptr);
    ERROR_IF(failed(result), "Could not get present modes. Cause: {}", result)
    THEN_ABORT(result);

    eastl::vector<vk::PresentModeKHR> presentModes(count);
    result = physicalDevice.getSurfacePresentModesKHR(surface, &count, presentModes.data());
    ERROR_IF(failed(result), "Could not get present modes. Cause: {}", result)
    THEN_ABORT(result);

    return presentModes;
}

[[nodiscard]] bool
GetQueuePresentSupport(const u32 queueFamilyIndex, vk::SurfaceKHR surface, const vk::PhysicalDevice physicalDevice)
{
    b32 supported = false;
    const vk::Result result = physicalDevice.getSurfaceSupportKHR(queueFamilyIndex, surface, &supported);
    ERROR_IF(failed(result), "Could not get queue family surface support. Cause: {}", result)
    THEN_ABORT(result);

    return supported;
}

[[nodiscard]] eastl::fixed_vector<vk::QueueFamilyProperties, 32>
GetQueueFamilyProperties(const vk::PhysicalDevice physicalDevice)
{
    // Devices rarely have more than 32 queue families. Thus fixed vector
    u32 count = 0;
    physicalDevice.getQueueFamilyProperties(&count, nullptr);

    eastl::fixed_vector<vk::QueueFamilyProperties, 32> queueFamilyProperties(count);
    physicalDevice.getQueueFamilyProperties(&count, queueFamilyProperties.data());

    return queueFamilyProperties;
}

// Size 384 return.
[[nodiscard]] eastl::vector<QueueFamilyInfo>
GetQueueFamilies(const vk::SurfaceKHR surface, const vk::PhysicalDevice physicalDevice)
{

    auto queueFamilyProperties = GetQueueFamilyProperties(physicalDevice);

    eastl::vector<QueueFamilyInfo> queueFamilyInfos;
    queueFamilyInfos.reserve(queueFamilyProperties.size());

    u32 familyIndex = 0;
    for (auto qfp : queueFamilyProperties)
    {
        QueueSupportFlags support = {};

        if (qfp.queueFlags | vk::QueueFlagBits::eGraphics)
        {
            support |= QueueSupportFlagBits::eGraphics;
        }

        if (qfp.queueFlags | vk::QueueFlagBits::eTransfer)
        {
            support |= QueueSupportFlagBits::eTransfer;
        }

        if (qfp.queueFlags | vk::QueueFlagBits::eCompute)
        {
            support |= QueueSupportFlagBits::eCompute;
        }

        if (GetQueuePresentSupport(familyIndex, surface, physicalDevice))
        {
            support |= QueueSupportFlagBits::ePresent;
        }

        queueFamilyInfos.push_back({
            .m_Index = familyIndex,
            .m_Count = qfp.queueCount,
            .m_Support = support,
        });

        familyIndex++;
    }

    return queueFamilyInfos;
}

PhysicalDevice::PhysicalDevice(const vk::SurfaceKHR surface, const vk::PhysicalDevice physicalDevice)
{
    physicalDevice.getProperties(&m_DeviceProperties);
    physicalDevice.getFeatures(&m_DeviceFeatures);

    m_SurfaceFormats = GetSurfaceFormats(physicalDevice, surface);
    m_PresentModes = GetSurfacePresentModes(physicalDevice, surface);
    m_QueueFamilies = GetQueueFamilies(surface, physicalDevice);

    m_PhysicalDevice = physicalDevice;
}

eastl::fixed_vector<vk::PhysicalDevice, 8>
EnumeratePhysicalDevices(const vk::Instance instance)
{
    u32 count = 0;
    vk::Result result = instance.enumeratePhysicalDevices(&count, nullptr);
    ERROR_IF(failed(result), "Could not fetch vulkan devices. Cause: {}", result)
    THEN_ABORT(result);

    eastl::fixed_vector<vk::PhysicalDevice, 8> physicalDevices(count);
    result = instance.enumeratePhysicalDevices(&count, physicalDevices.data());
    ERROR_IF(failed(result), "Could not fetch vulkan devices. Cause: {}", result)
    THEN_ABORT(result);

    return physicalDevices;
}

PhysicalDevices::PhysicalDevices(const Window *window, const Context *context)
{
    auto surface = window->m_Surface;
    auto physicalDevices = EnumeratePhysicalDevices(context->m_Instance);
    for (auto physicalDevice : physicalDevices)
    {
        this->emplace_back(surface, physicalDevice);
    }
}