module;

#include <cstdint>
#include <vulkan/vulkan.h>
#include <vulkan/vulkan_wayland.h>
#include <iostream>
#include <exception>
#include <vulkan/vk_enum_string_helper.h>
#include <cstring>
#include <print>
#include <cstdio>

#define GET_EXTENSION_FUNCTION(_id) ((PFN_##_id)(vkGetInstanceProcAddr(instance, #_id)))

module Crafter.Graphics;
using namespace Crafter;
const char* const instanceExtensionNames[] = {
    "VK_EXT_debug_utils",
    "VK_KHR_surface",
    "VK_KHR_wayland_surface"
};
const char* const deviceExtensionNames[] = {
    "VK_KHR_swapchain",
    "VK_KHR_spirv_1_4",
    "VK_EXT_mesh_shader",
    "VK_KHR_shader_float_controls"
};
const char* const layerNames[] = {
    "VK_LAYER_KHRONOS_validation"
};


void VulkanDevice::CHECK_VK_RESULT(VkResult result) {
    if (result != VK_SUCCESS)
    {
        throw std::runtime_error(string_VkResult(result));
    }    
} 

VkBool32 onError(VkDebugUtilsMessageSeverityFlagBitsEXT severity, VkDebugUtilsMessageTypeFlagsEXT type, const VkDebugUtilsMessengerCallbackDataEXT* callbackData, void* userData)
{
    printf("Vulkan ");

    switch (type)
    {
    case VK_DEBUG_UTILS_MESSAGE_TYPE_GENERAL_BIT_EXT :
        printf("general ");
        break;
    case VK_DEBUG_UTILS_MESSAGE_TYPE_VALIDATION_BIT_EXT :
        printf("validation ");
        break;
    case VK_DEBUG_UTILS_MESSAGE_TYPE_PERFORMANCE_BIT_EXT :
        printf("performance ");
        break;
    }

    switch (severity)
    {
    case VK_DEBUG_UTILS_MESSAGE_SEVERITY_VERBOSE_BIT_EXT :
        printf("(verbose): ");
        break;
    default :
    case VK_DEBUG_UTILS_MESSAGE_SEVERITY_INFO_BIT_EXT :
        printf("(info): ");
        break;
    case VK_DEBUG_UTILS_MESSAGE_SEVERITY_WARNING_BIT_EXT :
        printf("(warning): ");
        break;
    case VK_DEBUG_UTILS_MESSAGE_SEVERITY_ERROR_BIT_EXT :
        printf("(error): ");
        break;
    }

    printf("%s\n", callbackData->pMessage);

    return 0;
}

