#include "materialsystem/shaderinternals.h"
#include "tier1/utlvector.h"
#include "vulkan_state.h"
#include "shaderparser.h"
CVkShader::~CVkShader()
{

}
void CVkShader::AddLayout( int iIndex, int iStride )
{
	VkVertexInputBindingDescription layout = {};
	layout.binding = iIndex;
	layout.inputRate = VK_VERTEX_INPUT_RATE_VERTEX;
	layout.stride = iStride;
	m_layouts.AppendTail(layout);
}

void CVkShader::AddAttribute( int iBufferIndex, int iLocation, EVertexFormat eFormat, int iOffset )
{
	VkVertexInputAttributeDescription attribute = {};
	attribute.binding = iBufferIndex;
	attribute.location = iLocation;
	attribute.format = VulkanGetVertexFormat(eFormat);
	attribute.offset = iOffset;
	m_attributes.AppendTail(attribute);
}

void CVkShader::SetTopology( ETopologyMode eTopology )
{

}

void CVkShader::AddOutputImage( int iImageIndex, EImageFormat eFormat )
{
	m_eFormats.AppendTail(CVkImage::GetImageFormat(eFormat));
}

void CVkShader::SetDepthImage( EImageFormat eFormat )
{
	m_eDepthFormat = CVkImage::GetImageFormat(eFormat);
}
void CVkShader::DisablePixelShader( bool bDisable )
{

}


