/*
Crafter®.Graphics
Copyright (C) 2025 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 as published by the Free Software Foundation; either
version 3.0 of the License, or (at your option) any later version.

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;

#include <cstdint>
#include <vulkan/vulkan.h>
#include <cstring>
#include <iostream>
#include <vector>
#include "../../lib/VulkanInitializers.hpp"

export module Crafter.Graphics:DescriptorSet;
import Crafter.Event;
import :VulkanDevice;
import :VulkanShader;
import :WindowWaylandVulkan;
import :VulkanPipeline;

namespace Crafter {
    export struct DescriptorEntry {
        VkDescriptorType type;
        bool occured = true;
    };
    
    /**
    * @brief A holder for descriptor sets.
    */
    export template <typename MeshShader, typename FragmentShader>
    class DescriptorSet {
        public:
        VkDescriptorSet set[2];
        /**
        * @brief This event is triggered when a descriptor is aded to this static set which invalidates all previous descriptors, subscribe to this event to renew them.
        */
        inline static Event<void> onDescriptorRefresh;
        inline static std::vector<DescriptorSet*> sets;
        inline static VkDescriptorPool descriptorPool = VK_NULL_HANDLE;

            
        consteval static std::uint32_t GetUniqueDiscriptorCount() {
            DescriptorEntry 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}};
    
            for(const DescriptorBinding& binding : MeshShader::descriptors) {
                for(DescriptorEntry& type : types) {
                    if(type.type == binding.type) {
                        type.occured = true;
                    }
                }
            }
    
            for(const DescriptorBinding& binding : FragmentShader::descriptors) {
                for(DescriptorEntry& type : types) {
                    if(type.type == binding.type) {
                        type.occured = true;
                    }
                }
            }
    
            std::uint32_t size = 0;
            for(DescriptorEntry& type : types) {
                if(type.occured) {
                    size++;
                }
            }
    
            return size;
        }
        constexpr static std::uint32_t uniqueDescriptorCount = GetUniqueDiscriptorCount();
    
        consteval static std::array<VkDescriptorPoolSize, uniqueDescriptorCount> GetPoolSizes() {
            std::array<VkDescriptorPoolSize, uniqueDescriptorCount> types = {};
            for(std::uint32_t i = 0; i < uniqueDescriptorCount; i++){
                types[i].descriptorCount = 12345;
            }
    
            for(const DescriptorBinding& binding : MeshShader::descriptors) {
                bool found = false;
                for(VkDescriptorPoolSize& type : types) {
                    if(type.type == binding.type && type.descriptorCount != 12345) {
                        type.descriptorCount += 1;
                        found = true;
                    }
                }
                if(!found) {
                    for(std::uint32_t i = 0; i < uniqueDescriptorCount; i++){
                        if(types[i].descriptorCount == 12345) {
                            types[i].type = binding.type;
                            types[i].descriptorCount = 1;
                            break;
                        }
                    }
                }
            }
    
            for(const DescriptorBinding& binding : FragmentShader::descriptors) {
                bool found = false;
                for(VkDescriptorPoolSize& type : types) {
                    if(type.type == binding.type) {
                        type.descriptorCount += 1;
                        found = true;
                    }
                }
                if(!found) {
                    for(std::uint32_t i = 0; i < uniqueDescriptorCount; i++){
                        if(types[i].descriptorCount == 12345) {
                            types[i].type = binding.type;
                            types[i].descriptorCount = 1;
                            break;
                        }
                    }
                }
            }
    
            return types;
        }
        
        DescriptorSet() {
            sets.push_back(this);
    
            if(descriptorPool != VK_NULL_HANDLE) {
                vkDestroyDescriptorPool(VulkanDevice::device, descriptorPool, nullptr);
            }
    
            std::array<VkDescriptorPoolSize, uniqueDescriptorCount> poolSizes = GetPoolSizes();
            for(VkDescriptorPoolSize& size : poolSizes) {
                size.descriptorCount *= sets.size();
            }
            VkDescriptorPoolCreateInfo descriptorPoolInfo = vks::initializers::descriptorPoolCreateInfo(uniqueDescriptorCount, poolSizes.data(), sets.size()*2);
            VulkanDevice::CHECK_VK_RESULT(vkCreateDescriptorPool(VulkanDevice::device, &descriptorPoolInfo, nullptr, &descriptorPool));
    
            for(DescriptorSet* set : sets) {
                VkDescriptorSetAllocateInfo allocInfo = vks::initializers::descriptorSetAllocateInfo(descriptorPool, &VulkanPipeline<MeshShader, FragmentShader>::descriptorSetLayout[0], 2);
                VulkanDevice::CHECK_VK_RESULT(vkAllocateDescriptorSets(VulkanDevice::device, &allocInfo, set->set));
            }
            onDescriptorRefresh.Invoke();
        }

        ~DescriptorSet() {
            sets.erase(find(sets.begin(), sets.end(), this));
        }
    
        // void Write(VkWriteDescriptorSet* descriptors, std::uint32_t count) {
        //     vkUpdateDescriptorSets(VulkanDevice::device, count, descriptors, 0, nullptr);
        // }
        // void Write(std::uint32_t stage, VkDescriptorType type, std::uint32_t binding, VkDescriptorBufferInfo* buffer) {
        //     VkWriteDescriptorSet write = vks::initializers::writeDescriptorSet(set[stage], type, binding, buffer);
        //     vkUpdateDescriptorSets(VulkanDevice::device, 1, &write, 0, nullptr);
        // }
        // void Write(std::uint32_t stage, VkDescriptorType type, std::uint32_t binding, VkDescriptorImageInfo* buffer) {
        //     VkWriteDescriptorSet write = vks::initializers::writeDescriptorSet(set[stage], type, binding, buffer);
        //     vkUpdateDescriptorSets(VulkanDevice::device, 1, &write, 0, nullptr);
        // }
    };
}