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

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

namespace Crafter {
    export template <typename MeshShader, typename FragmentShader>
    class VulkanPipeline {
		private:
		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 VkDescriptorSetLayout descriptorSetLayout[2];

        static void CreatePipeline() {
			// 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));
        }
    };
}