// =============================================
//  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>

constexpr eastl::array DEVICE_EXTENSIONS = {VK_KHR_SWAPCHAIN_EXTENSION_NAME};

Device::Device(const Context *context, PhysicalDevice *physicalDevice,
               const vk::PhysicalDeviceFeatures *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::DeviceCreateInfo deviceCreateInfo = {
        .queueCreateInfoCount = cast<u32>(deviceQueueCreateInfos.size()),
        .pQueueCreateInfos = deviceQueueCreateInfos.data(),
        .enabledExtensionCount = cast<u32>(DEVICE_EXTENSIONS.size()),
        .ppEnabledExtensionNames = DEVICE_EXTENSIONS.data(),
        .pEnabledFeatures = enabledFeatures,
    };

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