/*
Crafter®.Graphics
Copyright (C) 2026 Catcrafts®
catcrafts.net

This library is free software; you can redistribute it and/or
modify it under the terms of the GNU Lesser General Public
License version 3.0 as published by the Free Software Foundation;

This library is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
Lesser General Public License for more details.

You should have received a copy of the GNU Lesser General Public
License along with this library; if not, write to the Free Software
Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301 USA
*/


module;

#include <vulkan/vulkan.h>
#include <vulkan/vk_enum_string_helper.h>

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

module Crafter.Graphics:VulkanDevice_impl;
import :VulkanDevice;
import std;


using namespace Crafter;
const char* const instanceExtensionNames[] = {
    "VK_EXT_debug_utils",
    "VK_KHR_surface",
    #ifdef CRAFTER_BUILD_CONFIGURATION_TARGET_x86_64_w64_mingw32
    "VK_KHR_win32_surface"
    #else
    "VK_KHR_wayland_surface"
    #endif
};
const char* const deviceExtensionNames[] = {
    "VK_KHR_swapchain",
    "VK_KHR_spirv_1_4",
    "VK_EXT_mesh_shader",
    "VK_KHR_shader_float_controls",
    "VK_KHR_dynamic_rendering",
    "VK_KHR_acceleration_structure",
    "VK_KHR_deferred_host_operations",
    "VK_KHR_ray_tracing_pipeline",
    "VK_KHR_ray_tracing_position_fetch"
};
const char* const layerNames[] = {
    "VK_LAYER_KHRONOS_validation"
};


void VulkanDevice::CheckVkResult(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;
    CheckVkResult(vkEnumerateInstanceLayerProperties(&instanceLayerCount, NULL));

    std::vector<VkLayerProperties> instanceLayerProperties(instanceLayerCount);
    CheckVkResult(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;
    }

    CheckVkResult(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;
    CheckVkResult(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];

        VkPhysicalDeviceProperties2 properties2{
           .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_PROPERTIES_2,
           .pNext = &rayTracingProperties
        };
        vkGetPhysicalDeviceProperties2(device, &properties2);

        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;

    VkPhysicalDeviceVulkan12Features features12 {
        .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_2_FEATURES,
        .runtimeDescriptorArray = VK_TRUE,
        .bufferDeviceAddress = VK_TRUE
    };

    VkPhysicalDeviceRayTracingPositionFetchFeaturesKHR vkPhysicalDeviceRayTracingPositionFetchFeatures {
        .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_RAY_TRACING_POSITION_FETCH_FEATURES_KHR,
        .pNext = &features12,
        .rayTracingPositionFetch = VK_TRUE
    };

    VkPhysicalDeviceRayTracingPipelineFeaturesKHR physicalDeviceRayTracingPipelineFeatures{
        .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_RAY_TRACING_PIPELINE_FEATURES_KHR,
        .pNext = &vkPhysicalDeviceRayTracingPositionFetchFeatures,
        .rayTracingPipeline = VK_TRUE
    };

    VkPhysicalDeviceAccelerationStructureFeaturesKHR deviceAccelerationStructureFeature = {
        .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_ACCELERATION_STRUCTURE_FEATURES_KHR,
        .pNext = &physicalDeviceRayTracingPipelineFeatures,
        .accelerationStructure = VK_TRUE
    };

    VkPhysicalDeviceDynamicRenderingFeaturesKHR dynamicRenderingFeature = {VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_DYNAMIC_RENDERING_FEATURES_KHR};
    dynamicRenderingFeature.dynamicRendering = VK_FALSE;
    dynamicRenderingFeature.pNext = &deviceAccelerationStructureFeature;

    VkPhysicalDeviceMeshShaderFeaturesEXT ext_feature = {VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_MESH_SHADER_FEATURES_EXT};
    ext_feature.meshShader = VK_FALSE;
    ext_feature.pNext = &dynamicRenderingFeature;

    VkPhysicalDeviceFeatures2 physical_features2 = {
        .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2,
        .features = {
            .samplerAnisotropy = VK_TRUE
        }
    };
    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;
    CheckVkResult(vkEnumerateDeviceLayerProperties(physDevice, &deviceLayerCount, NULL));

    std::vector<VkLayerProperties> deviceLayerProperties(deviceLayerCount);
    CheckVkResult(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 (std::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;
    }

    CheckVkResult(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;
    CheckVkResult(vkCreateCommandPool(device, &commandPoolcreateInfo, NULL, &commandPool));

	vkGetPhysicalDeviceMemoryProperties(physDevice, &memoryProperties);

    // vkCmdDrawMeshTasksEXTProc = reinterpret_cast<PFN_vkCmdDrawMeshTasksEXT>(vkGetDeviceProcAddr(device, "vkCmdDrawMeshTasksEXT"));
    // vkCmdBeginRenderingKHRProc = reinterpret_cast<PFN_vkCmdBeginRenderingKHR>(vkGetInstanceProcAddr(instance, "vkCmdBeginRenderingKHR"));
    // vkCmdEndRenderingKHRProc = reinterpret_cast<PFN_vkCmdEndRenderingKHR>(vkGetInstanceProcAddr(instance, "vkCmdEndRenderingKHR"));
    vkGetAccelerationStructureBuildSizesKHR = reinterpret_cast<PFN_vkGetAccelerationStructureBuildSizesKHR>(vkGetInstanceProcAddr(instance, "vkGetAccelerationStructureBuildSizesKHR"));
    vkCreateAccelerationStructureKHR = reinterpret_cast<PFN_vkCreateAccelerationStructureKHR>(vkGetInstanceProcAddr(instance, "vkCreateAccelerationStructureKHR"));
    vkCmdBuildAccelerationStructuresKHR =  reinterpret_cast<PFN_vkCmdBuildAccelerationStructuresKHR>(vkGetInstanceProcAddr(instance, "vkCmdBuildAccelerationStructuresKHR"));
    vkGetAccelerationStructureDeviceAddressKHR =  reinterpret_cast<PFN_vkGetAccelerationStructureDeviceAddressKHR>(vkGetInstanceProcAddr(instance, "vkGetAccelerationStructureDeviceAddressKHR"));
    vkCreateRayTracingPipelinesKHR =  reinterpret_cast<PFN_vkCreateRayTracingPipelinesKHR>(vkGetInstanceProcAddr(instance, "vkCreateRayTracingPipelinesKHR"));
    vkGetRayTracingShaderGroupHandlesKHR =  reinterpret_cast<PFN_vkGetRayTracingShaderGroupHandlesKHR>(vkGetInstanceProcAddr(instance, "vkGetRayTracingShaderGroupHandlesKHR"));
    vkCmdTraceRaysKHR = reinterpret_cast<PFN_vkCmdTraceRaysKHR>(vkGetInstanceProcAddr(instance, "vkCmdTraceRaysKHR"));
}

std::uint32_t VulkanDevice::GetMemoryType(uint32_t typeBits, VkMemoryPropertyFlags properties) {
	for (uint32_t i = 0; i < memoryProperties.memoryTypeCount; i++)
	{
		if ((typeBits & 1) == 1)
		{
			if ((memoryProperties.memoryTypes[i].propertyFlags & properties) == properties)
			{
				return i;
			}
		}
		typeBits >>= 1;
	}

	throw std::runtime_error("Could not find a matching memory type");
}