#include "../vulkan_state.h"
#include "../raster_libraries.h"
#include "tier1/utlbuffer.h"
#include "tier1/utlvector.h"
#include "tier2/ifilesystem.h"
#include "../shaderparser.h"


BEGIN_BUILD_PIPELINE_LIBRARY(VertexDescription)
	library.flags = VK_GRAPHICS_PIPELINE_LIBRARY_VERTEX_INPUT_INTERFACE_BIT_EXT;
	
	VkPipelineVertexInputStateCreateInfo vertexInput = {};
	VkPipelineInputAssemblyStateCreateInfo inputAssembly = {};

	vertexInput.sType = VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO;
	vertexInput.vertexBindingDescriptionCount = layouts.GetSize();
	vertexInput.pVertexBindingDescriptions = layouts.GetData();
	vertexInput.vertexAttributeDescriptionCount = attributes.GetSize();
	vertexInput.pVertexAttributeDescriptions = attributes.GetData();

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

	pipeline.pVertexInputState = &vertexInput;
	pipeline.pInputAssemblyState = &inputAssembly;
END_BUILD_PIPELINE_LIBRARY()

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

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

void CVkVertexDescriptionPipelineLibrary::SetTopology( ETopologyMode eTopology )
{
	m_eTopology = VulkanGetTopology(eTopology);
}
struct VulkanDescriptorInit_t
{
	CUtlVector<VkDescriptorSetLayoutBinding> m_bindings = {};
	VkDescriptorSetLayoutCreateInfo m_set = {};
};

BEGIN_BUILD_PIPELINE_LIBRARY(VertexTransform)
	library.flags = VK_GRAPHICS_PIPELINE_LIBRARY_PRE_RASTERIZATION_SHADERS_BIT_EXT;

	VkPipelineRasterizationStateCreateInfo rasterState = {};
	VkPipelineViewportStateCreateInfo viewportState = {};
	int i = 0;

	CUtlVector<VulkanDescriptor_t> vertexDescriptors = ShaderParser()->GetDescriptors(m_pShader, SHADER_STAGE_VERTEX);
	VulkanDescriptorInit_t inits[SHADER_STAGE_COUNT] = {};
	CUtlVector<CUtlBuffer<unsigned char>> spirvs = {};
	CUtlVector<VkPipelineShaderStageCreateInfo> stages = {};
	CUtlVector<VkShaderModuleCreateInfo> stageShaders = {};

	for ( auto desc: vertexDescriptors )
	{
		VkDescriptorSetLayoutBinding binding = {};
		binding.descriptorCount = 1;
		binding.descriptorType = desc.eDescriptorType;
		binding.binding = desc.uBinding;
		inits[desc.eDescriptorType].m_bindings.AppendTail(binding);
	}
	for ( i = 0; i < SHADER_STAGE_COUNT; i++)
	{
		inits[i].m_set.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
		inits[i].m_set.bindingCount = inits[i].m_bindings.GetSize();
		inits[i].m_set.pBindings = inits[i].m_bindings.GetData();
	
		vkCreateDescriptorSetLayout(g_vkDevice, &inits[i].m_set, NULL, &m_setLayouts[i] );
	}

	for ( i = 0; i < SHADER_STAGE_COUNT; i++)
	{
		if ( i == SHADER_STAGE_PIXEL )
			continue;

		CUtlBuffer<unsigned char> code = ShaderParser()->GetShaderCode(m_pShader, (EShaderStage)i);

		// We may fail loading the specific stage
		if (code.GetSize() == 0)
			continue;

		VkShaderModuleCreateInfo mod = {};
		mod.sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO;
		mod.codeSize = code.GetSize();


		VkPipelineShaderStageCreateInfo shader = {};
		shader.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
		shader.pName = "main";
		shader.stage = VulkanGetShaderStage((EShaderStage)i);
		
		spirvs.AppendTail(code);
		stageShaders.AppendTail(mod);
		stages.AppendTail(shader);
	}

	// Fix pointers
	for ( i = 0; i < stages.GetSize(); i++ )
	{
		stageShaders[i].pCode = (uint32_t*)spirvs[i].GetMemory();
		stages[i].pNext = &stageShaders[i];
	}

	VkPipelineLayoutCreateInfo stPipelineLayout = {};
	stPipelineLayout.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
	stPipelineLayout.setLayoutCount = SHADER_STAGE_COUNT;
	stPipelineLayout.pSetLayouts = m_setLayouts;
	stPipelineLayout.flags = VK_PIPELINE_LAYOUT_CREATE_INDEPENDENT_SETS_BIT_EXT ;
	vkCreatePipelineLayout(g_vkDevice, &stPipelineLayout, NULL, &m_layout);

	rasterState.sType = VK_STRUCTURE_TYPE_PIPELINE_RASTERIZATION_STATE_CREATE_INFO;
	rasterState.polygonMode = VK_POLYGON_MODE_FILL;
	rasterState.lineWidth = 1;

	viewportState.sType = VK_STRUCTURE_TYPE_PIPELINE_VIEWPORT_STATE_CREATE_INFO;

	VkDynamicState dynamicStates[] = {
		VK_DYNAMIC_STATE_VIEWPORT_WITH_COUNT,
		VK_DYNAMIC_STATE_SCISSOR_WITH_COUNT,
	};
	VkPipelineDynamicStateCreateInfo dynamicState = {};
	dynamicState.sType = VK_STRUCTURE_TYPE_PIPELINE_DYNAMIC_STATE_CREATE_INFO;
	dynamicState.dynamicStateCount = 2;
	dynamicState.pDynamicStates = dynamicStates;
	pipeline.stageCount = stages.GetSize();
	pipeline.pStages = stages.GetData();
	pipeline.pDynamicState = &dynamicState;
	pipeline.pRasterizationState = &rasterState;
	pipeline.pViewportState = &viewportState;
	pipeline.layout = m_layout;
