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

import net.halbear.Terrain4J.EngineCore.Logic.Rendering.VulkanContext;
import net.halbear.Terrain4J.EngineCore.Vulkan.Structure.VulkanInstance;
import org.lwjgl.PointerBuffer;
import org.lwjgl.system.MemoryStack;
import org.lwjgl.vulkan.*;
import org.tinylog.Logger;

import java.nio.IntBuffer;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.Set;

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

public class PhysicalDevice {
    private final VkExtensionProperties.Buffer vkDeviceExtensions;
    private final VkPhysicalDeviceMemoryProperties vkMemoryProperties;
    private final VkPhysicalDevice vkPhysicalDevice;
    private final VkPhysicalDeviceFeatures vkPhysicalDeviceFeatures;
    private final VkPhysicalDeviceProperties2 vkPhysicalDeviceProperties;
    private final VkQueueFamilyProperties.Buffer vkQueueFamilyProperties;

    protected static final Set<String> REQUIRED_EXTENSIONS;

    static{
        REQUIRED_EXTENSIONS = new HashSet<>();
        REQUIRED_EXTENSIONS.add(KHRSwapchain.VK_KHR_SWAPCHAIN_EXTENSION_NAME);
    }

    private PhysicalDevice(VkPhysicalDevice VulkanPhysicalDevice){
        try(var MemStack = MemoryStack.stackPush()){
            this.vkPhysicalDevice = VulkanPhysicalDevice;
            IntBuffer intBuffer = MemStack.mallocInt(1);
            vkPhysicalDeviceProperties = VkPhysicalDeviceProperties2.calloc().sType$Default();
            vkGetPhysicalDeviceProperties2(vkPhysicalDevice, vkPhysicalDeviceProperties);

            vkCheck(vkEnumerateDeviceExtensionProperties(vkPhysicalDevice, (String)null, intBuffer, null), "Failed to get extension property count");
            vkDeviceExtensions = VkExtensionProperties.calloc(intBuffer.get(0));
            vkCheck(vkEnumerateDeviceExtensionProperties(vkPhysicalDevice, (String)null, intBuffer, vkDeviceExtensions), "Failed to get extension property");

            vkGetPhysicalDeviceQueueFamilyProperties(vkPhysicalDevice, intBuffer, null);
            vkQueueFamilyProperties = VkQueueFamilyProperties.calloc(intBuffer.get(0));
            vkGetPhysicalDeviceQueueFamilyProperties(vkPhysicalDevice, intBuffer, vkQueueFamilyProperties);

            vkPhysicalDeviceFeatures = VkPhysicalDeviceFeatures.calloc();
            vkGetPhysicalDeviceFeatures(vkPhysicalDevice, vkPhysicalDeviceFeatures);

            vkMemoryProperties = VkPhysicalDeviceMemoryProperties.calloc();
            vkGetPhysicalDeviceMemoryProperties(vkPhysicalDevice, vkMemoryProperties);
        }
    }

    public static PhysicalDevice createPhysicalDevice(VulkanInstance instance, String PreferredDeviceName){
        Logger.debug("selecting physical devices");
        PhysicalDevice result = null;
        try (var MemStack = MemoryStack.stackPush()){
            PointerBuffer PointerPhyiscalDevices = getPhysicalDevices(instance, MemStack);
            int DeviceCount = PointerPhyiscalDevices.capacity();

            var PhysicalDevices = new ArrayList<PhysicalDevice>();
            for(int i = 0; i < DeviceCount; i++){
                var vkPhysicalDevice = new VkPhysicalDevice(PointerPhyiscalDevices.get(i),instance.getVkInstance());
                var PhysDevice = new PhysicalDevice(vkPhysicalDevice);

                String DeviceName = PhysDevice.FetchDeviceName();
                if(!PhysDevice.HasGraphicsQueueFamily()){
                    Logger.debug("Device [{}] does not support graphics queue family",DeviceName);
                    PhysDevice.cleanup();
                    continue;
                }

                if(!PhysDevice.SupportsExtensions(REQUIRED_EXTENSIONS)){
                    Logger.debug("Device [{}] does not support some required extensions",DeviceName);
                    PhysDevice.cleanup();
                    continue;
                }
                if(PreferredDeviceName != null && PreferredDeviceName.equals(DeviceName)){
                    result = PhysDevice;
                    break;
                }
                if(PhysDevice.vkPhysicalDeviceProperties.properties().deviceType() == VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU){
                    PhysicalDevices.addFirst(PhysDevice);
                } else{
                    PhysicalDevices.add(PhysDevice);
                }
            }
            result = result == null && !PhysicalDevices.isEmpty() ? PhysicalDevices.removeFirst() : result;
            PhysicalDevices.forEach(PhysicalDevice::cleanup);
            if (result == null){
                throw new RuntimeException("No Graphics Driver Found!");
            }
            Logger.debug("Selected Graphics Driver: [{}]", result.FetchDeviceName());
        }
        return result;
    }

