module;

#include <cstdint>
#include <vulkan/vulkan.h>
#include <array>

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

namespace Crafter {
    export template <typename MeshShader, typename FragmentShader>
    class VulkanPipeline {
        public:
        inline static VkPipeline pipeline;
        inline static VkPipelineLayout layout;
        inline static VkDescriptorPool descriptor_pool;
        inline static VkDescriptorSetLayout descriptor_set_layout;
        inline static VkDescriptorSet descriptor_set;
        static void CreatePipeline() {
            VkDescriptorPoolCreateInfo descriptor_pool_create_info = {VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO};
            descriptor_pool_create_info.maxSets = 2;
            descriptor_pool_create_info.poolSizeCount = 0;
            descriptor_pool_create_info.pPoolSizes = nullptr;
            VulkanDevice::CHECK_VK_RESULT(vkCreateDescriptorPool(VulkanDevice::device, &descriptor_pool_create_info, nullptr, &descriptor_pool));
    
            VkDescriptorSetLayoutCreateInfo descriptor_layout = {VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO};
            descriptor_layout.bindingCount = 0;
            descriptor_layout.pBindings = nullptr;
            VulkanDevice::CHECK_VK_RESULT(vkCreateDescriptorSetLayout(VulkanDevice::device, &descriptor_layout, nullptr, &descriptor_set_layout));
    
            VkDescriptorSetAllocateInfo alloc_info = {VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO};
            alloc_info.descriptorPool = descriptor_pool;
            alloc_info.descriptorSetCount = 1;
            alloc_info.pSetLayouts = &descriptor_set_layout;
            VulkanDevice::CHECK_VK_RESULT(vkAllocateDescriptorSets(VulkanDevice::device, &alloc_info, &descriptor_set));

            VkPipelineLayoutCreateInfo layout_info = {VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO};
            layout_info.setLayoutCount = 1;
            layout_info.pSetLayouts = &descriptor_set_layout;
            VulkanDevice::CHECK_VK_RESULT(vkCreatePipelineLayout(VulkanDevice::device, &layout_info, nullptr, &layout));
    
            VkPipelineRasterizationStateCreateInfo raster{VK_STRUCTURE_TYPE_PIPELINE_RASTERIZATION_STATE_CREATE_INFO};
	        raster.cullMode  = VK_CULL_MODE_BACK_BIT;
	        raster.frontFace = VK_FRONT_FACE_CLOCKWISE;
	        raster.lineWidth = 1.0f;
    
            // Our attachment will write to all color channels, but no blending is enabled.
	        VkPipelineColorBlendAttachmentState blend_attachment{};
	        blend_attachment.colorWriteMask = VK_COLOR_COMPONENT_R_BIT | VK_COLOR_COMPONENT_G_BIT | VK_COLOR_COMPONENT_B_BIT | VK_COLOR_COMPONENT_A_BIT;
    
            VkPipelineColorBlendStateCreateInfo blend{VK_STRUCTURE_TYPE_PIPELINE_COLOR_BLEND_STATE_CREATE_INFO};
            blend.attachmentCount = 1;
            blend.pAttachments    = &blend_attachment;
    
            // We will have one viewport and scissor box.
	        VkPipelineViewportStateCreateInfo viewport{VK_STRUCTURE_TYPE_PIPELINE_VIEWPORT_STATE_CREATE_INFO};
	        viewport.viewportCount = 1;
	        viewport.scissorCount  = 1;
    
            // Disable all depth testing.
            VkPipelineDepthStencilStateCreateInfo depth_stencil{VK_STRUCTURE_TYPE_PIPELINE_DEPTH_STENCIL_STATE_CREATE_INFO};
    
            // No multisampling.
            VkPipelineMultisampleStateCreateInfo multisample{VK_STRUCTURE_TYPE_PIPELINE_MULTISAMPLE_STATE_CREATE_INFO};
            multisample.rasterizationSamples = VK_SAMPLE_COUNT_1_BIT;
    
            // Specify that these states will be dynamic, i.e. not part of pipeline state object.
            std::array<VkDynamicState, 2> dynamics{VK_DYNAMIC_STATE_VIEWPORT, VK_DYNAMIC_STATE_SCISSOR};
    
        	VkPipelineDynamicStateCreateInfo dynamic{VK_STRUCTURE_TYPE_PIPELINE_DYNAMIC_STATE_CREATE_INFO};
	        dynamic.pDynamicStates    = dynamics.data();
	        dynamic.dynamicStateCount = static_cast<std::uint32_t>(dynamics.size());
    
            // Load our SPIR-V shaders.
            std::array<VkPipelineShaderStageCreateInfo, 2> shader_stages{};
    
           	//Mesh stage of the pipeline
	        shader_stages[0].sType  = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
	        shader_stages[0].stage  = MeshShader::_stage;
	        shader_stages[0].module = MeshShader::shader;
	        shader_stages[0].pName  = MeshShader::_entrypoint.value;

	        // Fragment stage of the pipeline
	        shader_stages[1].sType  = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
	        shader_stages[1].stage  = FragmentShader::_stage;
	        shader_stages[1].module = FragmentShader::shader;
	        shader_stages[1].pName  = FragmentShader::_entrypoint.value;
    
            VkGraphicsPipelineCreateInfo pipe{VK_STRUCTURE_TYPE_GRAPHICS_PIPELINE_CREATE_INFO};
            pipe.stageCount          = static_cast<std::uint32_t>(shader_stages.size());
            pipe.pStages             = shader_stages.data();
            pipe.pVertexInputState   = nullptr;
            pipe.pInputAssemblyState = nullptr;
            pipe.pRasterizationState = &raster;
            pipe.pColorBlendState    = &blend;
            pipe.pMultisampleState   = &multisample;
            pipe.pViewportState      = &viewport;
            pipe.pDepthStencilState  = &depth_stencil;
            pipe.pDynamicState       = &dynamic;
    
            // We need to specify the pipeline layout and the render pass description up front as well.
	        pipe.renderPass = WindowWaylandVulkan::renderPass;
	        pipe.layout     = layout;

	        VulkanDevice::CHECK_VK_RESULT(vkCreateGraphicsPipelines(VulkanDevice::device, VK_NULL_HANDLE, 1, &pipe, nullptr, &pipeline));
        }
    };
}