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

import net.halbear.Terrain4J.EngineCore.Vulkan.Structure.DeviceLayers.Device;
import org.lwjgl.system.MemoryStack;
import org.lwjgl.vulkan.VkDescriptorPoolCreateInfo;
import org.lwjgl.vulkan.VkDescriptorPoolSize;
import org.tinylog.Logger;

import java.nio.LongBuffer;
import java.util.List;

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

public class DescriptorPool {
    private final long VkDescriptorPool;
    private List<DescriptorTypeCount> DescriptorTypeCounts;

    public DescriptorPool(Device device, List<DescriptorTypeCount> DescriptorTypeCounts){
        Logger.debug("Creating Descriptor pool");
        this.DescriptorTypeCounts = DescriptorTypeCounts;
        try(var MemStack = MemoryStack.stackPush()){
            int MaxSets = 0;
            int TypeNum = DescriptorTypeCounts.size();
            var TypeCounts = VkDescriptorPoolSize.calloc(TypeNum,MemStack);
            for(int i = 0; i < TypeNum; i++){
                MaxSets+= DescriptorTypeCounts.get(i).Count();
                TypeCounts.get(i)
                        .type(DescriptorTypeCounts.get(i).DescriptorType())
                        .descriptorCount(DescriptorTypeCounts.get(i).Count());
            }
            var DescriptorPoolInfo = VkDescriptorPoolCreateInfo.calloc(MemStack)
                    .sType$Default()
                    .flags(VK_DESCRIPTOR_POOL_CREATE_FREE_DESCRIPTOR_SET_BIT)
                    .pPoolSizes(TypeCounts)
                    .maxSets(MaxSets);
            LongBuffer DescriptorPoolPointer = MemStack.mallocLong(1);
            vkCheck(vkCreateDescriptorPool(device.FetchVulkanDevice(), DescriptorPoolInfo, null, DescriptorPoolPointer)
                    ,"Failed to create descriptor pool");
            VkDescriptorPool = DescriptorPoolPointer.get(0);
        }
    }

    public void FreeDescriptorSet(Device device, Long VkDescriptorSet){
        try(var MemStack = MemoryStack.stackPush()){
            LongBuffer longBuffer = MemStack.mallocLong(1);
            longBuffer.put(0,VkDescriptorSet);
            vkCheck(vkFreeDescriptorSets(device.FetchVulkanDevice(), this.VkDescriptorPool, longBuffer),"Failed to free Descriptor Set");
        }
    }

    public List<DescriptorTypeCount> GetDescriptorTypeCounts(){return DescriptorTypeCounts;}
    public long GetVkDescriptorPool(){return VkDescriptorPool;}
    public void CleanUp(Device device){
        Logger.debug("Destroying Descriptor Pool");
        vkDestroyDescriptorPool(device.FetchVulkanDevice(), VkDescriptorPool, null);
    }

    public record DescriptorTypeCount(int DescriptorType, int Count){}
}