    public static PointerBuffer getPhysicalDevices(VulkanInstance instance, MemoryStack MemStack){
        PointerBuffer PointerPhysicalDevices;
        IntBuffer intBuffer = MemStack.mallocInt(1);
        vkCheck(vkEnumeratePhysicalDevices(instance.getVkInstance(),intBuffer,null),"Failed to fetch physical device count");
        int DeviceCount = intBuffer.get(0);
        Logger.debug("Detected Physical Device Count: {}",DeviceCount);
        PointerPhysicalDevices = MemStack.mallocPointer(DeviceCount);
        vkCheck(vkEnumeratePhysicalDevices(instance.getVkInstance(),intBuffer,PointerPhysicalDevices),"Failed to fetch physical devices");
        return PointerPhysicalDevices;
    }
    //Display Drivers have specific queues for specific tasks organised into "families", this makes sure the gpu can push queues that process graphics
    private boolean HasGraphicsQueueFamily(){
        boolean CanHanldeGraphicsQueues = false;
        int FamilyCount = vkQueueFamilyProperties != null ? vkQueueFamilyProperties.capacity() : 0;
        for(int i = 0; i < FamilyCount; i ++){
            VkQueueFamilyProperties FamilyProperties = vkQueueFamilyProperties.get(i);
            if((FamilyProperties.queueFlags() & VK_QUEUE_GRAPHICS_BIT) != 0){
                CanHanldeGraphicsQueues = true;
                break;
            }
        }
        return CanHanldeGraphicsQueues;
    }

    public boolean SupportsExtensions(Set<String> Extensions){
        var ExtensionsClone = new HashSet<>(Extensions);
        int ExtensionCount = vkDeviceExtensions != null ? vkDeviceExtensions.capacity() : 0;
        for(int i = 0; i < ExtensionCount; i++){
            String ExtensionName = vkDeviceExtensions.get(i).extensionNameString();
            ExtensionsClone.remove(ExtensionName);
        }

        boolean AllSupported = ExtensionsClone.isEmpty();
        if (!AllSupported){
            Logger.debug("A detected [{}] number of extensions are not compatible with the device [{}]",ExtensionsClone.iterator().next(), FetchDeviceName());
        }
        return AllSupported;
    }

    public void cleanup(){
        Logger.debug("freeing physical device [{}]", FetchDeviceName());
        vkMemoryProperties.free();
        vkPhysicalDeviceFeatures.free();
        vkQueueFamilyProperties.free();
        vkDeviceExtensions.free();
        vkPhysicalDeviceProperties.free();
        VulkanContext.VulkanLayersOpen--;
    }

    public String FetchDeviceName(){
        return vkPhysicalDeviceProperties.properties().deviceNameString();
    }
    public VkPhysicalDeviceMemoryProperties GetMemoryProperties(){
        return vkMemoryProperties;
    }

    public VkPhysicalDevice GetPhysicalDevice(){
        return vkPhysicalDevice;
    }
    public VkPhysicalDeviceFeatures GetPhysicalDeviceFeatures(){
        return vkPhysicalDeviceFeatures;
    }
    public VkPhysicalDeviceProperties2 GetPhysicalDeviceProperties(){
        return vkPhysicalDeviceProperties;
    }
    public VkQueueFamilyProperties.Buffer GetQueueFamilyProperties(){
        return vkQueueFamilyProperties;
    }
}
