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

#include "device.h"

#include "context.h"
#include "physical_device.h"
#include "queue_allocation.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,
};

void
Device::Init(const Context *context, PhysicalDevice *physicalDevice, Features *enabledFeatures,
               const eastl::vector<QueueAllocation> &queueAllocations, cstr name)
{
    this->Init(context, physicalDevice, enabledFeatures, queueAllocations, {}, name);
}

void
Device::Init(const Context *context, PhysicalDevice *physicalDevice, Features *enabledFeatures,
               const eastl::vector<QueueAllocation> &queueAllocations, const eastl::span<u8> &pipelineCacheData,
               cstr name)
{
    assert(!m_Device);

    m_PhysicalDevice = physicalDevice->m_PhysicalDevice;
    m_ValidationEnabled = context->IsInDebugMode();
    // 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.", name, physicalDevice->m_DeviceProperties.deviceName.data());

    SetName(m_Device, name);

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

    const VmaAllocatorCreateInfo allocatorCreateInfo = {
        .flags = VMA_ALLOCATOR_CREATE_BUFFER_DEVICE_ADDRESS_BIT,
        .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");

    vk::PipelineCacheCreateInfo pipelineCacheCreateInfo = {
        .initialDataSize = pipelineCacheData.size_bytes(),
        .pInitialData = pipelineCacheData.data(),
    };
    result = m_Device.createPipelineCache(&pipelineCacheCreateInfo, nullptr, &m_PipelineCache);
    ERROR_IF(Failed(result), "Pipeline Cache creation failed. Cause: {}", result)
    DO(m_Device.destroy(nullptr))
    THEN_ABORT(result)
    ELSE_VERBOSE("Pipeline Cache created.");

    DEBUG("Created Device '{}' Successfully", name);
}

void
Device::Destroy()
{
    if (!m_Device)
        return;
    
    m_Device.destroy(m_PipelineCache, nullptr);
    if (m_Allocator)
    {
        vmaDestroyAllocator(m_Allocator);
        m_Allocator = nullptr;
        DEBUG("Memory Allocator Destroyed");
    }
    m_Device.destroy(nullptr);
    DEBUG("Device Destroyed");

    m_Allocator = nullptr;
    m_PhysicalDevice = nullptr;
    m_Device = nullptr;
}

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

eastl::vector<u8>
Device::DumpPipelineCache() const
{
    usize pipelineCacheSize = 0;
    vk::Result result = m_Device.getPipelineCacheData(m_PipelineCache, &pipelineCacheSize, nullptr);
    ERROR_IF(Failed(result), "Pipeline Cache data fetch failed. Cause: {}", result);
    eastl::vector<u8> pipelineCacheData(pipelineCacheSize);
    result = m_Device.getPipelineCacheData(m_PipelineCache, &pipelineCacheSize, pipelineCacheData.data());
    ERROR_IF(Failed(result), "Pipeline Cache data fetch failed. Cause: {}", result);

    return pipelineCacheData;
}

void
Device::WaitIdle() const
{
    auto deviceIdleResult = m_Device.waitIdle();
    ERROR_IF(Failed(deviceIdleResult), "Device Idling Failed. Cause: {}", deviceIdleResult)
    THEN_ABORT(deviceIdleResult);
}