/*
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 <array>
#include "VulkanInitializers.hpp"
#include "VulkanBuffer.h"
#include <unordered_map>

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

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

	export struct DescriptorSet {
		VkDescriptorSet set[2];
        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);
		}
	};

    export template <typename MeshShader, typename FragmentShader>
    class VulkanPipeline {
		private:
		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++;
						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++;
						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;
		}

		template <typename Shader, VkShaderStageFlagBits Flag>
		consteval static std::array<VkDescriptorSetLayoutBinding, Shader::descriptorCount> GetDescriptorSet() {
			std::array<VkDescriptorSetLayoutBinding, Shader::descriptorCount> set;

			for(std::uint32_t i = 0; i < Shader::descriptors.size(); i++) {
				set[i] = {Shader::descriptors[i].slot, Shader::descriptors[i].type, 1, Flag, nullptr};
			}

			return set;
		}
        public:
        inline static VkPipeline pipeline;
        inline static VkPipelineLayout pipelineLayout;
		inline static VkDescriptorPool descriptorPool = VK_NULL_HANDLE;
		inline static VkDescriptorSetLayout descriptorSetLayout[2];

        static void CreatePipeline() {
			// Pool
			std::array<VkDescriptorPoolSize, uniqueDescriptorCount> poolSizes = GetPoolSizes();
			VkDescriptorPoolCreateInfo descriptorPoolInfo = vks::initializers::descriptorPoolCreateInfo(uniqueDescriptorCount, poolSizes.data(), 2);
			VulkanDevice::CHECK_VK_RESULT(vkCreateDescriptorPool(VulkanDevice::device, &descriptorPoolInfo, nullptr, &descriptorPool));

			// Layout
			constexpr std::array<VkDescriptorSetLayoutBinding, MeshShader::descriptorCount> setLayoutBindingsMesh = GetDescriptorSet<MeshShader, VK_SHADER_STAGE_MESH_BIT_EXT>();
			VkDescriptorSetLayoutCreateInfo descriptorLayoutInfoMesh = vks::initializers::descriptorSetLayoutCreateInfo(setLayoutBindingsMesh.data(), MeshShader::descriptorCount);
			VulkanDevice::CHECK_VK_RESULT(vkCreateDescriptorSetLayout(VulkanDevice::device, &descriptorLayoutInfoMesh, nullptr, &descriptorSetLayout[0]));

			constexpr std::array<VkDescriptorSetLayoutBinding, FragmentShader::descriptorCount> setLayoutBindingsFragment = GetDescriptorSet<FragmentShader, VK_SHADER_STAGE_FRAGMENT_BIT>();
			VkDescriptorSetLayoutCreateInfo descriptorLayoutInfoFragment = vks::initializers::descriptorSetLayoutCreateInfo(setLayoutBindingsFragment.data(), FragmentShader::descriptorCount);
			VulkanDevice::CHECK_VK_RESULT(vkCreateDescriptorSetLayout(VulkanDevice::device, &descriptorLayoutInfoFragment, nullptr, &descriptorSetLayout[1]));

        	// Layout
			VkPipelineLayoutCreateInfo pipelineLayoutInfo = vks::initializers::pipelineLayoutCreateInfo(&descriptorSetLayout[0], 2);
			VulkanDevice::CHECK_VK_RESULT(vkCreatePipelineLayout(VulkanDevice::device, &pipelineLayoutInfo, nullptr, &pipelineLayout));

			// Pipeline
			VkPipelineInputAssemblyStateCreateInfo inputAssemblyState = vks::initializers::pipelineInputAssemblyStateCreateInfo(VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST, 0, VK_FALSE);
			VkPipelineRasterizationStateCreateInfo rasterizationState = vks::initializers::pipelineRasterizationStateCreateInfo(VK_POLYGON_MODE_FILL, VK_CULL_MODE_BACK_BIT, VK_FRONT_FACE_CLOCKWISE, 0);
			VkPipelineColorBlendAttachmentState blendAttachmentState = vks::initializers::pipelineColorBlendAttachmentState(0xf, VK_FALSE);
			VkPipelineColorBlendStateCreateInfo colorBlendState = vks::initializers::pipelineColorBlendStateCreateInfo(1, &blendAttachmentState);
			VkPipelineDepthStencilStateCreateInfo depthStencilState = vks::initializers::pipelineDepthStencilStateCreateInfo(VK_TRUE, VK_TRUE, VK_COMPARE_OP_LESS_OR_EQUAL);
			VkPipelineViewportStateCreateInfo viewportState = vks::initializers::pipelineViewportStateCreateInfo(1, 1, 0);
			VkPipelineMultisampleStateCreateInfo multisampleState = vks::initializers::pipelineMultisampleStateCreateInfo(VK_SAMPLE_COUNT_1_BIT, 0);
			std::vector<VkDynamicState> dynamicStateEnables = { VK_DYNAMIC_STATE_VIEWPORT, VK_DYNAMIC_STATE_SCISSOR };
			VkPipelineDynamicStateCreateInfo dynamicState = vks::initializers::pipelineDynamicStateCreateInfo(dynamicStateEnables);
			std::array<VkPipelineShaderStageCreateInfo, 2> shaderStages;

			VkFormat format = VK_FORMAT_R8G8B8A8_UNORM;

			VkPipelineRenderingCreateInfoKHR pipeline_create{VK_STRUCTURE_TYPE_PIPELINE_RENDERING_CREATE_INFO_KHR};
			pipeline_create.pNext                   = VK_NULL_HANDLE;
			pipeline_create.colorAttachmentCount    = 1;
			pipeline_create.pColorAttachmentFormats = &format;
			pipeline_create.depthAttachmentFormat   = VulkanDevice::depthFormat;

			VkGraphicsPipelineCreateInfo pipelineCI = vks::initializers::pipelineCreateInfo(pipelineLayout, VK_NULL_HANDLE, 0);

			pipelineCI.pNext = &pipeline_create;
			pipelineCI.pRasterizationState = &rasterizationState;
			pipelineCI.pColorBlendState = &colorBlendState;
			pipelineCI.pMultisampleState = &multisampleState;
			pipelineCI.pViewportState = &viewportState;
			pipelineCI.pDepthStencilState = &depthStencilState;
			pipelineCI.pDynamicState = &dynamicState;
			pipelineCI.stageCount = static_cast<uint32_t>(shaderStages.size());
			pipelineCI.pStages = shaderStages.data();

			// Not using a vertex shader, mesh shading doesn't require vertex input state
			pipelineCI.pInputAssemblyState = nullptr;
			pipelineCI.pVertexInputState = nullptr;

			//Mesh stage of the pipeline
			shaderStages[0].sType  = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
			shaderStages[0].stage  = MeshShader::_stage;
			shaderStages[0].module = MeshShader::shader;
			shaderStages[0].pName  = MeshShader::_entrypoint.value;
			shaderStages[0].flags = 0;
			shaderStages[0].pSpecializationInfo = nullptr;
			shaderStages[0].pNext = nullptr;
	
			// Fragment stage of the pipeline
			shaderStages[1].sType  = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
			shaderStages[1].stage  = FragmentShader::_stage;
			shaderStages[1].module = FragmentShader::shader;
			shaderStages[1].pName  = FragmentShader::_entrypoint.value;
			shaderStages[1].flags = 0;
			shaderStages[1].pSpecializationInfo = nullptr;
			shaderStages[1].pNext = nullptr;
				
			VulkanDevice::CHECK_VK_RESULT(vkCreateGraphicsPipelines(VulkanDevice::device, VK_NULL_HANDLE, 1, &pipelineCI, nullptr, &pipeline));
        }

		static void GetDescriptorSet(DescriptorSet& set) {
			VkDescriptorSetAllocateInfo allocInfo = vks::initializers::descriptorSetAllocateInfo(descriptorPool, &descriptorSetLayout[0], 2);
			VulkanDevice::CHECK_VK_RESULT(vkAllocateDescriptorSets(VulkanDevice::device, &allocInfo, set.set));
		}
    };
}