END_BUILD_PIPELINE_LIBRARY()

void CVkVertexTransformPipelineLibrary::SetShader( CCompiledShader *pShader )
{
	m_pShader = pShader;
}

BEGIN_BUILD_PIPELINE_LIBRARY(PixelShader)
	printf("--- PixelShader ---\n");
	
	VkPipelineDepthStencilStateCreateInfo depthStencil = {};
	int i = 0;
	CUtlVector<VulkanDescriptor_t> vertexDescriptors = ShaderParser()->GetDescriptors(m_pShader, SHADER_STAGE_PIXEL);
	VulkanDescriptorInit_t inits[SHADER_STAGE_COUNT] = {};
	for ( auto desc: vertexDescriptors )
	{
		VkDescriptorSetLayoutBinding binding = {};
		binding.descriptorCount = 1;
		binding.descriptorType = desc.eDescriptorType;
		binding.binding = desc.uBinding;
		inits[desc.eDescriptorType].m_bindings.AppendTail(binding);
	}
	for ( i = 0; i < SHADER_STAGE_COUNT; i++)
	{
		inits[i].m_set.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
		inits[i].m_set.bindingCount = inits[i].m_bindings.GetSize();
		inits[i].m_set.pBindings = inits[i].m_bindings.GetData();
	
		vkCreateDescriptorSetLayout(g_vkDevice, &inits[i].m_set, NULL, &m_setLayouts[i] );
	}
	
	VkPipelineLayoutCreateInfo stPipelineLayout = {};
	stPipelineLayout.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
	stPipelineLayout.setLayoutCount = SHADER_STAGE_COUNT;
	stPipelineLayout.pSetLayouts = m_setLayouts;
	stPipelineLayout.flags = VK_PIPELINE_LAYOUT_CREATE_INDEPENDENT_SETS_BIT_EXT ;
	vkCreatePipelineLayout(g_vkDevice, &stPipelineLayout, NULL, &m_layout);

	CUtlBuffer<unsigned char> spirv = {};
	VkPipelineShaderStageCreateInfo shader = {};
	VkShaderModuleCreateInfo mod = {};
	for ( auto desc: vertexDescriptors )
	{
		VkDescriptorSetLayoutBinding binding = {};
		binding.descriptorCount = 1;
		binding.descriptorType = desc.eDescriptorType;
		binding.binding = desc.uBinding;
		inits[desc.eDescriptorType].m_bindings.AppendTail(binding);
	}
	for ( i = 0; i < SHADER_STAGE_COUNT; i++)
	{
		inits[i].m_set.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
		inits[i].m_set.bindingCount = inits[i].m_bindings.GetSize();
		inits[i].m_set.pBindings = inits[i].m_bindings.GetData();
	
		vkCreateDescriptorSetLayout(g_vkDevice, &inits[i].m_set, NULL, &m_setLayouts[i] );
	}

	spirv = ShaderParser()->GetShaderCode(m_pShader, SHADER_STAGE_PIXEL);

	if (spirv.GetSize() == 0)
		goto skipshader;

	mod.sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO;
	mod.codeSize = spirv.GetSize();
	mod.pCode = (uint32_t*)spirv.GetMemory();

	shader.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
	shader.pName = "main";
	shader.stage = VulkanGetShaderStage(SHADER_STAGE_PIXEL);
	shader.pNext = &mod;

	pipeline.stageCount = 1;
	pipeline.pStages = &shader;

	depthStencil.sType = VK_STRUCTURE_TYPE_PIPELINE_DEPTH_STENCIL_STATE_CREATE_INFO;
	depthStencil.minDepthBounds = 0;
	depthStencil.maxDepthBounds = 1;
	depthStencil.depthTestEnable = VK_TRUE;
	depthStencil.depthWriteEnable = VK_TRUE;
	depthStencil.depthBoundsTestEnable = VK_FALSE;
	depthStencil.stencilTestEnable = VK_FALSE;
	depthStencil.depthCompareOp = VK_COMPARE_OP_LESS;
	pipeline.pDepthStencilState = &depthStencil;

