#include <Vk/Device/Device.hpp>
#include <Utilities/Utilities.hpp>
#include <Settings/Settings.hpp>
#include <Log/Log.hpp>
#include <vulkan/vulkan_core.h>

const std::vector<const char*> Vk::Device::deviceExtensions = {
    VK_KHR_SWAPCHAIN_EXTENSION_NAME,
};

VkPhysicalDevice Vk::Device::pickPhysicalDevice()
{
    Log::info("Ranking physical devices...");

    std::vector<VkPhysicalDevice> devices = m_instance->getPhysicalDevices();

    std::optional<int> highestScore;
    VkPhysicalDevice bestDevice{};

    for (const VkPhysicalDevice &device : devices)
    {
        VkPhysicalDeviceProperties deviceProperties;
        vkGetPhysicalDeviceProperties(device, &deviceProperties);

        std::optional<int> score = rankPhysicalDevice(device);
        if (!score.has_value())
        {
            Log::info(formatString("$+ -> unsuitable", deviceProperties.deviceName));
            continue;
        }
        else
        {
            Log::info(formatString("$+ -> $+", deviceProperties.deviceName, *score));
        }

        if (!highestScore.has_value() || *highestScore < *score)
        {
            highestScore = score;
            bestDevice = device;
        }
    }

    if (!highestScore.has_value())
    {
        crash("Failed to find any suitable GPU!");
    }

    VkPhysicalDeviceProperties deviceProperties;
    vkGetPhysicalDeviceProperties(bestDevice, &deviceProperties);

    Log::info(formatString("Using device '$+'", deviceProperties.deviceName));
    return bestDevice;
}

Vk::Device::QueueFamilies Vk::Device::findQueueFamilies(VkPhysicalDevice device)
{
    QueueFamilies families;

    uint32_t queueFamilyCount = 0;
    vkGetPhysicalDeviceQueueFamilyProperties(device, &queueFamilyCount, nullptr);

    std::vector<VkQueueFamilyProperties> queueFamilies(queueFamilyCount);
    vkGetPhysicalDeviceQueueFamilyProperties(device, &queueFamilyCount, queueFamilies.data());

    uint32_t i = 0;
    for (const VkQueueFamilyProperties &queueFamily : queueFamilies)
    {
        VkBool32 presentSupport = false;
        vkGetPhysicalDeviceSurfaceSupportKHR(device, i, m_surface->inner(), &presentSupport);

        if (presentSupport && (queueFamily.queueFlags & VK_QUEUE_GRAPHICS_BIT))
        {
            families.graphicsFamily = i;
            families.presentFamily = i;
            return families;
        }

        i++;
    }

    i = 0;
    for (const VkQueueFamilyProperties &queueFamily : queueFamilies)
    {
        if (queueFamily.queueFlags & VK_QUEUE_GRAPHICS_BIT)
        {
            families.graphicsFamily = i;
        }

        VkBool32 presentSupport = false;
        vkGetPhysicalDeviceSurfaceSupportKHR(device, i, m_surface->inner(), &presentSupport);

        if (presentSupport)
        {
            families.presentFamily = i;
        }

        i++;
    }

    return families;
}

bool Vk::Device::areDeviceExtensionsSupported(VkPhysicalDevice device)
{
    uint32_t extensionCount;
    vkEnumerateDeviceExtensionProperties(device, nullptr, &extensionCount, nullptr);

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

    std::set<std::string> requiredExtensions(deviceExtensions.begin(), deviceExtensions.end());
    for (const VkExtensionProperties &extension : availableExtensions)
    {
        requiredExtensions.erase(extension.extensionName);
    }

    return requiredExtensions.empty();
}

Vk::Device::SwapchainSupportDetails Vk::Device::querySwapchainSupport(VkPhysicalDevice device)
{
    SwapchainSupportDetails details;

    vkGetPhysicalDeviceSurfaceCapabilitiesKHR(device, m_surface->inner(), &details.capabilities);

    uint32_t formatCount;
    vkGetPhysicalDeviceSurfaceFormatsKHR(device, m_surface->inner(), &formatCount, nullptr);

    if (formatCount != 0) {
        details.formats.resize(formatCount);
        vkGetPhysicalDeviceSurfaceFormatsKHR(device, m_surface->inner(), &formatCount, details.formats.data());
    }

    uint32_t presentModeCount;
    vkGetPhysicalDeviceSurfacePresentModesKHR(device, m_surface->inner(), &presentModeCount, nullptr);

    if (presentModeCount != 0) {
        details.presentModes.resize(presentModeCount);
        vkGetPhysicalDeviceSurfacePresentModesKHR(device, m_surface->inner(), &presentModeCount, details.presentModes.data());
    }

    return details;
}

