#include "vulkan_state.h"

void CVkComputeShader::Build()
{
	VkPipelineShaderStageCreateInfo stage = {};
	VkShaderModuleCreateInfo mod = {};
	ShaderObject_t *s = m_shader.FindShaderObject(SHADER_BACKEND_SPIRV_VULKAN, SHADER_STAGE_COMPUTE);
	VulkanInputMetaData_t *pMetaData = (VulkanInputMetaData_t*)m_shader.GetLumpPtr(s->m_nMetadataLump);
	VkPipelineLayoutCreateInfo stPipelineLayout = {};
	CUtlVector<CUtlVector<VkDescriptorSetLayoutBinding>> bindings = {};

	for ( int u = 0; u < pMetaData->nDescriptorsCount; u++ )
	{
		VulkanDescriptor_t stDescriptor = ((VulkanDescriptor_t*)m_shader.GetLumpPtr(pMetaData->pDescriptorSets))[u];
		bool bFound = false;
		if (bindings.GetSize()<=stDescriptor.uSet)
			bindings.Resize(stDescriptor.uSet+1);

		uint32_t i = 0;
		for ( auto &set: bindings )
		{
			for ( auto &b: set )
			{
				if (i != stDescriptor.uSet)
					continue;
				if (b.binding != stDescriptor.uBinding)
					continue;
				bFound = true;
				break;
			}
			i++;
			if (bFound)
				break;
		}
		if (bFound)
			continue;
		VkDescriptorSetLayoutBinding bind = {};
		bind.binding = stDescriptor.uBinding;
		bind.descriptorCount = stDescriptor.uCount;
		bind.descriptorType = stDescriptor.eDescriptorType;
		bind.stageFlags = VK_SHADER_STAGE_ALL;
		bindings[stDescriptor.uSet].AppendTail(bind);
		m_bindings.AppendTail(stDescriptor);
	}
	mod.sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO;
	mod.pCode = (uint32_t*)m_shader.GetLumpPtr(s->m_nDataLump);
	mod.codeSize = m_shader.GetLumpSize(s->m_nDataLump);
	stage.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
	stage.pNext = &mod;
	stage.pName = "main";
	stage.stage = VK_SHADER_STAGE_COMPUTE_BIT;

	if ( bindings.GetSize() >= 0 )
	{
		m_setLayouts.Reserve(bindings.GetSize());
		for ( int u = 0; u < bindings.GetSize(); u++ )
		{
			VkDescriptorSetLayoutCreateInfo stSetLayoutCreateInfo = {};
			VkDescriptorSetLayout l = NULL;
			stSetLayoutCreateInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
			stSetLayoutCreateInfo.pBindings = bindings[u].GetData();
			stSetLayoutCreateInfo.bindingCount = bindings[u].GetSize();
			stSetLayoutCreateInfo.flags = VK_DESCRIPTOR_SET_LAYOUT_CREATE_UPDATE_AFTER_BIND_POOL_BIT;
			vkCreateDescriptorSetLayout(m_hDevice, &stSetLayoutCreateInfo, NULL, &l);
			m_setLayouts.AppendTail(l);
		}
		stPipelineLayout.setLayoutCount = m_setLayouts.GetSize();
		stPipelineLayout.pSetLayouts = m_setLayouts.GetData();
	}

	stPipelineLayout.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
	vkCreatePipelineLayout(m_hDevice, &stPipelineLayout, NULL, &m_hPipelineLayout);

	VkComputePipelineCreateInfo ci = {};
	ci.sType = VK_STRUCTURE_TYPE_COMPUTE_PIPELINE_CREATE_INFO;
	ci.layout = m_hPipelineLayout;
	ci.stage = stage;

	vkCreateComputePipelines(m_hDevice, NULL, 1, &ci, NULL, &m_hPipeline);

}