void VulkanDevice::CreateDevice() {
    VkApplicationInfo app{VK_STRUCTURE_TYPE_APPLICATION_INFO};
	app.pApplicationName = "";
	app.pEngineName      = "Crafter.Graphics";
	app.apiVersion       = VK_MAKE_VERSION(1, 4, 0);

    VkInstanceCreateInfo instanceCreateInfo = {};
    instanceCreateInfo.sType = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO;
    instanceCreateInfo.pApplicationInfo = &app;
    instanceCreateInfo.enabledExtensionCount = sizeof(instanceExtensionNames) / sizeof(const char*);
    instanceCreateInfo.ppEnabledExtensionNames = instanceExtensionNames;

    size_t foundInstanceLayers = 0;

    std::uint32_t instanceLayerCount;
    CHECK_VK_RESULT(vkEnumerateInstanceLayerProperties(&instanceLayerCount, NULL));

    std::vector<VkLayerProperties> instanceLayerProperties(instanceLayerCount);
    CHECK_VK_RESULT(vkEnumerateInstanceLayerProperties(&instanceLayerCount, instanceLayerProperties.data()));

    for (uint32_t i = 0; i < instanceLayerCount; i++)
    {
        for (size_t j = 0; j < sizeof(layerNames) / sizeof(const char*); j++)
        {
            if (std::strcmp(instanceLayerProperties[i].layerName, layerNames[j]) == 0)
            {
                foundInstanceLayers++;
            }
        }
    }

    if (foundInstanceLayers >= sizeof(layerNames) / sizeof(const char*))
    {
        instanceCreateInfo.enabledLayerCount = sizeof(layerNames) / sizeof(const char*);
        instanceCreateInfo.ppEnabledLayerNames = layerNames;
    }

    CHECK_VK_RESULT(vkCreateInstance(&instanceCreateInfo, NULL, &instance));

    VkDebugUtilsMessengerCreateInfoEXT debugUtilsMessengerCreateInfo = {};
    debugUtilsMessengerCreateInfo.sType = VK_STRUCTURE_TYPE_DEBUG_UTILS_MESSENGER_CREATE_INFO_EXT;
    debugUtilsMessengerCreateInfo.messageSeverity = VK_DEBUG_UTILS_MESSAGE_SEVERITY_INFO_BIT_EXT | VK_DEBUG_UTILS_MESSAGE_SEVERITY_WARNING_BIT_EXT | VK_DEBUG_UTILS_MESSAGE_SEVERITY_ERROR_BIT_EXT;
    debugUtilsMessengerCreateInfo.messageType = VK_DEBUG_UTILS_MESSAGE_TYPE_GENERAL_BIT_EXT | VK_DEBUG_UTILS_MESSAGE_TYPE_VALIDATION_BIT_EXT | VK_DEBUG_UTILS_MESSAGE_TYPE_PERFORMANCE_BIT_EXT;
    debugUtilsMessengerCreateInfo.pfnUserCallback = onError;
    CHECK_VK_RESULT(GET_EXTENSION_FUNCTION(vkCreateDebugUtilsMessengerEXT)(instance, &debugUtilsMessengerCreateInfo, NULL, &debugMessenger));

    uint32_t physDeviceCount;
    vkEnumeratePhysicalDevices(instance, &physDeviceCount, NULL);

    std::vector<VkPhysicalDevice> physDevices(physDeviceCount);
    vkEnumeratePhysicalDevices(instance, &physDeviceCount, physDevices.data());

    uint32_t bestScore = 0;

    for (uint32_t i = 0; i < physDeviceCount; i++)
    {
        VkPhysicalDevice device = physDevices[i];

        VkPhysicalDeviceProperties properties;
        vkGetPhysicalDeviceProperties(device, &properties);

        uint32_t score;

        switch (properties.deviceType)
        {
        default :
            continue;
        case VK_PHYSICAL_DEVICE_TYPE_OTHER :
            score = 1;
            break;
        case VK_PHYSICAL_DEVICE_TYPE_INTEGRATED_GPU :
            score = 4;
            break;
        case VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU :
            score = 5;
            break;
        case VK_PHYSICAL_DEVICE_TYPE_VIRTUAL_GPU :
            score = 3;
            break;
        case VK_PHYSICAL_DEVICE_TYPE_CPU :
            score = 2;
            break;
        }

        if (score > bestScore)
        {
            physDevice = device;
            bestScore = score;
        }
    }

    
    uint32_t queueFamilyCount;
    vkGetPhysicalDeviceQueueFamilyProperties(physDevice, &queueFamilyCount, NULL);

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

    for (uint32_t i = 0; i < queueFamilyCount; i++)
    {
        if (queueFamilies[i].queueFlags & VK_QUEUE_GRAPHICS_BIT)
        {
            queueFamilyIndex = i;
            break;
        }
    }
        
    float priority = 1;

    VkDeviceQueueCreateInfo queueCreateInfo = {};
    queueCreateInfo.sType = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO;
    queueCreateInfo.queueFamilyIndex = queueFamilyIndex;
    queueCreateInfo.queueCount = 1;
    queueCreateInfo.pQueuePriorities = &priority;

    VkPhysicalDeviceMeshShaderFeaturesEXT ext_feature = {VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_MESH_SHADER_FEATURES_EXT};
    ext_feature.meshShader = VK_TRUE;
    VkPhysicalDeviceFeatures2 physical_features2 = {VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2};
    physical_features2.pNext = &ext_feature;

    VkDeviceCreateInfo deviceCreateInfo = {};
    deviceCreateInfo.sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO;
    deviceCreateInfo.queueCreateInfoCount = 1;
    deviceCreateInfo.pQueueCreateInfos = &queueCreateInfo;
    deviceCreateInfo.enabledExtensionCount = sizeof(deviceExtensionNames) / sizeof(const char*);
    deviceCreateInfo.ppEnabledExtensionNames = deviceExtensionNames;
    deviceCreateInfo.pNext = &physical_features2;

    uint32_t deviceLayerCount;
    CHECK_VK_RESULT(vkEnumerateDeviceLayerProperties(physDevice, &deviceLayerCount, NULL));

    std::vector<VkLayerProperties> deviceLayerProperties(deviceLayerCount);
    CHECK_VK_RESULT(vkEnumerateDeviceLayerProperties(physDevice, &deviceLayerCount, deviceLayerProperties.data()));

    size_t foundDeviceLayers = 0;

    for (uint32_t i = 0; i < deviceLayerCount; i++)
    {
        for (size_t j = 0; j < sizeof(layerNames) / sizeof(const char*); j++)
        {
            if (strcmp(deviceLayerProperties[i].layerName, layerNames[j]) == 0)
            {
                foundDeviceLayers++;
            }
        }
    }

    if (foundDeviceLayers >= sizeof(layerNames) / sizeof(const char*))
    {
        deviceCreateInfo.enabledLayerCount = sizeof(layerNames) / sizeof(const char*);
        deviceCreateInfo.ppEnabledLayerNames = layerNames;
    }

    CHECK_VK_RESULT(vkCreateDevice(physDevice, &deviceCreateInfo, NULL, &device));
    vkGetDeviceQueue(device, queueFamilyIndex, 0, &queue);

    VkCommandPoolCreateInfo commandPoolcreateInfo = {};
    commandPoolcreateInfo.sType = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO;
    commandPoolcreateInfo.queueFamilyIndex = queueFamilyIndex;
    commandPoolcreateInfo.flags = VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT;
    CHECK_VK_RESULT(vkCreateCommandPool(device, &commandPoolcreateInfo, NULL, &commandPool));
}