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

#include "device.h"

#include "context.h"
#include "physical_device.h"

#include <EASTL/array.h>
#include <EASTL/fixed_vector.h>

// TODO: This will need to be flexible for devices that don't support some of the extensions.
constexpr eastl::array DEVICE_EXTENSIONS = {
    VK_KHR_SWAPCHAIN_EXTENSION_NAME,
};

Device::Device(const Context *context, PhysicalDevice *physicalDevice, Features *enabledFeatures,
               const eastl::vector<QueueAllocation> &queueAllocations, NameString &&name)
    : m_Name(std::move(name))
    , m_PhysicalDevice(physicalDevice->m_PhysicalDevice)
{
    // Shouldn't have more than 4 deviceQueueFamilies in use anyway. Else we can heap
    eastl::fixed_vector<vk::DeviceQueueCreateInfo, 4> deviceQueueCreateInfos;
    deviceQueueCreateInfos.reserve(queueAllocations.size());

    u32 numPriorities = 0;
    for (auto [_, count] : queueAllocations)
    {
        numPriorities = std::max(count, numPriorities);
    }
    // Shouldn't have more than 4 queues either.
    eastl::fixed_vector<f32, 4> priorities(numPriorities, 1.0f);

    for (auto [family, count] : queueAllocations)
    {
        deviceQueueCreateInfos.push_back({
            .queueFamilyIndex = family,
            .queueCount = count,
            .pQueuePriorities = priorities.data(),
        });
    }

    vk::PhysicalDeviceFeatures *deviceFeatures = &enabledFeatures->m_Vulkan10Features;
    vk::PhysicalDeviceVulkan11Features *vulkan11Features = &enabledFeatures->m_Vulkan11Features;
    vk::PhysicalDeviceVulkan12Features *vulkan12Features = &enabledFeatures->m_Vulkan12Features;
    vk::PhysicalDeviceVulkan13Features *vulkan13Features = &enabledFeatures->m_Vulkan13Features;

    vulkan11Features->pNext = vulkan12Features;
    vulkan12Features->pNext = vulkan13Features;

    vk::DeviceCreateInfo deviceCreateInfo = {
        .pNext = vulkan11Features,
        .queueCreateInfoCount = Cast<u32>(deviceQueueCreateInfos.size()),
        .pQueueCreateInfos = deviceQueueCreateInfos.data(),
        .enabledExtensionCount = Cast<u32>(DEVICE_EXTENSIONS.size()),
        .ppEnabledExtensionNames = DEVICE_EXTENSIONS.data(),
        .pEnabledFeatures = deviceFeatures,
    };

    vk::Result result = m_PhysicalDevice.createDevice(&deviceCreateInfo, nullptr, &m_Device);
    ERROR_IF(Failed(result), "Could not initialize Vulkan Device. Cause: {}", result)
    THEN_ABORT(result)
    ELSE_DEBUG("{} ({}) Initialized.", m_Name, physicalDevice->m_DeviceProperties.deviceName.data());

    VmaVulkanFunctions vmaVulkanFunctions = {
        .vkGetInstanceProcAddr = vk::defaultDispatchLoaderDynamic.vkGetInstanceProcAddr,
        .vkGetDeviceProcAddr = vk::defaultDispatchLoaderDynamic.vkGetDeviceProcAddr,
    };

    const VmaAllocatorCreateInfo allocatorCreateInfo = {
        .physicalDevice = m_PhysicalDevice,
        .device = m_Device,
        .pVulkanFunctions = &vmaVulkanFunctions,
        .instance = context->m_Instance,
        .vulkanApiVersion = ASTER_API_VERSION,
    };

    result = Cast<vk::Result>(vmaCreateAllocator(&allocatorCreateInfo, &m_Allocator));
    ERROR_IF(Failed(result), "Memory allocator creation failed. Cause: {}", result)
    DO(m_Device.destroy(nullptr))
    THEN_ABORT(result)
    ELSE_VERBOSE("Memory Allocator Created");

    DEBUG("Created '{}' Successfully", m_Name);
}

Device::~Device()
{
    if (m_Allocator)
    {
        vmaDestroyAllocator(m_Allocator);
        m_Allocator = nullptr;
        DEBUG("Memory Allocator Destroyed");
    }
    m_Device.destroy(nullptr);
    DEBUG("Device '{}' Destroyed", m_Name);
    m_PhysicalDevice = nullptr;
}

vk::Queue
Device::GetQueue(const u32 familyIndex, const u32 queueIndex) const
{
    vk::Queue queue;
    m_Device.getQueue(familyIndex, queueIndex, &queue);
    return queue;
}