#pragma once

#include <vulkan/vulkan.h>
#include <log.hpp>
#include <array>
#include <vector>
#include <string>
#include <algorithm>

namespace vk
{
    struct Device
    {
        Device() {}

        Device(VkPhysicalDevice inPhysicalDevice)
        {
            physicalDevice = inPhysicalDevice;

            vkGetPhysicalDeviceProperties(physicalDevice, &deviceProperties);
            vkGetPhysicalDeviceFeatures(physicalDevice, &deviceFeatures);
            vkGetPhysicalDeviceMemoryProperties(physicalDevice, &deviceMemoryProperties);

            uint32_t extensionCount;
            vkEnumerateDeviceExtensionProperties(physicalDevice, nullptr, &extensionCount, nullptr);

            std::vector<VkExtensionProperties> extensions(extensionCount);
            vkEnumerateDeviceExtensionProperties(physicalDevice, nullptr, &extensionCount, extensions.data());

            for(auto extension : extensions)
                supportedExtensions.push_back(extension.extensionName);
        }

        ~Device()
        {
            if(logicalDevice != VK_NULL_HANDLE)
            {
                vkDestroyCommandPool(logicalDevice, commandPool, nullptr);

                vkDestroyDevice(logicalDevice, nullptr);
            }
        }

        void CreateLogical(VkPhysicalDeviceFeatures enabledDeviceFeatures = VkPhysicalDeviceFeatures())
        {
            uint32_t queueFamilyCount = 0;
            vkGetPhysicalDeviceQueueFamilyProperties(physicalDevice, &queueFamilyCount, nullptr);

            std::vector<VkQueueFamilyProperties> queueFamilies(queueFamilyCount);
            vkGetPhysicalDeviceQueueFamilyProperties(physicalDevice, &queueFamilyCount, queueFamilies.data());
            for(uint32_t i = 0; i < queueFamilies.size(); i++)
            {
                if(queueFamilies[i].queueFlags & VK_QUEUE_GRAPHICS_BIT)
                {
                    queueFamilyIndices.graphics = i;
                }
                else if(queueFamilies[i].queueFlags & VK_QUEUE_TRANSFER_BIT)
                {
                    queueFamilyIndices.transfer = i;
                }
            }

            std::array<VkDeviceQueueCreateInfo, 2> queueCreateInfos = {
                vk::initializers::DeviceQueue(queueFamilyIndices.graphics),
                vk::initializers::DeviceQueue(queueFamilyIndices.transfer)
            };

            std::array<const char*, 1> enabledExtensions = { 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.enabledExtensionCount = static_cast<uint32_t>(enabledExtensions.size());
            deviceCreateInfo.ppEnabledExtensionNames = enabledExtensions.data();
            deviceCreateInfo.pEnabledFeatures = &enabledDeviceFeatures;

            vkCreateDevice(physicalDevice, &deviceCreateInfo, nullptr, &logicalDevice);

            vkGetDeviceQueue(logicalDevice, queueFamilyIndices.graphics, 0, &graphicsQueue);
            vkGetDeviceQueue(logicalDevice, queueFamilyIndices.transfer, 0, &transferQueue);

            VkCommandPoolCreateInfo commandPoolCreateInfo = {};
            commandPoolCreateInfo.sType = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO;
            commandPoolCreateInfo.flags = VK_COMMAND_POOL_CREATE_TRANSIENT_BIT | VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT;
            commandPoolCreateInfo.queueFamilyIndex = queueFamilyIndices.graphics;

            vkCreateCommandPool(logicalDevice, &commandPoolCreateInfo, nullptr, &commandPool);
        }

        bool IsExtensionSupported(const std::string& extension)
        {
            return std::find(supportedExtensions.begin(), supportedExtensions.end(), extension) != supportedExtensions.end();
        }

        VkDevice logicalDevice = VK_NULL_HANDLE;
        VkPhysicalDevice physicalDevice;

        VkPhysicalDeviceProperties deviceProperties;
        VkPhysicalDeviceFeatures deviceFeatures;
        VkPhysicalDeviceMemoryProperties deviceMemoryProperties;

        std::vector<std::string> supportedExtensions;

        struct
        {
            uint32_t graphics, transfer;
        } queueFamilyIndices;

        VkQueue graphicsQueue, transferQueue;

        VkCommandPool commandPool;
    };
}