package net.halbear.Terrain4J.EngineCore.Vulkan.Rendering.Shader;

import net.halbear.Terrain4J.EngineCore.Logic.EngineConfig;
import net.halbear.Terrain4J.EngineCore.Vulkan.Structure.DeviceLayers.Device;
import net.halbear.Terrain4J.EngineCore.Vulkan.Structure.DeviceLayers.PhysicalDevice;
import org.lwjgl.vulkan.VkPhysicalDeviceFeatures2;
import org.lwjgl.vulkan.VkPhysicalDeviceLimits;
import org.tinylog.Logger;

import java.util.*;

import static org.lwjgl.vulkan.VK10.*;

public class DescriptorAllocator {
    private final Map<Integer,Integer> DescriptorLimits;
    private final List<DescriptorPoolInfo> DescriptorPoolList;
    private final Map<String,DescriptorSetInfo> DescriptorSetInfoMap;

    public DescriptorAllocator(PhysicalDevice physicalDevice, Device device){
        Logger.debug("Creating Descriptor Allocator");
        DescriptorPoolList = new ArrayList<>();
        DescriptorLimits = CreateDescriptorLimits(physicalDevice);
        DescriptorPoolList.add(CreateDescriptorPoolInformation(device, DescriptorLimits));
        DescriptorSetInfoMap = new HashMap<>();
    }

    public synchronized DescriptorSet AddDescriptorSet(Device device, String ID, DescriptorSetLayout descriptorSetLayout){
        return AddDescriptorSets(device, ID, 1, descriptorSetLayout)[0];
    }
    public synchronized DescriptorSet[] AddDescriptorSets(Device device, String ID, int count, DescriptorSetLayout descriptorSetLayout){
        DescriptorPoolInfo TargetPool = null;
        int PoolPosition = 0;
        for(DescriptorPoolInfo descriptorPoolInfo : DescriptorPoolList){
            for(DescriptorSetLayout.LayoutInformation LayoutInfo : descriptorSetLayout.GetLayoutInfos()){
                int DescriptorType = LayoutInfo.DescriptorType();
                Integer Available = descriptorPoolInfo.DescriptorCount.get(DescriptorType);
                if(Available == null){
                    throw new RuntimeException("Unknown Type [" + DescriptorType + "]");
                }
                Integer MaxTotal = DescriptorLimits.get(DescriptorType);
                if(count > MaxTotal){
                    throw new RuntimeException("Cannot Create More than ["+MaxTotal+"] for descriptor type [" + DescriptorType + "]");
                }
                if(Available < count){
                    TargetPool = null;
                    break;
                } else{
                    TargetPool = descriptorPoolInfo;
                }

            }
            PoolPosition++;
        }
        if(TargetPool == null){
            TargetPool = CreateDescriptorPoolInformation(device, DescriptorLimits);
            DescriptorPoolList.add(TargetPool);
            PoolPosition++;
        }
        var Result = new DescriptorSet[count];
        for(int i = 0; i < count; i++){
            DescriptorSet descriptorSet = new DescriptorSet(device, TargetPool.descriptorPool(), descriptorSetLayout);
            Result[i] = descriptorSet;
        }
        DescriptorSetInfoMap.put(ID, new DescriptorSetInfo(Result, PoolPosition));

        for(DescriptorSetLayout.LayoutInformation LayoutInfo : descriptorSetLayout.GetLayoutInfos()){
            int DescriptorType = LayoutInfo.DescriptorType();
            TargetPool.DescriptorCount.put(DescriptorType,TargetPool.DescriptorCount.get(DescriptorType) - count);
        }
        return Result;
    }

    private static DescriptorPoolInfo CreateDescriptorPoolInformation(Device device, Map<Integer,Integer>DescriptorLimits){
        Map<Integer,Integer> DescriptorCount = new HashMap<>();
        List<DescriptorPool.DescriptorTypeCount> DescriptorTypeCounts = new ArrayList<>();
        DescriptorLimits.forEach((key,value)->{
            DescriptorCount.put(key,value);
            DescriptorTypeCounts.add(new DescriptorPool.DescriptorTypeCount(key,value));
        });
        var DescriptorPool = new DescriptorPool(device, DescriptorTypeCounts);
        return new DescriptorPoolInfo(DescriptorCount, DescriptorPool);
    }

    private static Map<Integer,Integer> CreateDescriptorLimits(PhysicalDevice physicalDevice){
        var EngConfig = EngineConfig.getInstance();
        int MaxDescriptors = EngConfig.GetMaxDescriptors();
        VkPhysicalDeviceLimits limits = physicalDevice.GetPhysicalDeviceProperties().properties().limits();
        Map<Integer, Integer> DescriptorLimits = new HashMap<>();
        DescriptorLimits.put(VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, (int)Math.min(MaxDescriptors, Integer.toUnsignedLong(limits.maxDescriptorSetUniformBuffers())));
        DescriptorLimits.put(VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, (int)Math.min(MaxDescriptors, Integer.toUnsignedLong(limits.maxDescriptorSetSamplers())));
        DescriptorLimits.put(VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, (int)Math.min(MaxDescriptors, Integer.toUnsignedLong(limits.maxDescriptorSetStorageBuffers())));
        return DescriptorLimits;
    }


    public synchronized void FreeDescriptorSet(Device device, String ID){
        DescriptorSetInfo descriptorSetInfo = DescriptorSetInfoMap.get(ID);
        if(descriptorSetInfo == null){
            Logger.info("Descriptor Set of ID [{}} does not exist",ID);
            return;
        }
        if(descriptorSetInfo.PoolPosition >= DescriptorPoolList.size()){
            Logger.info("Could not locate a pool associated with Descriptor Set ID [{}]",ID);
            return;
        }
        DescriptorPoolInfo descriptorPoolInfo = DescriptorPoolList.get(descriptorSetInfo.PoolPosition);
        Arrays.asList(descriptorSetInfo.DescriptorSets()).forEach(descriptorSet ->
                descriptorPoolInfo.descriptorPool.FreeDescriptorSet(device, descriptorSet.GetVkDescriptorSet()));
    }

    public synchronized DescriptorSet GetDescriptorSet(String ID, int Position){
        DescriptorSet Result = null;
        DescriptorSetInfo descriptorSetInfo = DescriptorSetInfoMap.get(ID);
        if(descriptorSetInfo != null){
            Result = descriptorSetInfo.DescriptorSets()[Position];
        }
        return Result;
    }

    public synchronized DescriptorSet GetDescriptorSet(String ID){
        return GetDescriptorSet(ID, 0);
    }

    record DescriptorPoolInfo(Map<Integer,Integer>DescriptorCount, DescriptorPool descriptorPool){}
    record DescriptorSetInfo(DescriptorSet[] DescriptorSets, int PoolPosition){}

    public synchronized void CleanUp(Device device){
        Logger.debug("Destroying Descriptor Allocator");
        DescriptorSetInfoMap.clear();
        DescriptorPoolList.forEach(descriptor -> descriptor.descriptorPool.CleanUp(device));
    }
}
