#include "VulkanDevice.h"
#include "VulkanTools.h"
#include <cassert>
#include <stdexcept>

VULKANBASE_NAMESPACE_BEGIN

VulkanDevice::VulkanDevice()
    : m_commandPool(VK_NULL_HANDLE)
{
}

VulkanDevice::VulkanDevice(VkPhysicalDevice physicalDevice)
{
    /// 检查物理设备是否存在，后续使用log方式记录
    assert(physicalDevice);
    m_physicalDevice = physicalDevice;

    /// 获取设备信息
    vkGetPhysicalDeviceProperties(physicalDevice, &m_properties);
    /// 获取设备支持的功能
    vkGetPhysicalDeviceFeatures(physicalDevice, &m_features);
    /// 获取设备内存信息，用于创建内存
    vkGetPhysicalDeviceMemoryProperties(physicalDevice, &m_memoryProperties);
    /// 队列族信息，用于设备创建时获取设置获取的队列
    uint32_t queueFamilyCount;
    vkGetPhysicalDeviceQueueFamilyProperties(physicalDevice, &queueFamilyCount, nullptr);
    /// 检查设备队列族数量，必须大于0
    assert(queueFamilyCount > 0);
    m_queueFamilyProperties.resize(queueFamilyCount);
    vkGetPhysicalDeviceQueueFamilyProperties(physicalDevice, &queueFamilyCount, m_queueFamilyProperties.data());
}

VulkanDevice::~VulkanDevice()
{
    if (m_commandPool)
    {
        vkDestroyCommandPool(m_logicalDevice, m_commandPool, nullptr);
    }
    if (m_logicalDevice)
    {
        vkDestroyDevice(m_logicalDevice, nullptr);
    }
}

/// @brief 据分配的物理设备创建逻辑设备，同时获取默认队列族索引
/// @param enabledFeatures 在创建设备时启用某些功能
/// @param enabledExtensions 在创建设备时启用某些扩展
/// @param requestedQueueTypes 指定要从设备请求的队列类型
/// @return 逻辑设备是否成功创建
VkResult VulkanDevice::createLogicalDevice(VkPhysicalDeviceFeatures enabledFeatures, std::vector<const char *> enabledExtensions, VkQueueFlags requestedQueueTypes)
{
    std::vector<VkDeviceQueueCreateInfo> queueCreateInfos{};
    const float defaultQueuePriority(0.0f);
    // Graphics queue
    if (requestedQueueTypes & VK_QUEUE_GRAPHICS_BIT)
    {
        m_queueFamilyIndices.graphics = getQueueFamilyIndex(VK_QUEUE_GRAPHICS_BIT);
        VkDeviceQueueCreateInfo queueInfo{};
        queueInfo.sType = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO;
        queueInfo.queueFamilyIndex = m_queueFamilyIndices.graphics;
        queueInfo.queueCount = 1;
        queueInfo.pQueuePriorities = &defaultQueuePriority;
        queueCreateInfos.push_back(queueInfo);
    }
    else
    {
        m_queueFamilyIndices.graphics = 0;
    }

    // Dedicated compute queue
    if (requestedQueueTypes & VK_QUEUE_COMPUTE_BIT)
    {
        m_queueFamilyIndices.compute = getQueueFamilyIndex(VK_QUEUE_COMPUTE_BIT);
        if (m_queueFamilyIndices.compute != m_queueFamilyIndices.graphics)
        {
            // If compute family index differs, we need an additional queue create info for the compute queue
            VkDeviceQueueCreateInfo queueInfo{};
            queueInfo.sType = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO;
            queueInfo.queueFamilyIndex = m_queueFamilyIndices.compute;
            queueInfo.queueCount = 1;
            queueInfo.pQueuePriorities = &defaultQueuePriority;
            queueCreateInfos.push_back(queueInfo);
        }
    }
    else
    {
        m_queueFamilyIndices.compute = m_queueFamilyIndices.graphics;
    }

    // Create the logical device representation
    std::vector<const char *> deviceExtensions(enabledExtensions);
    deviceExtensions.push_back(VK_KHR_SWAPCHAIN_EXTENSION_NAME);

    VkDeviceCreateInfo deviceCreateInfo = {};
    deviceCreateInfo.sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO;
    deviceCreateInfo.queueCreateInfoCount = static_cast<uint32_t>(queueCreateInfos.size());
    ;
    deviceCreateInfo.pQueueCreateInfos = queueCreateInfos.data();
    deviceCreateInfo.pEnabledFeatures = &enabledFeatures;

    if (deviceExtensions.size() > 0)
    {
        deviceCreateInfo.enabledExtensionCount = (uint32_t)deviceExtensions.size();
        deviceCreateInfo.ppEnabledExtensionNames = deviceExtensions.data();
    }

    VkResult result = vkCreateDevice(m_physicalDevice, &deviceCreateInfo, nullptr, &m_logicalDevice);

    if (result == VK_SUCCESS)
    {
        m_commandPool = createCommandPool(m_queueFamilyIndices.graphics);
    }

    m_enabledFeatures = enabledFeatures;

    return result;
}