void CVkShader::Build()
{
	VkPipelineLayoutCreateInfo layoutInfo = {};
	VkGraphicsPipelineCreateInfo createInfo = {};
	CUtlVector<VkPipelineShaderStageCreateInfo> stages = {};
	CUtlVector<VkShaderModuleCreateInfo> modules = {};
	VkPipelineVertexInputStateCreateInfo vertexInput = {};
	VkPipelineInputAssemblyStateCreateInfo inputAssembly = {};
	VkPipelineDynamicStateCreateInfo dynamicState = {};
	VkPipelineRasterizationStateCreateInfo rasterState = {};
	VkPipelineViewportStateCreateInfo viewportState = {};
	VkPipelineMultisampleStateCreateInfo msaa = {};
	VkPipelineDepthStencilStateCreateInfo depthStencil = {};
	VkPipelineColorBlendStateCreateInfo blend = {};
	VkPipelineRenderingCreateInfo render = {};
	CUtlVector<VkPipelineColorBlendAttachmentState> attachments = {};
	VkDynamicState dynamicStates[] = {
		VK_DYNAMIC_STATE_VIEWPORT_WITH_COUNT,
		VK_DYNAMIC_STATE_SCISSOR_WITH_COUNT,
	};
	VkPipelineLayoutCreateInfo stPipelineLayout = {};
	CUtlVector<VkDescriptorSetLayoutBinding> bindings = {};

	// TODO: Filter by vulkan shaders at some points
	stages.Resize(m_shader.m_objects.GetSize());
	modules.Resize(m_shader.m_objects.GetSize());
	for ( int i = 0; i < m_shader.m_objects.GetSize(); i++ )
	{
		VulkanInputMetaData_t *pMetaData = (VulkanInputMetaData_t*)m_shader.GetLumpPtr(m_shader.m_objects[i].m_nMetadataLump);
		for ( int u = 0; u < pMetaData->nDescriptorsCount; u++ )
		{
			VulkanDescriptor_t stDescriptor = ((VulkanDescriptor_t*)m_shader.GetLumpPtr(pMetaData->pDescriptorSets))[u];
			bool bFound = false;
			for ( auto &b: bindings )
			{
				if (b.binding != stDescriptor.uBinding)
					continue;
				bFound = true;
				break;
			}
			if (bFound)
				break;
			VkDescriptorSetLayoutBinding bind = {};
			bind.binding = stDescriptor.uBinding;
			bind.descriptorCount = 1;
			bind.descriptorType = stDescriptor.eDescriptorType;
			bind.stageFlags = VK_SHADER_STAGE_ALL;
			bindings.AppendTail(bind);
			m_bindings.AppendTail(stDescriptor);
		}
		modules[i].sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO;
		modules[i].pCode = (uint32_t*)m_shader.GetLumpPtr(m_shader.m_objects[i].m_nDataLump);
		modules[i].codeSize = m_shader.GetLumpSize(m_shader.m_objects[i].m_nDataLump);
		stages[i].sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
		stages[i].pNext = &modules[i];
		stages[i].pName = "main";
		stages[i].stage = VulkanGetShaderStage(m_shader.m_objects[i].m_eStage);
	}

	if ( bindings.GetSize() >= 0 )
	{
		m_setLayouts.Resize(1);
		VkDescriptorSetLayoutCreateInfo stSetLayoutCreateInfo = {};
		stSetLayoutCreateInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
		stSetLayoutCreateInfo.pBindings = bindings.GetData();
		stSetLayoutCreateInfo.bindingCount = bindings.GetSize();
		stSetLayoutCreateInfo.flags = VK_DESCRIPTOR_SET_LAYOUT_CREATE_UPDATE_AFTER_BIND_POOL_BIT;
		vkCreateDescriptorSetLayout(m_hDevice, &stSetLayoutCreateInfo, NULL, m_setLayouts.GetData());
		stPipelineLayout.setLayoutCount = 1;
		stPipelineLayout.pSetLayouts = m_setLayouts.GetData();
	}

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

	vertexInput.sType = VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO;
	vertexInput.vertexBindingDescriptionCount = m_layouts.GetSize();
	vertexInput.pVertexBindingDescriptions = m_layouts.GetData();
	vertexInput.vertexAttributeDescriptionCount = m_attributes.GetSize();
	vertexInput.pVertexAttributeDescriptions = m_attributes.GetData();

	inputAssembly.sType = VK_STRUCTURE_TYPE_PIPELINE_INPUT_ASSEMBLY_STATE_CREATE_INFO;
	inputAssembly.primitiveRestartEnable = VK_FALSE;
	inputAssembly.topology = VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST;

	dynamicState.sType = VK_STRUCTURE_TYPE_PIPELINE_DYNAMIC_STATE_CREATE_INFO;
	dynamicState.dynamicStateCount = 2;
	dynamicState.pDynamicStates = dynamicStates;

	rasterState.sType = VK_STRUCTURE_TYPE_PIPELINE_RASTERIZATION_STATE_CREATE_INFO;
	rasterState.polygonMode = VK_POLYGON_MODE_FILL;
	rasterState.lineWidth = 1;
	rasterState.cullMode = VK_CULL_MODE_NONE;
	rasterState.frontFace = VK_FRONT_FACE_CLOCKWISE;

	viewportState.sType = VK_STRUCTURE_TYPE_PIPELINE_VIEWPORT_STATE_CREATE_INFO;

	msaa.sType = VK_STRUCTURE_TYPE_PIPELINE_MULTISAMPLE_STATE_CREATE_INFO;
	msaa.rasterizationSamples = VK_SAMPLE_COUNT_1_BIT;

	render.sType = VK_STRUCTURE_TYPE_PIPELINE_RENDERING_CREATE_INFO;
	render.colorAttachmentCount = m_eFormats.GetSize();
	render.pColorAttachmentFormats = m_eFormats.GetData();

	depthStencil.sType = VK_STRUCTURE_TYPE_PIPELINE_DEPTH_STENCIL_STATE_CREATE_INFO;
	if (m_eDepthFormat == VK_FORMAT_D32_SFLOAT)
	{
		depthStencil.depthTestEnable = VK_TRUE;
		depthStencil.depthWriteEnable = VK_TRUE;
		depthStencil.depthCompareOp = VK_COMPARE_OP_LESS;
		render.depthAttachmentFormat = m_eDepthFormat;
	}

	for ( auto e: m_eFormats )
	{
		VkPipelineColorBlendAttachmentState a = {};
		a.colorWriteMask = VK_COLOR_COMPONENT_R_BIT | VK_COLOR_COMPONENT_G_BIT | VK_COLOR_COMPONENT_B_BIT | VK_COLOR_COMPONENT_A_BIT;
		a.blendEnable = VK_FALSE;
		attachments.AppendTail(a);
	}

	blend.sType = VK_STRUCTURE_TYPE_PIPELINE_COLOR_BLEND_STATE_CREATE_INFO;
	blend.attachmentCount = attachments.GetSize();
	blend.pAttachments = attachments.GetData();

	createInfo.sType = VK_STRUCTURE_TYPE_GRAPHICS_PIPELINE_CREATE_INFO;
	createInfo.stageCount = stages.GetSize();
	createInfo.pStages = stages.GetData();
	createInfo.layout = m_hPipelineLayout;
	createInfo.pVertexInputState = &vertexInput;
	createInfo.pInputAssemblyState = &inputAssembly;
	createInfo.pDynamicState = &dynamicState;
	createInfo.pRasterizationState = &rasterState;
	createInfo.pViewportState = &viewportState;
	createInfo.pDepthStencilState = &depthStencil;
	createInfo.pColorBlendState = &blend;
	createInfo.pMultisampleState = &msaa;
	createInfo.pNext = &render;
	vkCreateGraphicsPipelines(m_hDevice, NULL, 1, &createInfo, NULL, &m_hPipeline);
}

uint32_t CVkShader::PSGetResourceByName( const char *szName )
{

}

uint32_t CVkShader::VSGetResourceByName( const char *szName )
{

}