std::optional<int> Vk::Device::rankPhysicalDevice(VkPhysicalDevice device)
{
    VkPhysicalDeviceProperties deviceProperties;
    vkGetPhysicalDeviceProperties(device, &deviceProperties);

    VkPhysicalDeviceFeatures deviceFeatures;
    vkGetPhysicalDeviceFeatures(device, &deviceFeatures);

    int score;
    switch (deviceProperties.deviceType)
    {
    case VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU:
        score = 1000;
        break;
    case VK_PHYSICAL_DEVICE_TYPE_INTEGRATED_GPU:
        score = 100;
        break;
    case VK_PHYSICAL_DEVICE_TYPE_VIRTUAL_GPU:
        score = 10;
        break;
    case VK_PHYSICAL_DEVICE_TYPE_CPU:
        score = 1;
        break;
    default:
        return std::nullopt;
    }

    if (deviceProperties.limits.maxComputeWorkGroupSize[0] < 64)
    {
        return std::nullopt;
    }

    QueueFamilies families = findQueueFamilies(device);
    if (!families.isComplete())
    {
        return std::nullopt;
    }

    if (families.graphicsFamily == families.presentFamily)
    {
        score++;
    }

    if (!areDeviceExtensionsSupported(device))
    {
        return std::nullopt;
    }

    SwapchainSupportDetails swapChainSupport = querySwapchainSupport(device);
    if (swapChainSupport.formats.empty() || swapChainSupport.presentModes.empty())
    {
        return std::nullopt;
    }

    return score;
}

VkQueue Vk::Device::Queue::inner()
{
    return m_queue;
}

Vk::Device::Device(std::shared_ptr<Vk::Instance> instance, std::shared_ptr<Surface> surface) :
    m_instance(instance),
    m_surface(surface)
{
    m_physicalDevice = pickPhysicalDevice();

    vkGetPhysicalDeviceFeatures(m_physicalDevice, &features);
    vkGetPhysicalDeviceProperties(m_physicalDevice, &properties);

    {
        maxWorkGroupCount = glm::uvec3(
            properties.limits.maxComputeWorkGroupCount[0],
            properties.limits.maxComputeWorkGroupCount[1],
            properties.limits.maxComputeWorkGroupCount[2]
        );

        maxWorkGroupSize = glm::uvec3(
            properties.limits.maxComputeWorkGroupSize[0],
            properties.limits.maxComputeWorkGroupSize[1],
            properties.limits.maxComputeWorkGroupSize[2]
        );

        maxWorkGroupInvocations = properties.limits.maxComputeWorkGroupInvocations;

        std::stringstream ss;
        ss << "Properties of selected device:" << std::endl;
        ss << "  - Timestamp query support? " << (properties.limits.timestampComputeAndGraphics ? "YES" : "NO") << std::endl;
        if (properties.limits.timestampComputeAndGraphics)
        {
            ss << "    - Timestamp period: " << properties.limits.timestampPeriod << "ns" << std::endl;
        }

        ss << "  - Max workgroup count: " << maxWorkGroupCount << std::endl;
        ss << "  - Max workgroup size: " << maxWorkGroupSize << std::endl;
        ss << "  - Max workgroup invocations: " << maxWorkGroupInvocations << std::endl;

        Log::info(ss.str());
    }

    SwapchainSupportDetails swapchainSupportDetails = querySwapchainSupport(m_physicalDevice);
    Log::info(formatString("Swapchain support details on selected device: $+", swapchainSupportDetails));

    QueueFamilies families = findQueueFamilies(m_physicalDevice);
    m_queueFamilies = families;
    float queuePriority = 1.0f;

    std::set<uint32_t> uniqueQueueFamilies = {families.graphicsFamily.value(), families.presentFamily.value()};
    std::vector<VkDeviceQueueCreateInfo> queueCreateInfos;

    for (uint32_t uniqueFamily : uniqueQueueFamilies)
    {
        VkDeviceQueueCreateInfo queueCreateInfo{};
        queueCreateInfo.sType = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO;
        queueCreateInfo.queueFamilyIndex = uniqueFamily;
        queueCreateInfo.queueCount = 1;
        queueCreateInfo.pQueuePriorities = &queuePriority;
        queueCreateInfos.push_back(queueCreateInfo);
    }

    VkPhysicalDeviceFeatures requestedFeatures{};
    requestedFeatures.independentBlend = VK_TRUE;

    VkDeviceCreateInfo createInfo{};
    createInfo.sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO;
    createInfo.pQueueCreateInfos = queueCreateInfos.data();
    createInfo.queueCreateInfoCount = queueCreateInfos.size();
    createInfo.pEnabledFeatures = &requestedFeatures;
    createInfo.enabledExtensionCount = deviceExtensions.size();
    createInfo.ppEnabledExtensionNames = deviceExtensions.data();

    if (settings.validationLayers)
    {
        createInfo.ppEnabledLayerNames = Vk::Instance::validationLayers.data();
        createInfo.enabledLayerCount = Vk::Instance::validationLayers.size();
    }

    if (vkCreateDevice(m_physicalDevice, &createInfo, nullptr, &m_device) != VK_SUCCESS)
    {
        crash("Failed to create logical device.");
    }

    vkGetPhysicalDeviceMemoryProperties(m_physicalDevice, &m_memProperties);

    vkGetDeviceQueue(m_device, families.graphicsFamily.value(), 0, &graphicsQueue.m_queue);
    vkGetDeviceQueue(m_device, families.presentFamily.value(), 0, &presentQueue.m_queue);

    VmaAllocatorCreateInfo allocatorCreateInfo = {};
    allocatorCreateInfo.physicalDevice = m_physicalDevice;
    allocatorCreateInfo.device = m_device;
    allocatorCreateInfo.instance = instance->inner();

    if (vmaCreateAllocator(&allocatorCreateInfo, &m_allocator) != VK_SUCCESS)
    {
        crash("Failed to create an allocator.");
    }
}