VkDevice VulkanDevice::getLogicalDevice()
{
    return m_logicalDevice;
}

VkPhysicalDevice VulkanDevice::getPhysicalDevice()
{
    return m_physicalDevice;
}

VkPhysicalDeviceProperties VulkanDevice::VulkanDevice::getPhysicalDeviceProperties()
{
    return m_properties;
}
VkPhysicalDeviceFeatures VulkanDevice::getPhysicalDeviceFeatures()
{
    return m_features;
}
VkPhysicalDeviceMemoryProperties VulkanDevice::getPhysicalDeviceMemoryProperties()
{
    return m_memoryProperties;
}
std::vector<VkQueueFamilyProperties> VulkanDevice::getQueueFamilyProperties()
{
    return m_queueFamilyProperties;
}
VkCommandPool VulkanDevice::getCommandPool()
{
    return m_commandPool;
}

uint32_t VulkanDevice::getMemoryType(uint32_t typeBits, VkMemoryPropertyFlags properties, VkBool32 *memTypeFound)
{
    for (uint32_t i = 0; i < m_memoryProperties.memoryTypeCount; i++)
    {
        if ((typeBits & 1) == 1)
        {
            if ((m_memoryProperties.memoryTypes[i].propertyFlags & properties) == properties)
            {
                if (memTypeFound)
                {
                    *memTypeFound = true;
                }
                return i;
            }
        }
        typeBits >>= 1;
    }

    if (memTypeFound)
    {
        *memTypeFound = false;
        return 0;
    }
    else
    {
        throw std::runtime_error("Could not find a matching memory type");
    }
}

uint32_t VulkanDevice::getQueueFamilyIndex(VkQueueFlagBits queueFlags)
{
    /// 获取支持计算的队列组索引
    if (queueFlags & VK_QUEUE_COMPUTE_BIT)
    {
        for (uint32_t i = 0; i < static_cast<uint32_t>(m_queueFamilyProperties.size()); i++)
        {
            if ((m_queueFamilyProperties[i].queueFlags & queueFlags) && ((m_queueFamilyProperties[i].queueFlags & VK_QUEUE_GRAPHICS_BIT) == 0))
            {
                return i;
                break;
            }
        }
    }

    /// 对于其他的队列类型，如果当前没有单独的计算队列，返回支持标志类型的第一个队列
    for (uint32_t i = 0; i < static_cast<uint32_t>(m_queueFamilyProperties.size()); i++)
    {
        if (m_queueFamilyProperties[i].queueFlags & queueFlags)
        {
            return i;
            break;
        }
    }

    throw std::runtime_error("Could not find a matching queue family index");
}
uint32_t VulkanDevice::getGraphicsQueueFamilyIndex()
{
    return m_queueFamilyIndices.graphics;
}
uint32_t VulkanDevice::getComputeQueueFamilyIndex()
{
    return m_queueFamilyIndices.compute;
}

