/*
Crafter®.Graphics
Copyright (C) 2026 Catcrafts®
catcrafts.net

This library is free software; you can redistribute it and/or
modify it under the terms of the GNU Lesser General Public
License version 3.0 as published by the Free Software Foundation;

This library is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
Lesser General Public License for more details.

You should have received a copy of the GNU Lesser General Public
License along with this library; if not, write to the Free Software
Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301 USA
*/

module;
#ifdef CRAFTER_GRAPHICS_VULKAN
#include <vulkan/vulkan.h>
#endif
export module Crafter.Graphics:DescriptorPoolVulkan;
#ifdef CRAFTER_GRAPHICS_VULKAN
import std;
import :VulkanDevice;
import :Types;
import :DescriptorSetLayoutVulkan;
import Crafter.Event;

export namespace Crafter {
    struct DescriptorEntry {
        VkDescriptorType type;
        bool occured = true;
    };

    class DescriptorPool {
        public:
        std::vector<VkDescriptorSet> sets;
        VkDescriptorPool descriptorPool = VK_NULL_HANDLE;

        template <typename Shader>
        consteval static void GetOccuringDescriptors(std::array<DescriptorEntry, 20>& types) {
            for (const VkDescriptorSetLayoutBinding& binding : Shader::descriptors) {
                for (DescriptorEntry& type : types) {
                    if (type.type == binding.descriptorType) {
                        type.occured = true;
                    }
                }
            }
        }

        template <typename... Shaders>
        consteval static std::uint32_t GetUniqueDiscriptorCount() {
            std::array<DescriptorEntry, 20> types = {{{VK_DESCRIPTOR_TYPE_SAMPLER, 0},{VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, 0},{VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE, 0},{VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, 0},{VK_DESCRIPTOR_TYPE_UNIFORM_TEXEL_BUFFER, 0},{VK_DESCRIPTOR_TYPE_STORAGE_TEXEL_BUFFER, 0},{VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 0},{VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 0},{VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER_DYNAMIC, 0},{VK_DESCRIPTOR_TYPE_STORAGE_BUFFER_DYNAMIC, 0},{VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT, 0},{VK_DESCRIPTOR_TYPE_INLINE_UNIFORM_BLOCK, 0},{VK_DESCRIPTOR_TYPE_ACCELERATION_STRUCTURE_KHR, 0},{VK_DESCRIPTOR_TYPE_ACCELERATION_STRUCTURE_NV, 0},{VK_DESCRIPTOR_TYPE_SAMPLE_WEIGHT_IMAGE_QCOM, 0},{VK_DESCRIPTOR_TYPE_BLOCK_MATCH_IMAGE_QCOM, 0},{VK_DESCRIPTOR_TYPE_MUTABLE_EXT, 0},{VK_DESCRIPTOR_TYPE_PARTITIONED_ACCELERATION_STRUCTURE_NV, 0},{VK_DESCRIPTOR_TYPE_INLINE_UNIFORM_BLOCK_EXT, 0},{VK_DESCRIPTOR_TYPE_MUTABLE_VALVE, 0}}};
    
            (GetOccuringDescriptors<Shaders>(types), ... );
    
            std::uint32_t size = 0;
            for(DescriptorEntry& type : types) {
                if(type.occured) {
                    size++;
                }
            }
    
            return size;
        }
    
        template <typename... Shaders>
        consteval static std::array<VkDescriptorPoolSize, GetUniqueDiscriptorCount<Shaders...>()> GetPoolSizes() {
            std::array<VkDescriptorPoolSize, GetUniqueDiscriptorCount<Shaders...>()> types = {};
            for(std::uint32_t i = 0; i < GetUniqueDiscriptorCount<Shaders...>(); i++){
                types[i].descriptorCount = 12345;
            }

            ([&] {
                    for (const VkDescriptorSetLayoutBinding& binding : Shaders::descriptors) {
                        bool found = false;
                        for(VkDescriptorPoolSize& type : types) {
                            if(type.type == binding.descriptorType && type.descriptorCount != 12345) {
                                type.descriptorCount += binding.descriptorCount;
                                found = true;
                            }
                        }
                        if(!found) {
                            for(std::uint32_t i = 0; i <  GetUniqueDiscriptorCount<Shaders...>(); i++){
                                if(types[i].descriptorCount == 12345) {
                                    types[i].type = binding.descriptorType;
                                    types[i].descriptorCount = binding.descriptorCount;
                                    break;
                                }
                            }
                        }
                    }
                }(), 
            ...);
    
            return types;
        }

        constexpr static std::vector<VkDescriptorPoolSize> GetPoolSizes(const std::span<const DescriptorSetLayoutVulkan> shaders) {
            std::vector<VkDescriptorPoolSize> types;

            for(const DescriptorSetLayoutVulkan& shader : shaders) {
                for (const VkDescriptorSetLayoutBinding& binding : shader.descriptors) {
                    for(VkDescriptorPoolSize& type : types) {
                        if(type.type == binding.descriptorType) {
                            type.descriptorCount += binding.descriptorCount;
                            goto inner;
                        }
                    }
                    types.emplace_back(binding.descriptorType, binding.descriptorCount);
                    inner:;
                }
            }
    
            return types;
        }
    
        public:
        void BuildPool(std::span<const VkDescriptorPoolSize> poolSizes, std::span<const VkDescriptorSetLayout> setLayouts) {
            if(descriptorPool != VK_NULL_HANDLE) {
                vkDestroyDescriptorPool(VulkanDevice::device, descriptorPool, nullptr);
            }


            VkDescriptorPoolCreateInfo descriptorPoolInfo {
                .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO,
                .maxSets = static_cast<std::uint32_t>(sets.size()),
                .poolSizeCount = static_cast<std::uint32_t>(poolSizes.size()),
                .pPoolSizes = poolSizes.data()
            };

            VulkanDevice::CheckVkResult(vkCreateDescriptorPool(VulkanDevice::device, &descriptorPoolInfo, nullptr, &descriptorPool));

            VkDescriptorSetAllocateInfo allocInfo {
                .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO,
                .descriptorPool = descriptorPool,
                .descriptorSetCount = static_cast<std::uint32_t>(sets.size()),
                .pSetLayouts = setLayouts.data(),
            };

            VulkanDevice::CheckVkResult(vkAllocateDescriptorSets(VulkanDevice::device, &allocInfo, sets.data()));
        }
    };
}

#endif