/*
	VulkanFramework
	by John Ryland
	Copyright (c) 2023
*/

////////////////////////////////////////////////////////////////////////////////////
//	Vulkan Device

#include "VulkanDevice.h"
#include <stdexcept>
#include <cstring>
#include <cstdio>
#include <cassert>
#include <memory>

//#ifdef __APPLE__
//#include "MoltenVK/vk_mvk_moltenvk.h"
//#endif

#ifdef NDEBUG
static const bool g_enableDebugReporting = false;
static const bool g_enableValidationLayers = false;
#else
static const bool g_enableDebugReporting = true;
static const bool g_enableValidationLayers = true;
#endif

namespace Vulkan {

Device::Device()
    : m_instance(nullptr)
    , m_device(nullptr)
    , m_descriptorPool(nullptr)
    , m_descriptorSetLayout(nullptr)
    , m_commandPool(nullptr)
    , m_debugReport(nullptr)
    , m_allocator(nullptr)
    , m_graphicsQueueFamily(0)
    , m_graphicsQueue(nullptr)
    , m_physicalDevice(nullptr)
{
}

// virtual
Device::~Device()
{
    Destroy();
}

// virtual
void Device::Create(const char** extensions, uint32_t extensionsCount)
{
    CreateInstance(extensions, extensionsCount, g_enableDebugReporting, g_enableValidationLayers);
    SelectGPU();
    DumpPhysicalDeviceExtensions();
    SelectGraphicsQueueFamily();
    CreateLogicalDevice();
    vkGetDeviceQueue(m_device, m_graphicsQueueFamily, 0, &m_graphicsQueue);
    CreateDescriptorPool();
    CreateDescriptorSetLayout();
    CreateCommandPool();
}

// virtual
void Device::Destroy()
{
    if (m_debugReport)
    {
        // Remove the debug report callback
        auto vkDestroyDebugReportCallbackEXT = (PFN_vkDestroyDebugReportCallbackEXT)vkGetInstanceProcAddr(m_instance, "vkDestroyDebugReportCallbackEXT");
        vkDestroyDebugReportCallbackEXT(m_instance, m_debugReport, m_allocator);
    }
    if (m_commandPool)
        vkDestroyCommandPool(m_device, m_commandPool, m_allocator);
    if (m_descriptorSetLayout)
        vkDestroyDescriptorSetLayout(m_device, m_descriptorSetLayout, m_allocator);
    if (m_descriptorPool)
        vkDestroyDescriptorPool(m_device, m_descriptorPool, m_allocator);
    if (m_device)
        vkDestroyDevice(m_device, m_allocator);
    if (m_instance)
        vkDestroyInstance(m_instance, m_allocator);
    m_descriptorPool = nullptr;
    m_descriptorSetLayout = nullptr;
    m_commandPool = nullptr;
    m_debugReport = nullptr;
    m_device = nullptr;
    m_instance = nullptr;
}

// virtual
void Device::CheckResult(VkResult error, const char* message)
{
    if (error == VK_SUCCESS) {
        fprintf(stderr, "[vulkan] Success: NOT %s \n", (message) ? message : "-");
        return;
    }
    fprintf(stderr, "[vulkan] Error: VkResult = %d\n", error);
    if (message)
        fprintf(stderr, "[vulkan] Message: %s\n\n", message);
    if (error < 0)
        abort();
}

static VKAPI_ATTR
VkBool32 VKAPI_CALL DebugReport(VkDebugReportFlagsEXT       flags,
                                VkDebugReportObjectTypeEXT  objectType,
                                uint64_t                    object,
                                size_t                      location,
                                int32_t                     messageCode,
                                const char*                 pLayerPrefix,
                                const char*                 pMessage,
                                void*                       pUserData)
{
    (void)flags; (void)object; (void)location; (void)messageCode; (void)pUserData; (void)pLayerPrefix; // Unused arguments
    fprintf(stderr, "[vulkan] Debug report from ObjectType: %i\n[vulkan] Message: %s\n\n", objectType, pMessage);
    return VK_FALSE;
}

// virtual
void Device::CreateInstance(const char** extensions, uint32_t extensionsCount, bool enableDebug, bool enableValidationLayers)
{
    const char* const validation_layers[] = {
        "VK_LAYER_KHRONOS_validation"
    };

    VkInstanceCreateInfo create_info = {};
    create_info.sType = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO;

    if (enableValidationLayers)
    {
        // Enabling validation layers
        create_info.enabledLayerCount       = 1;
        create_info.ppEnabledLayerNames     = validation_layers;
    }

#if 1 // def __APPLE__
    //create_info.flags = VK_INSTANCE_CREATE_ENUMERATE_PORTABILITY_BIT_KHR;

    // Enable debug report extension (we need additional storage, so we duplicate the user array to add our new extension to it)
    // extensionsCount = 0;
    auto extensions_ext = std::make_unique<const char *[]>(extensionsCount + 4);
    memcpy(extensions_ext.get(), extensions, extensionsCount * sizeof(const char *));
    //extensions_ext[extensionsCount + 0] = VK_KHR_GET_PHYSICAL_DEVICE_PROPERTIES_2_EXTENSION_NAME;
    //extensions_ext[extensionsCount + 1] = VK_KHR_PORTABILITY_ENUMERATION_EXTENSION_NAME;
    //extensions_ext[extensionsCount + 2] = "VK_EXT_metal_surface";
    extensions_ext[extensionsCount + 0] = "VK_EXT_debug_report";
    extensions_ext[extensionsCount + 1] = "VK_KHR_portability_enumeration";

    create_info.enabledExtensionCount = extensionsCount + 2;
    create_info.ppEnabledExtensionNames = extensions_ext.get();
    CheckResult(vkCreateInstance(&create_info, m_allocator, &m_instance), "failed to create instance!");

/*
    MVKConfiguration config = {};
    size_t siz = sizeof(config);
    vkGetMoltenVKConfigurationMVK(m_instance, &config, &siz);
    config.apiVersionToAdvertise = 310;
    vkSetMoltenVKConfigurationMVK(m_instance, &config, &siz);
*/

    // Get the function pointer (required for any extensions)
    auto vkCreateDebugReportCallbackEXT = (PFN_vkCreateDebugReportCallbackEXT)vkGetInstanceProcAddr(m_instance, "vkCreateDebugReportCallbackEXT");
    assert(vkCreateDebugReportCallbackEXT != nullptr);

    // Setup the debug report callback
    VkDebugReportCallbackCreateInfoEXT debug_report_ci = {};
    debug_report_ci.sType = VK_STRUCTURE_TYPE_DEBUG_REPORT_CALLBACK_CREATE_INFO_EXT;
    debug_report_ci.flags = VK_DEBUG_REPORT_ERROR_BIT_EXT | VK_DEBUG_REPORT_WARNING_BIT_EXT | VK_DEBUG_REPORT_PERFORMANCE_WARNING_BIT_EXT
        | VK_INSTANCE_CREATE_ENUMERATE_PORTABILITY_BIT_KHR;
    debug_report_ci.pfnCallback = DebugReport;
    debug_report_ci.pUserData = nullptr;
    CheckResult(vkCreateDebugReportCallbackEXT(m_instance, &debug_report_ci, m_allocator, &m_debugReport), "failed to register debug report callback!");
#else
    if (enableDebug)
    {
        // Enable debug report extension (we need additional storage, so we duplicate the user array to add our new extension to it)
        auto extensions_ext = std::make_unique<const char*[]>(extensionsCount + 1);
        memcpy(extensions_ext.get(), extensions, extensionsCount * sizeof(const char*));
        extensions_ext[extensionsCount]     = "VK_EXT_debug_report";
        extensions_ext[extensionsCount + 1] = "VK_KHR_portability_enumeration";

        create_info.enabledExtensionCount   = extensionsCount + 2;
        create_info.ppEnabledExtensionNames = extensions_ext.get();
        CheckResult(vkCreateInstance(&create_info, m_allocator, &m_instance), "failed to create instance!");

        // Get the function pointer (required for any extensions)
        auto vkCreateDebugReportCallbackEXT = (PFN_vkCreateDebugReportCallbackEXT)vkGetInstanceProcAddr(m_instance, "vkCreateDebugReportCallbackEXT");
        assert(vkCreateDebugReportCallbackEXT != nullptr);

        // Setup the debug report callback
        VkDebugReportCallbackCreateInfoEXT debug_report_ci = {};
        debug_report_ci.sType               = VK_STRUCTURE_TYPE_DEBUG_REPORT_CALLBACK_CREATE_INFO_EXT;
        debug_report_ci.flags               = VK_DEBUG_REPORT_ERROR_BIT_EXT | VK_DEBUG_REPORT_WARNING_BIT_EXT | VK_DEBUG_REPORT_PERFORMANCE_WARNING_BIT_EXT;
        debug_report_ci.pfnCallback         = DebugReport;
        debug_report_ci.pUserData           = nullptr;
        CheckResult(vkCreateDebugReportCallbackEXT(m_instance, &debug_report_ci, m_allocator, &m_debugReport), "failed to register debug report callback!");
    }
    else
    {
        auto extensions_ext = std::make_unique<const char*[]>(extensionsCount + 1);
        memcpy(extensions_ext.get(), extensions, extensionsCount * sizeof(const char*));
        extensions_ext[extensionsCount]     = "VK_KHR_portability_enumeration";
        create_info.enabledExtensionCount   = extensionsCount + 1;
        create_info.ppEnabledExtensionNames = extensions_ext.get();
        CheckResult(vkCreateInstance(&create_info, m_allocator, &m_instance), "failed to create instance!");
    }
#endif
}

// virtual
void Device::SelectGPU()
{
    uint32_t count;
    CheckResult(vkEnumeratePhysicalDevices(m_instance, &count, nullptr));
    CheckResult(count > 0 ? VK_SUCCESS : VK_ERROR_UNKNOWN, "no Vulkan supported devices found!"); // Ensure there is atleast one available
    auto gpus = std::make_unique<VkPhysicalDevice[]>(count);
    CheckResult(vkEnumeratePhysicalDevices(m_instance, &count, gpus.get()));
    for (uint32_t i = 0; i < count; i++)
    {
        VkPhysicalDeviceProperties properties;
        vkGetPhysicalDeviceProperties(gpus[i], &properties);
        if (properties.deviceType == VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU)
        {
            m_physicalDevice = gpus[i]; // Use the 1st discrete gpu if available
            return;
        }
    }
    m_physicalDevice = gpus[0]; // Default to first gpu if no discrete gpus
}

// virtual
void Device::DumpPhysicalDeviceExtensions()
{
    uint32_t extensionCount = 0;
    const char* layerName = nullptr;
    CheckResult(vkEnumerateDeviceExtensionProperties(m_physicalDevice, layerName, &extensionCount, nullptr));
    auto extensions = std::make_unique<VkExtensionProperties[]>(extensionCount);
    CheckResult(vkEnumerateDeviceExtensionProperties(m_physicalDevice, layerName, &extensionCount, extensions.get()));
    printf("Extensions:\n");
    for (uint32_t i = 0; i < extensionCount; i++)
    {
        printf(" - %s\n", extensions[i].extensionName);
    }
}

// virtual
void Device::SelectGraphicsQueueFamily()
{
    uint32_t queueFamilyCount = 0;
    vkGetPhysicalDeviceQueueFamilyProperties(m_physicalDevice, &queueFamilyCount, nullptr);
    auto queueFamilies = std::make_unique<VkQueueFamilyProperties[]>(queueFamilyCount);
    vkGetPhysicalDeviceQueueFamilyProperties(m_physicalDevice, &queueFamilyCount, queueFamilies.get());
    for (uint32_t i = 0; i < queueFamilyCount; i++)
    {
        if (queueFamilies[i].queueFlags & VK_QUEUE_GRAPHICS_BIT)
        {
            m_graphicsQueueFamily = i;
            return;
        }
    }
    m_graphicsQueueFamily = (uint32_t)-1;
    CheckResult(VK_ERROR_UNKNOWN, "no graphics queue family found!");
}

// virtual
void Device::CreateLogicalDevice()
{
    const char* const device_extensions[] = {
        VK_KHR_SWAPCHAIN_EXTENSION_NAME,
        "VK_KHR_portability_subset"
    };

    const float queue_priority[] = { 1.0f };
    VkDeviceQueueCreateInfo queue_info[1] = {};
    queue_info[0].sType = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO;
    queue_info[0].queueFamilyIndex = m_graphicsQueueFamily;
    queue_info[0].queueCount = 1;
    queue_info[0].pQueuePriorities = queue_priority;

    VkDeviceCreateInfo create_info = {};
    create_info.sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO;
    create_info.queueCreateInfoCount = sizeof(queue_info) / sizeof(queue_info[0]);
    create_info.pQueueCreateInfos = queue_info;
    create_info.enabledExtensionCount = 2;
    create_info.ppEnabledExtensionNames = device_extensions;
    CheckResult(vkCreateDevice(m_physicalDevice, &create_info, m_allocator, &m_device), "failed to create logical device!");
}

// virtual
void Device::CreateDescriptorPool()
{
    static const uint32_t s_defaultPoolSets = 1000;
    static const uint32_t s_defaultPoolSize = 1000;

    VkDescriptorPoolSize pool_sizes[] =
    {
        { VK_DESCRIPTOR_TYPE_SAMPLER,                s_defaultPoolSize },
        { VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, s_defaultPoolSize },
        { VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE,          s_defaultPoolSize },
        { VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,          s_defaultPoolSize },
        { VK_DESCRIPTOR_TYPE_UNIFORM_TEXEL_BUFFER,   s_defaultPoolSize },
        { VK_DESCRIPTOR_TYPE_STORAGE_TEXEL_BUFFER,   s_defaultPoolSize },
        { VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,         s_defaultPoolSize },
        { VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,         s_defaultPoolSize },
        { VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER_DYNAMIC, s_defaultPoolSize },
        { VK_DESCRIPTOR_TYPE_STORAGE_BUFFER_DYNAMIC, s_defaultPoolSize },
        { VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT,       s_defaultPoolSize }
    };
    VkDescriptorPoolCreateInfo pool_info = {};
    pool_info.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;
    pool_info.flags = VK_DESCRIPTOR_POOL_CREATE_FREE_DESCRIPTOR_SET_BIT;
    pool_info.maxSets = s_defaultPoolSets * std::size(pool_sizes);
    pool_info.poolSizeCount = (uint32_t)std::size(pool_sizes);
    pool_info.pPoolSizes = pool_sizes;
    CheckResult(vkCreateDescriptorPool(m_device, &pool_info, m_allocator, &m_descriptorPool), "failed to create descriptor pool!");
}

// virtual
void Device::CreateDescriptorSetLayout()
{
    VkDescriptorSetLayoutBinding uboLayoutBinding {};
    uboLayoutBinding.binding            = 0;
    uboLayoutBinding.descriptorCount    = 1;
    uboLayoutBinding.descriptorType     = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
    uboLayoutBinding.pImmutableSamplers = nullptr;
    uboLayoutBinding.stageFlags         = VK_SHADER_STAGE_VERTEX_BIT;

    VkDescriptorSetLayoutBinding samplerLayoutBinding {};
    samplerLayoutBinding.binding            = 1;
    samplerLayoutBinding.descriptorCount    = 1;
    samplerLayoutBinding.descriptorType     = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
    samplerLayoutBinding.pImmutableSamplers = nullptr;
    samplerLayoutBinding.stageFlags         = VK_SHADER_STAGE_FRAGMENT_BIT;

    VkDescriptorSetLayoutBinding bindings[2] = { uboLayoutBinding, samplerLayoutBinding };
    VkDescriptorSetLayoutCreateInfo layoutInfo {};
    layoutInfo.sType        = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
    layoutInfo.bindingCount = static_cast<uint32_t>(std::size(bindings));
    layoutInfo.pBindings    = bindings;
    CheckResult(vkCreateDescriptorSetLayout(m_device, &layoutInfo, m_allocator, &m_descriptorSetLayout), "failed to create descriptor set layout!");
}

// virtual
void Device::CreateCommandPool()
{
    VkCommandPoolCreateInfo poolInfo {};
    poolInfo.sType            = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO;
    poolInfo.flags            = VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT;
    poolInfo.queueFamilyIndex = m_graphicsQueueFamily;
    CheckResult(vkCreateCommandPool(m_device, &poolInfo, m_allocator, &m_commandPool), "failed to create graphics command pool!");
}

// virtual
uint32_t Device::FindMemoryType(uint32_t              typeFilter,
                                VkMemoryPropertyFlags properties)
{
    VkPhysicalDeviceMemoryProperties memProperties;
    vkGetPhysicalDeviceMemoryProperties(m_physicalDevice, &memProperties);
    for (uint32_t i = 0; i < memProperties.memoryTypeCount; i++)
    {
        if ((typeFilter & (1 << i)) && (memProperties.memoryTypes[i].propertyFlags & properties) == properties)
            return i;
    }
    CheckResult(VK_ERROR_UNKNOWN, "failed to find suitable memory type!");
    return (uint32_t)-1;
}

// virtual
VkFormat Device::FindFormat(VkFormat* candidateFormats, size_t size, VkFormatFeatureFlags properties)
{
    for (uint32_t i = 0; i < size; ++i)
    {
        VkFormatProperties props;
        vkGetPhysicalDeviceFormatProperties(m_physicalDevice, candidateFormats[i], &props);
        if ((props.optimalTilingFeatures & properties) == properties)
            return candidateFormats[i];
    }
    CheckResult(VK_ERROR_UNKNOWN, "failed to find supported format!");
    return VK_FORMAT_UNDEFINED;
}

Device::SingleTimeCommand Device::BeginSingleTimeCommands()
{
    VkCommandBuffer commandBuffer;

    VkCommandBufferAllocateInfo allocInfo {};
    allocInfo.sType              = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
    allocInfo.level              = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
    allocInfo.commandPool        = m_commandPool;
    allocInfo.commandBufferCount = 1;
    vkAllocateCommandBuffers(m_device, &allocInfo, &commandBuffer);

    VkCommandBufferBeginInfo beginInfo {};
    beginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
    beginInfo.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT;
    vkBeginCommandBuffer(commandBuffer, &beginInfo);

    return SingleTimeCommand{ this, commandBuffer };
}

Device::SingleTimeCommand::~SingleTimeCommand()
{
    vkEndCommandBuffer(m_commandBuffer);

    VkSubmitInfo submitInfo {};
    submitInfo.sType              = VK_STRUCTURE_TYPE_SUBMIT_INFO;
    submitInfo.commandBufferCount = 1;
    submitInfo.pCommandBuffers    = &m_commandBuffer;
    vkQueueSubmit(m_owner->m_graphicsQueue, 1, &submitInfo, VK_NULL_HANDLE);

    vkQueueWaitIdle(m_owner->m_graphicsQueue);

    vkFreeCommandBuffers(m_owner->m_device, m_owner->m_commandPool, 1, &m_commandBuffer);
}

} // Vulkan namespace