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

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

import java.nio.ByteBuffer;
import java.nio.FloatBuffer;
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.KHRPortabilitySubset.VK_KHR_PORTABILITY_SUBSET_EXTENSION_NAME;
import static org.lwjgl.vulkan.VK13.*;

public class Device {
    private final VkDevice VulkanDevice;
    private final boolean SamplesAnisotropy;

    public Device(PhysicalDevice PhysDevice){
        Logger.debug("Creating Device Interface Layer");
        try (var MemStack = MemoryStack.stackPush()) {
            PointerBuffer RequiredExtensions = CreateRequiredExtensions(PhysDevice, MemStack);
            var QueuePropertyBuffer = PhysDevice.GetQueueFamilyProperties();
            int QueueFamilyCount = QueuePropertyBuffer.capacity();
            var QueueCreationInfoBuffer = VkDeviceQueueCreateInfo.calloc(QueueFamilyCount, MemStack);
            for(int i = 0; i < QueueFamilyCount; i++){
                FloatBuffer Priorities = MemStack.callocFloat(QueuePropertyBuffer.get(i).queueCount());
                QueueCreationInfoBuffer.get(i)
                        .sType$Default()
                        .queueFamilyIndex(i)
                        .pQueuePriorities(Priorities);
            }

            var features13 = VkPhysicalDeviceVulkan13Features.calloc(MemStack)
                    .sType$Default()
                    .dynamicRendering(true)
                    .synchronization2(true);
            var features2 = VkPhysicalDeviceFeatures2.calloc(MemStack).sType$Default();
            var features = features2.features();

            VkPhysicalDeviceFeatures SupportedFeatures = PhysDevice.GetPhysicalDeviceFeatures();
            SamplesAnisotropy = SupportedFeatures.samplerAnisotropy();
            if(SamplesAnisotropy){
                features.samplerAnisotropy(true);
            }
            features2.pNext(features13.address());

            var DeviceCreateInfo = VkDeviceCreateInfo.calloc(MemStack)
                    .sType$Default()
                    .pNext(features2.address())
                    .ppEnabledExtensionNames(RequiredExtensions)
                    .pQueueCreateInfos(QueueCreationInfoBuffer);
            PointerBuffer pp = MemStack.mallocPointer(1);

            VkPhysicalDeviceProperties deviceProps = VkPhysicalDeviceProperties.malloc();
            vkGetPhysicalDeviceProperties(PhysDevice.GetPhysicalDevice(), deviceProps);

            if (!deviceProps.limits().timestampComputeAndGraphics()) {
                Logger.warn("GPU does not support timestamp queries!");
            }

            vkCheck(vkCreateDevice(PhysDevice.GetPhysicalDevice(), DeviceCreateInfo, null, pp), "Failed to create device");
            VulkanDevice = new VkDevice(pp.get(0), PhysDevice.GetPhysicalDevice(), DeviceCreateInfo);
        }
    }

    private static PointerBuffer CreateRequiredExtensions(PhysicalDevice PhysDevice, MemoryStack MemStack){
        Set<String> DeviceExtensions = GetDeviceExtensions(PhysDevice);
        //cross platform support for that fuckass Metal API mac uses
        boolean UsePortability = DeviceExtensions.contains(VK_KHR_PORTABILITY_SUBSET_EXTENSION_NAME) && VulkanUtils.getOS() == VulkanUtils.OSType.MACOS;

        var ExtentionList = new ArrayList<ByteBuffer>();
        for(String Extension : PhysicalDevice.REQUIRED_EXTENSIONS){
            ExtentionList.add(MemStack.ASCII(Extension));
        }
        if (UsePortability){
            ExtentionList.add(MemStack.ASCII(VK_KHR_PORTABILITY_SUBSET_EXTENSION_NAME));
        }

        PointerBuffer RequiredExtensions = MemStack.mallocPointer(ExtentionList.size());
        ExtentionList.forEach(RequiredExtensions::put);
        RequiredExtensions.flip();

        return RequiredExtensions;
    }

    private static Set<String> GetDeviceExtensions(PhysicalDevice PhysDevice){
        Set<String> DeviceExtensions = new HashSet<>();
        try(var MemStack = MemoryStack.stackPush()){
            IntBuffer ExtensionsBufferCount  = MemStack.callocInt(1);
            vkEnumerateDeviceExtensionProperties(PhysDevice.GetPhysicalDevice(),(String) null, ExtensionsBufferCount, null);
            int ExtensionCount = ExtensionsBufferCount.get(0);
            Logger.trace("Device supports [{}] number of extensions", ExtensionCount);

            try(var PropertiesBuffer = VkExtensionProperties.calloc(ExtensionCount)){
                vkEnumerateDeviceExtensionProperties(PhysDevice.GetPhysicalDevice(), (String) null, ExtensionsBufferCount,PropertiesBuffer);
                for(int i = 0; i < ExtensionCount; i++){
                    VkExtensionProperties Properties =  PropertiesBuffer.get(i);
                    String ExtensionName = Properties.extensionNameString();
                    DeviceExtensions.add(ExtensionName);
                    Logger.trace("[{}] is a supported extension", ExtensionName);
                }
            }
        }
       return DeviceExtensions;
    }

    public void cleanup(){
        Logger.debug("Removing Vulkan Device");
        vkDestroyDevice(VulkanDevice, null);
        VulkanContext.VulkanLayersOpen--;
    }

    public boolean SamplesAnisotropy(){return SamplesAnisotropy;}

    public VkDevice FetchVulkanDevice()
    {
        return VulkanDevice;
    }

    public void waitIdle(){
        vkDeviceWaitIdle(VulkanDevice);
    }


}


