#include "vulkan-context.hpp"

#include <array>
#include <print>
#include <string>
#include <utility>

#include <imgui_impl_vulkan.h>

namespace {
inline void check_vk_result(VkResult err) {
    if (err == 0) {
        return;
    }

    std::println(stderr, "[vulkan] Error: VkResult = {}", static_cast<int>(err));
    if (err < 0) {
        std::abort();
    }
}

bool is_extension_available(const std::vector<VkExtensionProperties>& properties, const std::string& extension) {
    for (const auto& property : properties) {
        if (extension == static_cast<const char*>(property.extensionName)) {
            return true;
        }
    }

    return false;
}

std::vector<VkExtensionProperties> get_instance_extension_properties() {
    uint32_t num_properties = 0;
    std::vector<VkExtensionProperties> properties;
    vkEnumerateInstanceExtensionProperties(nullptr, &num_properties, nullptr);
    properties.resize(num_properties);
    check_vk_result(vkEnumerateInstanceExtensionProperties(nullptr, &num_properties, properties.data()));
    return properties;
}

std::vector<VkExtensionProperties> get_device_extension_properties(VkPhysicalDevice device) {
    uint32_t num_properties = 0;
    std::vector<VkExtensionProperties> properties;
    vkEnumerateDeviceExtensionProperties(device, nullptr, &num_properties, nullptr);
    properties.resize(num_properties);
    check_vk_result(vkEnumerateDeviceExtensionProperties(device, nullptr, &num_properties, properties.data()));
    return properties;
}

VkInstance create_instance(VkAllocationCallbacks* allocator, std::vector<const char*> extensions) {
    VkInstanceCreateInfo instance_create_info = {};
    instance_create_info.sType = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO;

    auto instance_properties = get_instance_extension_properties();

    if (is_extension_available(instance_properties, VK_KHR_GET_PHYSICAL_DEVICE_PROPERTIES_2_EXTENSION_NAME)) {
        extensions.push_back(VK_KHR_GET_PHYSICAL_DEVICE_PROPERTIES_2_EXTENSION_NAME);
    }

    if (is_extension_available(instance_properties, VK_KHR_PORTABILITY_ENUMERATION_EXTENSION_NAME)) {
        extensions.push_back(VK_KHR_PORTABILITY_ENUMERATION_EXTENSION_NAME);
        instance_create_info.flags |= VK_INSTANCE_CREATE_ENUMERATE_PORTABILITY_BIT_KHR;
    }

    instance_create_info.enabledExtensionCount = (uint32_t) extensions.size();
    instance_create_info.ppEnabledExtensionNames = extensions.data();

    VkInstance instance = nullptr;
    check_vk_result(vkCreateInstance(&instance_create_info, allocator, &instance));
    return instance;
}

VkDevice create_logical_device(VkAllocationCallbacks* allocator, VkPhysicalDevice physical_device, uint32_t queue_family,
    std::vector<const char*> extensions) {
    const std::array<float, 1> queue_priority = {1.0F};

    std::array<VkDeviceQueueCreateInfo, 1> queue_info = {};
    queue_info[0].sType = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO;
    queue_info[0].queueFamilyIndex = queue_family;
    queue_info[0].queueCount = 1;
    queue_info[0].pQueuePriorities = queue_priority.data();

    VkDeviceCreateInfo create_info = {};
    create_info.sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO;
    create_info.queueCreateInfoCount = queue_info.size();
    create_info.pQueueCreateInfos = queue_info.data();
    create_info.enabledExtensionCount = extensions.size();
    create_info.ppEnabledExtensionNames = extensions.data();

    VkDevice logical_device = {};
    check_vk_result(vkCreateDevice(physical_device, &create_info, allocator, &logical_device));
    return logical_device;
}
} // namespace

VulkanContext::VulkanContext(std::vector<const char*> instance_extensions) noexcept {
#ifdef IMGUI_IMPL_VULKAN_USE_VOLK
    check_vk_result(volkInitialize());
#endif

    instance = create_instance(allocator, std::move(instance_extensions));

#ifdef IMGUI_IMPL_VULKAN_USE_VOLK
    volkLoadInstance(instance);
#endif

    physical_device = ImGui_ImplVulkanH_SelectPhysicalDevice(instance);
    assert(physical_device != VK_NULL_HANDLE);

    queue_family = ImGui_ImplVulkanH_SelectQueueFamilyIndex(physical_device);
    assert(queue_family != UINT32_MAX);

    std::vector<const char*> device_extensions = {"VK_KHR_swapchain"};
    auto device_properties = get_device_extension_properties(physical_device);

#ifdef VK_KHR_PORTABILITY_SUBSET_EXTENSION_NAME
    if (IsExtensionAvailable(device_properties, VK_KHR_PORTABILITY_SUBSET_EXTENSION_NAME)) {
        device_extensions.push_back(VK_KHR_PORTABILITY_SUBSET_EXTENSION_NAME);
    }
#endif

    logical_device = create_logical_device(allocator, physical_device, queue_family, device_extensions);

    vkGetDeviceQueue(logical_device, queue_family, 0, &queue);

    std::array<VkDescriptorPoolSize, 1> pool_sizes = {
        {{.type = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
            .descriptorCount = IMGUI_IMPL_VULKAN_MINIMUM_IMAGE_SAMPLER_POOL_SIZE}},
    };

    VkDescriptorPoolCreateInfo pool_info = {};
    pool_info.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;
    pool_info.flags = VK_DESCRIPTOR_POOL_CREATE_FREE_DESCRIPTOR_SET_BIT;
    pool_info.maxSets = 0;
    for (auto& size : pool_sizes) {
        pool_info.maxSets += size.descriptorCount;
    }
    pool_info.poolSizeCount = pool_sizes.size();
    pool_info.pPoolSizes = pool_sizes.data();
    check_vk_result(vkCreateDescriptorPool(logical_device, &pool_info, allocator, &descriptor_pool));
}