Vk::Device::~Device()
{
    vmaDestroyAllocator(m_allocator);
    vkDestroyDevice(m_device, nullptr);
}

VkDevice Vk::Device::inner()
{
    return m_device;
}

Vk::Device::SwapchainSupportDetails Vk::Device::getSwapchainSupportDetails()
{
    return querySwapchainSupport(m_physicalDevice);
}

VkFormat Vk::Device::findSupportedFormat(
    const std::vector<VkFormat>& candidates,
    VkImageTiling tiling,
    VkFormatFeatureFlags features
)
{
    for (VkFormat format : candidates)
    {
        VkFormatProperties props;
        vkGetPhysicalDeviceFormatProperties(m_physicalDevice, format, &props);

        if (tiling == VK_IMAGE_TILING_LINEAR && (props.linearTilingFeatures & features) == features)
        {
            return format;
        }
        else if (tiling == VK_IMAGE_TILING_OPTIMAL && (props.optimalTilingFeatures & features) == features)
        {
            return format;
        }
    }

    crash("Failed to find supported image format!");
}

VkFormat Vk::Device::findDepthFormat()
{
    return findSupportedFormat(
        {VK_FORMAT_D32_SFLOAT, VK_FORMAT_D32_SFLOAT_S8_UINT, VK_FORMAT_D24_UNORM_S8_UINT},
        VK_IMAGE_TILING_OPTIMAL,
        VK_FORMAT_FEATURE_DEPTH_STENCIL_ATTACHMENT_BIT
    );
}

Vk::Device::QueueFamilies Vk::Device::getQueueFamilies()
{
    return m_queueFamilies;
}

void Vk::Device::waitIdle()
{
    vkDeviceWaitIdle(m_device);
}

VmaAllocator Vk::Device::getAllocator()
{
    return m_allocator;
}

uint32_t Vk::Device::getNumMemoryHeaps()
{
    return m_memProperties.memoryHeapCount;
}

bool Vk::Device::areTimerQueriesEnabled()
{
    return properties.limits.timestampComputeAndGraphics;
}

std::ostream& operator<<(std::ostream &os, const Vk::Device::SwapchainSupportDetails &details)
{
    os << "SwapchainSupportDetails:" << std::endl;

    os << "  Surface formats:" << std::endl;
    for (const VkSurfaceFormatKHR &format : details.formats)
    {
        os << "    - format=" << string_VkFormat(format.format) <<
            " colorSpace=" << string_VkColorSpaceKHR(format.colorSpace) << std::endl;
    }

    os << "  Present modes:" << std::endl;
    for (const VkPresentModeKHR &presentMode : details.presentModes)
    {
        os << "    - " << string_VkPresentModeKHR(presentMode) << std::endl;
    }

    return os;
}