skipshader:

	library.flags = VK_GRAPHICS_PIPELINE_LIBRARY_FRAGMENT_SHADER_BIT_EXT;
	pipeline.layout = m_layout;
END_BUILD_PIPELINE_LIBRARY()

void CVkPixelShaderPipelineLibrary::SetShader( CCompiledShader *pShader )
{
	m_pShader = pShader;
}

BEGIN_BUILD_PIPELINE_LIBRARY(PixelOutput)
	printf("--- PixelOutput ---\n");
	VkPipelineMultisampleStateCreateInfo msaa = {};
	VkPipelineRenderingCreateInfo render = {};
	VkPipelineDepthStencilStateCreateInfo depthStencil = {};
	VkPipelineColorBlendStateCreateInfo blend = {};
	CUtlVector<VkPipelineColorBlendAttachmentState> attachments = {};

	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();
	render.depthAttachmentFormat = VK_FORMAT_D32_SFLOAT;;
	
	depthStencil.sType = VK_STRUCTURE_TYPE_PIPELINE_DEPTH_STENCIL_STATE_CREATE_INFO;
	depthStencil.minDepthBounds = 0;
	depthStencil.maxDepthBounds = 1;
	depthStencil.depthTestEnable = VK_TRUE;
	depthStencil.depthWriteEnable = VK_TRUE;
	depthStencil.depthBoundsTestEnable = VK_FALSE;
	depthStencil.stencilTestEnable = VK_FALSE;
	depthStencil.depthCompareOp = VK_COMPARE_OP_LESS;

	for ( auto e: m_eFormats )
	{
		VkPipelineColorBlendAttachmentState a = {};
		a.blendEnable = VK_TRUE;
		a.blendEnable = VK_TRUE;
		a.srcColorBlendFactor = VK_BLEND_FACTOR_SRC_ALPHA;
		a.dstColorBlendFactor = VK_BLEND_FACTOR_ONE_MINUS_SRC_ALPHA;
		a.colorBlendOp = VK_BLEND_OP_ADD;
		a.srcAlphaBlendFactor = VK_BLEND_FACTOR_ONE;
		a.dstAlphaBlendFactor = VK_BLEND_FACTOR_ZERO;
		a.alphaBlendOp = VK_BLEND_OP_ADD;
		attachments.AppendTail(a);
	}
	blend.sType = VK_STRUCTURE_TYPE_PIPELINE_COLOR_BLEND_STATE_CREATE_INFO;
	blend.attachmentCount = m_eFormats.GetSize();
	blend.pAttachments = attachments.GetData();

	pipeline.pDepthStencilState = &depthStencil;
	pipeline.pColorBlendState = &blend;
	pipeline.pMultisampleState = &msaa;
	render.pNext = pipeline.pNext;
	pipeline.pNext = &render;
	library.flags = VK_GRAPHICS_PIPELINE_LIBRARY_FRAGMENT_OUTPUT_INTERFACE_BIT_EXT;

END_BUILD_PIPELINE_LIBRARY()
	
void CVkPixelOutputPipelineLibrary::AddAttachment( EImageFormat eFormat )
{
	m_eFormats.AppendTail(CVkImage::GetImageFormat(eFormat));	
}