package net.halbear.Terrain4J.EngineCore.Vulkan.Structure;

import net.halbear.Terrain4J.EngineCore.Logic.Rendering.VulkanContext;
import net.halbear.Terrain4J.EngineCore.Vulkan.Util.VulkanUtils;
import org.lwjgl.PointerBuffer;
import org.lwjgl.glfw.GLFWVulkan;
import org.lwjgl.system.MemoryStack;
import org.lwjgl.system.MemoryUtil;
import org.lwjgl.vulkan.*;
import org.tinylog.Logger;

import java.nio.ByteBuffer;
import java.nio.IntBuffer;
import java.nio.LongBuffer;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

import static net.halbear.Terrain4J.EngineCore.Vulkan.Util.VulkanUtils.vkCheck;
import static org.lwjgl.vulkan.EXTDebugUtils.*;
import static org.lwjgl.vulkan.VK13.*;


public class VulkanInstance {
    public static final int MESSAGE_SEVERITY_BITMASK = VK_DEBUG_UTILS_MESSAGE_SEVERITY_ERROR_BIT_EXT |
            VK_DEBUG_UTILS_MESSAGE_SEVERITY_WARNING_BIT_EXT;
    public static final int MESSAGE_TYPE_BITMASK = VK_DEBUG_UTILS_MESSAGE_TYPE_GENERAL_BIT_EXT |
            VK_DEBUG_UTILS_MESSAGE_TYPE_VALIDATION_BIT_EXT |
            VK_DEBUG_UTILS_MESSAGE_TYPE_PERFORMANCE_BIT_EXT;
    private static final String DBG_CALL_BACK_PREF = "VkDebugUtilsCallback, {}";
    private static final String VALIDATION_LAYER = "VK_LAYER_KHRONOS_validation";
    private static final String PORTABILITY_EXTENSION = "VK_KHR_portability_enumeration";
    private final VkInstance vkInstance;
    private VkDebugUtilsMessengerCreateInfoEXT DebugUtils;
    private long DebugHandle;

