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

import net.halbear.Terrain4J.EngineCore.Vulkan.Structure.DeviceLayers.Device;
import net.halbear.Terrain4J.EngineCore.Vulkan.Structure.DisplayToScreen.ImageView;
import net.halbear.Terrain4J.EngineCore.Vulkan.Util.VulkanBuffer;
import org.lwjgl.system.MemoryStack;
import org.lwjgl.vulkan.VkDescriptorBufferInfo;
import org.lwjgl.vulkan.VkDescriptorImageInfo;
import org.lwjgl.vulkan.VkDescriptorSetAllocateInfo;
import org.lwjgl.vulkan.VkWriteDescriptorSet;

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

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

public class DescriptorSet {
    protected long VkDescriptorSet;

    public DescriptorSet(Device device,DescriptorPool descriptorPool, DescriptorSetLayout descriptorSetLayout){
        try(var MemStack = MemoryStack.stackPush()){
            LongBuffer DescriptorSetLayoutPointer = MemStack.mallocLong(1);
            DescriptorSetLayoutPointer.put(0,descriptorSetLayout.GetVkDescriptorLayout());
            var AllocInfo = VkDescriptorSetAllocateInfo.calloc(MemStack)
                    .sType$Default()
                    .descriptorPool(descriptorPool.GetVkDescriptorPool())
                    .pSetLayouts(DescriptorSetLayoutPointer);
            LongBuffer DescriptorSetPointer = MemStack.mallocLong(1);
            vkCheck(vkAllocateDescriptorSets(device.FetchVulkanDevice(), AllocInfo, DescriptorSetPointer)
                    ,"Failed to create descriptor sets");
            VkDescriptorSet = DescriptorSetPointer.get(0);
        }
    }

    public void SetBuffer(Device device, VulkanBuffer VkBuffer, long Range, int Binding, int Type){
        try(var MemStack = MemoryStack.stackPush()){
            var BufferInfo = VkDescriptorBufferInfo.calloc(1,MemStack)
                    .buffer(VkBuffer.GetBuffer())
                    .offset(0)
                    .range(Range);
            var DescriptorBuffer = VkWriteDescriptorSet.calloc(1,MemStack);
            DescriptorBuffer.get(0)
                    .sType$Default()
                    .dstSet(VkDescriptorSet)
                    .dstBinding(Binding)
                    .descriptorType(Type)
                    .descriptorCount(1)
                    .pBufferInfo(BufferInfo);
            vkUpdateDescriptorSets(device.FetchVulkanDevice(),DescriptorBuffer,null);
        }
    }

    public void SetImage(Device device, ImageView imageView, TextureSampler textureSampler, int BaseBinding){
        List<ImageView> imageViews = new ArrayList<>();
        imageViews.add(imageView);
        SetImages(device, imageViews, textureSampler, BaseBinding);
    }
    public void SetImages(Device device, List<ImageView> imageViews, TextureSampler textureSampler, int BaseBinding){
        try(var MemStack = MemoryStack.stackPush()){
            int ImageCount = imageViews.size();
            var DescriptorBuffer = VkWriteDescriptorSet.calloc(ImageCount,MemStack);
            for(int i = 0; i < ImageCount; i++){
                ImageView imageView = imageViews.get(i);
                var ImageInfo = VkDescriptorImageInfo.calloc(1,MemStack)
                        .imageView(imageView.GetVulkanImageView())
                        .sampler(textureSampler.GetVkSampler());
                if(imageView.IsDepthImage()){
                    ImageInfo.imageLayout(VK_IMAGE_LAYOUT_DEPTH_STENCIL_READ_ONLY_OPTIMAL);
                } else{
                    ImageInfo.imageLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);
                }
                DescriptorBuffer.get(i)
                        .sType$Default()
                        .dstSet(VkDescriptorSet)
                        .dstBinding(BaseBinding + i)
                        .descriptorType(VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER)
                        .descriptorCount(1)
                        .pImageInfo(ImageInfo);
            }
            vkUpdateDescriptorSets(device.FetchVulkanDevice(),DescriptorBuffer,null);
        }
    }

    public void SetImageArray(Device device, List<ImageView> imageViews, TextureSampler textureSampler, int BaseBinding){
        try(var MemStack = MemoryStack.stackPush()){
            int ImageCount = imageViews.size();
            VkDescriptorImageInfo.Buffer imageInfos = VkDescriptorImageInfo.calloc(ImageCount,MemStack);
            for(int i = 0; i < ImageCount; i++){
                ImageView imageView = imageViews.get(i);
                VkDescriptorImageInfo imageInfo = imageInfos.get(i);
                imageInfo.imageView(imageView.GetVulkanImageView()).sampler(textureSampler.GetVkSampler());

                if(imageView.IsDepthImage()){
                    imageInfo.imageLayout(VK_IMAGE_LAYOUT_DEPTH_STENCIL_READ_ONLY_OPTIMAL);
                } else{
                    imageInfo.imageLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);
                }
            }

            var DescriptorBuffer = VkWriteDescriptorSet.calloc(1,MemStack);
            DescriptorBuffer.get(0)
                    .sType$Default()
                    .dstSet(VkDescriptorSet)
                    .dstBinding(BaseBinding)
                    .dstArrayElement(0)
                    .descriptorType(VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER)
                    .descriptorCount(ImageCount)
                    .pImageInfo(imageInfos);
            vkUpdateDescriptorSets(device.FetchVulkanDevice(),DescriptorBuffer,null);
        }
    }

    public long GetVkDescriptorSet(){return VkDescriptorSet;}
}
