#include "vulkan_state.h"
CVkMaterial::CVkMaterial( IShader *pShader )
{
	m_pVkShader = (CVkShader*)pShader;

	VkDescriptorPoolSize pools[1] =
	{
		{
			VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
			128,
		}
	};
	VkDescriptorPoolCreateInfo stPool = {};
	stPool.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;
	stPool.poolSizeCount = 1;
	stPool.pPoolSizes = pools;
	stPool.maxSets = 1;
	stPool.flags = VK_DESCRIPTOR_POOL_CREATE_UPDATE_AFTER_BIND_BIT_EXT;
	vkCreateDescriptorPool(m_pVkShader->m_hDevice, &stPool, NULL, &m_hPool);

	VkDescriptorSetAllocateInfo stInfo = {};
	stInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
	stInfo.descriptorSetCount = 1;
	stInfo.descriptorPool = m_hPool;
	stInfo.pSetLayouts = m_pVkShader->m_setLayouts.GetData();
	vkAllocateDescriptorSets(m_pVkShader->m_hDevice, &stInfo, &m_hSet);
	
	for ( auto b: m_pVkShader->m_bindings )
	{
		VkWriteDescriptorSet write = {};
		write.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
		write.dstSet = m_hSet;
		write.dstBinding = b.uBinding;
		write.dstArrayElement = 0;
		write.descriptorType = b.eDescriptorType;
		write.descriptorCount = 1;
		switch (b.eDescriptorType)
		{
		case VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER:
			write.pBufferInfo = new VkDescriptorBufferInfo; 
			break;
		default:
			break;
		};
		m_writes.AppendTail(write);
	}
}

CVkMaterial::~CVkMaterial() 
{

}

void CVkMaterial::Frame()
{
	vkUpdateDescriptorSets(m_pVkShader->m_hDevice, m_writes.GetSize(), m_writes.GetData(), 0, 0);
};

void CVkMaterial::VSSetShaderResource( uint32_t uRegister, IRenderingObject *pResource )
{
	SetShaderResource(uRegister, SHADER_STAGE_VERTEX, pResource);
}

void CVkMaterial::PSSetShaderResource( uint32_t uRegister, IRenderingObject *pResource )
{
	SetShaderResource(uRegister, SHADER_STAGE_PIXEL, pResource);
}

void CVkMaterial::VSSetConstantsBuffer( uint32_t uRegister, IBuffer *pConstants )
{
	SetShaderResource(uRegister, SHADER_STAGE_VERTEX, pConstants);
}

void CVkMaterial::PSSetConstantsBuffer( uint32_t uRegister, IBuffer *pConstants )
{
	SetShaderResource(uRegister, SHADER_STAGE_PIXEL, pConstants);
}

void CVkMaterial::SetShaderResource( uint32_t uRegister, uint32_t uSet, IRenderingObject *pObject)
{
	union {
		IRenderingObject *pVkObject;
		CVkBuffer *pBuffer;
	}; 
	pVkObject = pObject;
	if (dynamic_cast<IBuffer*>(pObject))
	{
		((VkDescriptorBufferInfo*)m_writes[uRegister].pBufferInfo)->buffer = pBuffer->m_buffer;
		((VkDescriptorBufferInfo*)m_writes[uRegister].pBufferInfo)->range = pBuffer->m_nSize;
	}
	Frame();
}