VkResult VulkanDevice::createBuffer(VkBufferUsageFlags usageFlags, VkMemoryPropertyFlags memoryPropertyFlags, VkDeviceSize size, VkBuffer *buffer, VkDeviceMemory *memory, void *data)
{
    // 创建buffer句柄
    VkBufferCreateInfo bufferCreateInfo{};
    bufferCreateInfo.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO;
    bufferCreateInfo.usage = usageFlags;
    bufferCreateInfo.size = size;
    bufferCreateInfo.sharingMode = VK_SHARING_MODE_EXCLUSIVE;
    VK_CHECK_RESULT(vkCreateBuffer(m_logicalDevice, &bufferCreateInfo, nullptr, buffer));

    // 创建buffer的设备内存分配信息
    VkMemoryRequirements memReqs;
    VkMemoryAllocateInfo memAlloc{};
    memAlloc.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
    vkGetBufferMemoryRequirements(m_logicalDevice, *buffer, &memReqs);
    memAlloc.allocationSize = memReqs.size;
    /// 获取符合buffer的设备内存类型索引
    memAlloc.memoryTypeIndex = getMemoryType(memReqs.memoryTypeBits, memoryPropertyFlags);
    VK_CHECK_RESULT(vkAllocateMemory(m_logicalDevice, &memAlloc, nullptr, memory));

    // 如何buffer指针已经存在，复制或者映射该buffer
    if (data != nullptr)
    {
        void *mapped;
        VK_CHECK_RESULT(vkMapMemory(m_logicalDevice, *memory, 0, size, 0, &mapped));
        memcpy(mapped, data, size);
        // 如果host coherency 未设置, 对buffer进行flush，让设备侧可见
        if ((memoryPropertyFlags & VK_MEMORY_PROPERTY_HOST_COHERENT_BIT) == 0)
        {
            VkMappedMemoryRange mappedRange{};
            mappedRange.sType = VK_STRUCTURE_TYPE_MAPPED_MEMORY_RANGE;
            mappedRange.memory = *memory;
            mappedRange.offset = 0;
            mappedRange.size = size;
            vkFlushMappedMemoryRanges(m_logicalDevice, 1, &mappedRange);
        }
        vkUnmapMemory(m_logicalDevice, *memory);
    }

    // 绑定设备内存到buffer object
    VK_CHECK_RESULT(vkBindBufferMemory(m_logicalDevice, *buffer, *memory, 0));

    return VK_SUCCESS;
}

VkCommandPool VulkanDevice::createCommandPool(uint32_t queueFamilyIndex, VkCommandPoolCreateFlags createFlags)
{
    /// 创建命令缓冲池
    VkCommandPoolCreateInfo cmdPoolInfo = {};
    cmdPoolInfo.sType = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO;
    cmdPoolInfo.queueFamilyIndex = queueFamilyIndex;
    cmdPoolInfo.flags = createFlags;
    VkCommandPool cmdPool;
    VK_CHECK_RESULT(vkCreateCommandPool(m_logicalDevice, &cmdPoolInfo, nullptr, &cmdPool));
    return cmdPool;
}

VkCommandBuffer VulkanDevice::createCommandBuffer(VkCommandBufferLevel level, bool begin)
{
    VkCommandBufferAllocateInfo cmdBufAllocateInfo{};
    cmdBufAllocateInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
    cmdBufAllocateInfo.commandPool = m_commandPool;
    cmdBufAllocateInfo.level = level;
    cmdBufAllocateInfo.commandBufferCount = 1;

    VkCommandBuffer cmdBuffer;
    VK_CHECK_RESULT(vkAllocateCommandBuffers(m_logicalDevice, &cmdBufAllocateInfo, &cmdBuffer));

    // 开始记录指令buffer
    if (begin)
    {
        beginCommandBuffer(cmdBuffer);
    }

    return cmdBuffer;
}

void VulkanDevice::beginCommandBuffer(VkCommandBuffer commandBuffer)
{
    VkCommandBufferBeginInfo commandBufferBI{};
    commandBufferBI.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
    VK_CHECK_RESULT(vkBeginCommandBuffer(commandBuffer, &commandBufferBI));
}

void VulkanDevice::flushCommandBuffer(VkCommandBuffer commandBuffer, VkQueue queue, bool free)
{
    VK_CHECK_RESULT(vkEndCommandBuffer(commandBuffer));

    VkSubmitInfo submitInfo{};
    submitInfo.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;
    submitInfo.commandBufferCount = 1;
    submitInfo.pCommandBuffers = &commandBuffer;

    // 创建同步栅栏，确保命令buffer执行完毕
    VkFenceCreateInfo fenceInfo{};
    fenceInfo.sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO;
    VkFence fence;
    VK_CHECK_RESULT(vkCreateFence(m_logicalDevice, &fenceInfo, nullptr, &fence));

    // 提交队列
    VK_CHECK_RESULT(vkQueueSubmit(queue, 1, &submitInfo, fence));
    // 等待栅栏发出信号说明命令buffer执行完毕
    VK_CHECK_RESULT(vkWaitForFences(m_logicalDevice, 1, &fence, VK_TRUE, 100000000000));
    // 同步栅栏使命结束，销毁
    vkDestroyFence(m_logicalDevice, fence, nullptr);
    // 需要的时候，销毁命令buffer
    if (free)
    {
        vkFreeCommandBuffers(m_logicalDevice, m_commandPool, 1, &commandBuffer);
    }
}

VULKANBASE_NAMESPACE_END