#include "DeviceFunctions.hpp"

#include <stdexcept>
#include <string>
#include <optional>

VkExtent2D DeviceFunctions::choose_extent(const VkSurfaceCapabilitiesKHR& capabilities, GLFWwindow* window)
{
    // The acceptable rendering sparsity is calculated here
    if (capabilities.currentExtent.width != std::numeric_limits<uint32_t>::max())
        return capabilities.currentExtent;
    else
    {
        int width, height;
        glfwGetFramebufferSize(window, &width, &height);
        VkExtent2D actualExtent = {
            static_cast<uint32_t>(width),
            static_cast<uint32_t>(height)
        };
        actualExtent.width = std::max(capabilities.minImageExtent.width, std::min(capabilities.maxImageExtent.width, actualExtent.width));
        return actualExtent;
    }
}

uint32_t DeviceFunctions::find_memory_type(VkPhysicalDevice physical_device, uint32_t typeFilter,
                                           VkMemoryPropertyFlags properties)
{
    VkPhysicalDeviceMemoryProperties memory_properties{};
    vkGetPhysicalDeviceMemoryProperties(physical_device, &memory_properties);

    std::optional<uint32_t> mem_type = 0;
    for (uint32_t i = 0; i < memory_properties.memoryTypeCount; ++i)
    {
        if (typeFilter & (1 << i) && (memory_properties.memoryTypes[i].propertyFlags & properties) == properties)
            mem_type = i;
    }
    
    if (mem_type.has_value() == false)
        throw std::runtime_error("failed to find suitable memory type!");
    return mem_type.value();
}

void DeviceFunctions::device_memcpy(VkDevice device, uint32_t transfer_queue_family, VkBuffer src, VkBuffer dst, size_t size)
{
    VkCommandPool copy_pool;
    VkCommandPoolCreateInfo copy_pool_info{};
    copy_pool_info.sType = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO;
    copy_pool_info.queueFamilyIndex = transfer_queue_family;
    copy_pool_info.flags = VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT;
    VK_CHECK(vkCreateCommandPool(device, &copy_pool_info, nullptr, &copy_pool))
    
    VkCommandBufferAllocateInfo command_buffer_allocate_info{};
    command_buffer_allocate_info.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
    command_buffer_allocate_info.commandPool = copy_pool;
    command_buffer_allocate_info.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
    command_buffer_allocate_info.commandBufferCount = 1;
    
    VkCommandBuffer copy_buffer;
    VK_CHECK(vkAllocateCommandBuffers(device, &command_buffer_allocate_info, &copy_buffer))
    
    VkCommandBufferBeginInfo command_buffer_begin_info{};
    command_buffer_begin_info.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
    command_buffer_begin_info.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT;
    
    VK_CHECK(vkBeginCommandBuffer(copy_buffer, &command_buffer_begin_info))
    
    VkBufferCopy buffer_copy;
    buffer_copy.srcOffset = 0;
    buffer_copy.dstOffset = 0;
    buffer_copy.size = size;
    vkCmdCopyBuffer(copy_buffer, src, dst, 1, &buffer_copy);

    vkEndCommandBuffer(copy_buffer);

    VkSubmitInfo submit_info{};
    submit_info.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;
    submit_info.commandBufferCount = 1;
    submit_info.pCommandBuffers = &copy_buffer;

    VkQueue transfer_queue;
    vkGetDeviceQueue(device, transfer_queue_family, 0, &transfer_queue);

    vkQueueSubmit(transfer_queue, 1, &submit_info, VK_NULL_HANDLE);
    vkQueueWaitIdle(transfer_queue);
    
    vkFreeCommandBuffers(device, copy_pool, 1, &copy_buffer);
    vkDestroyCommandPool(device, copy_pool, nullptr);
}

VkBuffer DeviceFunctions::create_buffer(VkDevice device, VkDeviceSize size, VkBufferUsageFlags usage,
    const std::vector<uint32_t>& queue_families)
{
    VkBufferCreateInfo buffer_info{};
    buffer_info.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO;
    buffer_info.size = size;
    buffer_info.usage = usage;
    
    if (queue_families.size() > 1)
    {
        buffer_info.sharingMode = VK_SHARING_MODE_CONCURRENT;
        buffer_info.queueFamilyIndexCount = static_cast<uint32_t>(queue_families.size());
        buffer_info.pQueueFamilyIndices = queue_families.data();
    }
    else
    {
        buffer_info.sharingMode = VK_SHARING_MODE_EXCLUSIVE;
        buffer_info.queueFamilyIndexCount = 0;
        buffer_info.pQueueFamilyIndices = nullptr;
    }

    VkBuffer buffer;
    VK_CHECK(vkCreateBuffer(device, &buffer_info, nullptr, &buffer))
    return buffer;
}

VkDeviceMemory DeviceFunctions::allocate_device_memory(VkPhysicalDevice physical_device, VkDevice device, VkBuffer buffer, VkMemoryPropertyFlags property)
{
    VkMemoryRequirements requirements{};
    vkGetBufferMemoryRequirements(device, buffer, &requirements);
    
    VkMemoryAllocateInfo allocate_info{};
    allocate_info.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
    allocate_info.allocationSize = requirements.size;
    allocate_info.memoryTypeIndex = find_memory_type(physical_device, requirements.memoryTypeBits, property);

    VkDeviceMemory device_memory;
    VK_CHECK(vkAllocateMemory(device, &allocate_info, nullptr, &device_memory))
    return device_memory;
}