    public VulkanInstance(boolean validate){
        Logger.debug("Initialising Vulkan Instance");
        try (MemoryStack MemStack = MemoryStack.stackPush()) {
            ByteBuffer appShortName = MemStack.UTF8("Terrain4J");
            var appInfo = VkApplicationInfo.calloc(MemStack).sType$Default().pApplicationName(appShortName).applicationVersion(1).pEngineName(appShortName).engineVersion(3).apiVersion(VK_API_VERSION_1_3);

            //Vulkan operates in "Layers" and does minimal validation as it expects the developer to be competent, so we add these ourseves for better debugging
            //refer to https://vulkan-tutorial.com/Drawing_a_triangle/Setup/Validation_layers for an explanation
            List<String> ValidationLayers = getSupportedValidationLayers();
            int ValidationLayerCount = ValidationLayers.size();
            boolean SupportsValidation = validate;
            if (validate && ValidationLayerCount == 0){
                SupportsValidation = false;
                Logger.warn("No Layers Found to Validate, continuing without validation");
            }
            Logger.debug("Validation:{}", SupportsValidation);
            //once its done the validation layers, transform it into a pointer (null terminated string to be exact)
            PointerBuffer RequiredLayers = null;
            if (SupportsValidation){
                RequiredLayers = MemStack.mallocPointer(ValidationLayerCount);
                for(int i = 0; i < ValidationLayerCount; i++){
                    Logger.debug("Using Validation Layer:{}", ValidationLayers.get(i));
                    RequiredLayers.put(i, MemStack.ASCII(ValidationLayers.get(i)));
                }
            }
            //enable MacOS Support because they use sum fuckin buuuullshiiiit
            Set<String> InstanceExtensions = getSupportedInstanceExtensions();
            boolean CrossPlatform = InstanceExtensions.contains(PORTABILITY_EXTENSION) && VulkanUtils.getOS() == VulkanUtils.OSType.MACOS;

            PointerBuffer GLFW_Extensions = GLFWVulkan.glfwGetRequiredInstanceExtensions();
            if (GLFW_Extensions == null){
                throw new RuntimeException("Failed to find the GLFW platform surface extensions for vulkan");
            }

            var ExtraExtensions = new ArrayList<ByteBuffer>();
            if(SupportsValidation){
                ExtraExtensions.add(MemStack.UTF8(EXTDebugUtils.VK_EXT_DEBUG_UTILS_EXTENSION_NAME));
            }
            if(CrossPlatform){
                ExtraExtensions.add(MemStack.UTF8(PORTABILITY_EXTENSION));
            }
            int ExtraExtensionsCount = ExtraExtensions.size();
            PointerBuffer RequiredExtensions = MemStack.mallocPointer(GLFW_Extensions.remaining() + ExtraExtensionsCount);
            RequiredExtensions.put(GLFW_Extensions);
            for(int i = 0; i < ExtraExtensionsCount; i++){
                RequiredExtensions.put(ExtraExtensions.get(i));
            }
            RequiredExtensions.flip();
            long extension = MemoryUtil.NULL;
            if(SupportsValidation){
                DebugUtils = CreateDebugCallBack();
                extension = DebugUtils.address();
            }
            var InstanceInfo = VkInstanceCreateInfo.calloc(MemStack)
                    .sType$Default()
                    .pApplicationInfo(appInfo)
                    .ppEnabledLayerNames(RequiredLayers)
                    .ppEnabledExtensionNames(RequiredExtensions);
            if(CrossPlatform){
                InstanceInfo.flags(0x00000001);
            }

            PointerBuffer PointerInstance = MemStack.mallocPointer(1);
            vkCheck(vkCreateInstance(InstanceInfo,null,PointerInstance), "Error Creating Vulkan Instance");
            vkInstance = new VkInstance(PointerInstance.get(0),InstanceInfo);

            DebugHandle = VK_NULL_HANDLE;
            if(SupportsValidation){
                LongBuffer LongBuffer = MemStack.mallocLong(1);
                vkCheck(vkCreateDebugUtilsMessengerEXT(vkInstance, DebugUtils, null, LongBuffer), "Error Creating Handle");
                DebugHandle = LongBuffer.get(0);
            }
        }
    }
    private Set<String> getSupportedInstanceExtensions(){
        Set<String> InstanceExtensions = new HashSet<>();
        try(var MemStack = MemoryStack.stackPush()){
            IntBuffer ExtensionsBuffer = MemStack.callocInt(1);
            vkEnumerateInstanceExtensionProperties((String)null,ExtensionsBuffer, null);
            int ExtensionsCount = ExtensionsBuffer.get(0);
            Logger.debug("Instance Extension Count:[{}]", ExtensionsCount);

            var InstanceExtensionsProps = VkExtensionProperties.calloc(ExtensionsCount, MemStack);
            vkEnumerateInstanceExtensionProperties((String)null,ExtensionsBuffer, InstanceExtensionsProps);
            for(int i = 0; i < ExtensionsCount; i++){
                VkExtensionProperties Properties = InstanceExtensionsProps.get(i);
                String ExtensionName = Properties.extensionNameString();
                InstanceExtensions.add(ExtensionName);
                Logger.trace("Supported Instance Extension [{}]", ExtensionName);
                          }
        }
        return InstanceExtensions;
    }
    private List<String> getSupportedValidationLayers(){
        try(var MemStack = MemoryStack.stackPush()) {
            IntBuffer LayersArray = MemStack.callocInt(1);
            vkEnumerateInstanceLayerProperties(LayersArray, null);
            int Layers = LayersArray.get(0);
            Logger.debug("Validation Layer Count:{}", Layers);
            var PropertiesBuffer = VkLayerProperties.calloc(Layers, MemStack);
            vkEnumerateInstanceLayerProperties(LayersArray, PropertiesBuffer);
            List<String> SupportedLayers = new ArrayList<>();
            for (int i = 0; i < Layers; i++) {
                VkLayerProperties properties = PropertiesBuffer.get(i);
                String LayerName = properties.layerNameString();
                SupportedLayers.add(LayerName);
                Logger.trace("Supported Layer [{}]", LayerName);
            }
            List<String> LayersToUse = new ArrayList<>();
            if (SupportedLayers.contains(VALIDATION_LAYER)) {
                LayersToUse.add(VALIDATION_LAYER);
            }
            return LayersToUse;
        }
    }
    public void cleanup() {
        Logger.debug("Closing Vulkan instance");
        if(DebugHandle != VK_NULL_HANDLE){
            vkDestroyDebugUtilsMessengerEXT(vkInstance, DebugHandle, null);
        }
        vkDestroyInstance(vkInstance, null);
        if(DebugUtils != null){
            DebugUtils.pfnUserCallback().free();
            DebugUtils.free();
        }
        VulkanContext.VulkanLayersOpen--;
    }

    public VkInstance getVkInstance() {
        return vkInstance;
    }

    private static VkDebugUtilsMessengerCreateInfoEXT CreateDebugCallBack(){
        return VkDebugUtilsMessengerCreateInfoEXT
                .calloc()
                .sType$Default()
                .messageSeverity(MESSAGE_SEVERITY_BITMASK)
                .messageType(MESSAGE_TYPE_BITMASK)
                .pfnUserCallback((messageSeverity, messageTypes, pCallbackData,pUserData) -> {
                   VkDebugUtilsMessengerCallbackDataEXT callBackData = VkDebugUtilsMessengerCallbackDataEXT.create(pCallbackData);
                    if ((messageSeverity & VK_DEBUG_UTILS_MESSAGE_SEVERITY_INFO_BIT_EXT) != 0) {
                        Logger.info(DBG_CALL_BACK_PREF, callBackData.pMessageString());
                    } else if ((messageSeverity & VK_DEBUG_UTILS_MESSAGE_SEVERITY_WARNING_BIT_EXT) != 0) {
                        Logger.warn(DBG_CALL_BACK_PREF, callBackData.pMessageString());
                    } else if ((messageSeverity & VK_DEBUG_UTILS_MESSAGE_SEVERITY_ERROR_BIT_EXT) != 0) {
                        Logger.error(DBG_CALL_BACK_PREF, callBackData.pMessageString());
                    } else {
                        Logger.debug(DBG_CALL_BACK_PREF, callBackData.pMessageString());
                    }
                   return VK_FALSE;
                });
    }
}
