#include "VulkanContext.hpp"

#include <SDL3/SDL.h>
#include <SDL3/SDL_vulkan.h>

#include <stdexcept>

QueueFamilyIndices findQueueFamilies(
    VkPhysicalDevice device,
    VkSurfaceKHR surface)
{
    QueueFamilyIndices indices;

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

    std::vector<VkQueueFamilyProperties> queueFamilies(
        queueFamilyCount
    );

    vkGetPhysicalDeviceQueueFamilyProperties(
        device,
        &queueFamilyCount,
        queueFamilies.data()
    );

    for (uint32_t i = 0; i < queueFamilyCount; ++i)
    {
        if (queueFamilies[i].queueFlags &
            VK_QUEUE_GRAPHICS_BIT)
        {
            indices.graphicsFamily = i;
        }

        VkBool32 presentSupport = VK_FALSE;
        vkGetPhysicalDeviceSurfaceSupportKHR(
            device,
            i,
            surface,
            &presentSupport
        );

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

        if (indices.isComplete())
        {
            break;
        }
    }

    return indices;
}

SwapchainSupportDetails querySwapchainSupport(
    VkPhysicalDevice device,
    VkSurfaceKHR surface)
{
    SwapchainSupportDetails details;

    vkGetPhysicalDeviceSurfaceCapabilitiesKHR(
        device,
        surface,
        &details.capabilities
    );

    uint32_t formatCount;
    vkGetPhysicalDeviceSurfaceFormatsKHR(
        device,
        surface,
        &formatCount,
        nullptr
    );

    if (formatCount != 0)
    {
        details.formats.resize(formatCount);
        vkGetPhysicalDeviceSurfaceFormatsKHR(
            device,
            surface,
            &formatCount,
            details.formats.data()
        );
    }

    uint32_t presentModeCount;
    vkGetPhysicalDeviceSurfacePresentModesKHR(
        device,
        surface,
        &presentModeCount,
        nullptr
    );

    if (presentModeCount != 0)
    {
        details.presentModes.resize(presentModeCount);
        vkGetPhysicalDeviceSurfacePresentModesKHR(
            device,
            surface,
            &presentModeCount,
            details.presentModes.data()
        );
    }

    return details;
}

VulkanContext::VulkanContext(
    SDL_Window* window,
    const char* appName,
    uint32_t versionMajor,
    uint32_t versionMinor,
    uint32_t versionPatch)
{
    uint32_t extensionCount = 0;
    const char* const* extensions =
        SDL_Vulkan_GetInstanceExtensions(
            &extensionCount
        );

    VkApplicationInfo appInfo{};
    appInfo.sType =
        VK_STRUCTURE_TYPE_APPLICATION_INFO;
    appInfo.pApplicationName = appName;
    appInfo.applicationVersion =
        VK_MAKE_VERSION(
            versionMajor,
            versionMinor,
            versionPatch
        );
    appInfo.apiVersion = VK_API_VERSION_1_3;

    VkInstanceCreateInfo createInfo{};
    createInfo.sType =
        VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO;
    createInfo.pApplicationInfo = &appInfo;
    createInfo.enabledExtensionCount =
        extensionCount;
    createInfo.ppEnabledExtensionNames =
        extensions;

    if (vkCreateInstance(
            &createInfo,
            nullptr,
            &m_instance) != VK_SUCCESS)
    {
        throw std::runtime_error(
            "Failed to create Vulkan instance"
        );
    }

    if (!SDL_Vulkan_CreateSurface(
            window,
            m_instance,
            nullptr,
            &m_surface))
    {
        vkDestroyInstance(m_instance, nullptr);
        throw std::runtime_error(
            "Failed to create Vulkan surface"
        );
    }

    uint32_t gpuCount = 0;
    vkEnumeratePhysicalDevices(
        m_instance,
        &gpuCount,
        nullptr
    );

    if (gpuCount == 0)
    {
        vkDestroySurfaceKHR(
            m_instance,
            m_surface,
            nullptr
        );
        vkDestroyInstance(m_instance, nullptr);
        throw std::runtime_error(
            "No Vulkan GPUs found"
        );
    }

    std::vector<VkPhysicalDevice> gpus(gpuCount);
    vkEnumeratePhysicalDevices(
        m_instance,
        &gpuCount,
        gpus.data()
    );

    for (const auto& gpu : gpus)
    {
        QueueFamilyIndices testIndices =
            findQueueFamilies(gpu, m_surface);

        if (testIndices.isComplete())
        {
            m_physicalDevice = gpu;
            m_queueIndices = testIndices;
            break;
        }
    }

    if (m_physicalDevice == VK_NULL_HANDLE)
    {
        vkDestroySurfaceKHR(
            m_instance,
            m_surface,
            nullptr
        );
        vkDestroyInstance(m_instance, nullptr);
        throw std::runtime_error(
            "No suitable GPU found"
        );
    }

    VkPhysicalDeviceProperties props{};
    vkGetPhysicalDeviceProperties(
        m_physicalDevice,
        &props
    );

    std::cout << "Selected GPU: "
              << props.deviceName
              << std::endl;

    float queuePriority = 1.0f;
    std::vector<VkDeviceQueueCreateInfo>
        queueCreateInfos;
    std::vector<uint32_t> uniqueQueues;

    uniqueQueues.push_back(
        *m_queueIndices.graphicsFamily
    );

    if (*m_queueIndices.presentFamily !=
        *m_queueIndices.graphicsFamily)
    {
        uniqueQueues.push_back(
            *m_queueIndices.presentFamily
        );
    }

    for (uint32_t queueFamily : uniqueQueues)
    {
        VkDeviceQueueCreateInfo queueInfo{};
        queueInfo.sType =
            VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO;
        queueInfo.queueFamilyIndex = queueFamily;
        queueInfo.queueCount = 1;
        queueInfo.pQueuePriorities = &queuePriority;
        queueCreateInfos.push_back(queueInfo);
    }

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

    VkPhysicalDeviceFeatures deviceFeatures{};

    VkDeviceCreateInfo deviceInfo{};
    deviceInfo.sType =
        VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO;
    deviceInfo.queueCreateInfoCount =
        static_cast<uint32_t>(
            queueCreateInfos.size()
        );
    deviceInfo.pQueueCreateInfos =
        queueCreateInfos.data();
    deviceInfo.enabledExtensionCount =
        static_cast<uint32_t>(
            deviceExtensions.size()
        );
    deviceInfo.ppEnabledExtensionNames =
        deviceExtensions.data();
    deviceInfo.pEnabledFeatures = &deviceFeatures;

    if (vkCreateDevice(
            m_physicalDevice,
            &deviceInfo,
            nullptr,
            &m_device) != VK_SUCCESS)
    {
        vkDestroySurfaceKHR(
            m_instance,
            m_surface,
            nullptr
        );
        vkDestroyInstance(m_instance, nullptr);
        throw std::runtime_error(
            "Failed to create logical device"
        );
    }

    std::cout << "Graphics queue: "
              << *m_queueIndices.graphicsFamily
              << std::endl;
    std::cout << "Present queue: "
              << *m_queueIndices.presentFamily
              << std::endl;

    vkGetDeviceQueue(
        m_device,
        *m_queueIndices.graphicsFamily,
        0,
        &m_graphicsQueue
    );

    vkGetDeviceQueue(
        m_device,
        *m_queueIndices.presentFamily,
        0,
        &m_presentQueue
    );

    VkCommandPoolCreateInfo poolInfo{};
    poolInfo.sType =
        VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO;
    poolInfo.flags =
        VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT;
    poolInfo.queueFamilyIndex =
        *m_queueIndices.graphicsFamily;

    if (vkCreateCommandPool(
            m_device,
            &poolInfo,
            nullptr,
            &m_commandPool) != VK_SUCCESS)
    {
        vkDestroyDevice(m_device, nullptr);
        vkDestroySurfaceKHR(
            m_instance,
            m_surface,
            nullptr
        );
        vkDestroyInstance(m_instance, nullptr);
        throw std::runtime_error(
            "Failed to create command pool"
        );
    }

    std::cout << "Vulkan context created"
              << std::endl;
}

VulkanContext::~VulkanContext()
{
    if (m_device != VK_NULL_HANDLE)
    {
        vkDestroyCommandPool(
            m_device,
            m_commandPool,
            nullptr
        );
        vkDestroyDevice(m_device, nullptr);
        m_device = VK_NULL_HANDLE;
    }

    if (m_surface != VK_NULL_HANDLE)
    {
        vkDestroySurfaceKHR(
            m_instance,
            m_surface,
            nullptr
        );
        m_surface = VK_NULL_HANDLE;
    }

    if (m_instance != VK_NULL_HANDLE)
    {
        vkDestroyInstance(m_instance, nullptr);
        m_instance = VK_NULL_HANDLE;
    }
}

void VulkanContext::release()
{
    if (m_device != VK_NULL_HANDLE)
    {
        vkDeviceWaitIdle(m_device);
        vkDestroyCommandPool(
            m_device,
            m_commandPool,
            nullptr
        );
        vkDestroyDevice(m_device, nullptr);
        m_device = VK_NULL_HANDLE;
    }

    if (m_surface != VK_NULL_HANDLE)
    {
        vkDestroySurfaceKHR(
            m_instance,
            m_surface,
            nullptr
        );
        m_surface = VK_NULL_HANDLE;
    }

    if (m_instance != VK_NULL_HANDLE)
    {
        vkDestroyInstance(m_instance, nullptr);
        m_instance = VK_NULL_HANDLE;
    }
}