#include "vulkan_state.h"
CVkMaterial::CVkMaterial( IBaseShader *pShader )
{
	VkDevice hDevice = NULL;

	CVkShader *pS = dynamic_cast<CVkShader*>(pShader);
	CVkComputeShader *pCS = dynamic_cast<CVkComputeShader*>(pShader);
	if ( pS )
	{
		if (pS->m_setLayouts.GetSize() == 0 )
			return;
		hDevice = pS->m_hDevice;
	}
	if (pCS)
	{
		hDevice = pCS->m_hDevice;
	}

	VkDescriptorPoolSize pools[5] =
	{
		{
			VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
			1024,
		},
		{
			VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
			1024,
		},
		{
			VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
			1024,
		},
		{
			VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE,
			1024,
		},
		{
			VK_DESCRIPTOR_TYPE_SAMPLER,
			1024,
		}
	};
	VkDescriptorPoolCreateInfo stPool = {};
	stPool.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;
	stPool.poolSizeCount = 5;
	stPool.pPoolSizes = pools;
	stPool.maxSets = 16;
	stPool.flags = VK_DESCRIPTOR_POOL_CREATE_UPDATE_AFTER_BIND_BIT_EXT;
	vkCreateDescriptorPool(hDevice, &stPool, NULL, &m_hPool);


	if ( pS )
	{
		VkDescriptorSetAllocateInfo stInfo = {};
		stInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
		stInfo.descriptorPool = m_hPool;
		stInfo.descriptorSetCount = pS->m_setLayouts.GetSize();
		stInfo.pSetLayouts = pS->m_setLayouts.GetData();
		m_hSets.Resize(pS->m_setLayouts.GetSize());
		VkResult r = vkAllocateDescriptorSets(hDevice, &stInfo, m_hSets.GetData());	
		VULKAN_RESULT_PRINT(r, vkAllocateDescriptorSets);
	}

	if ( pCS )
	{
		VkDescriptorSetAllocateInfo stInfo = {};
		stInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
		stInfo.descriptorPool = m_hPool;
		stInfo.descriptorSetCount = pCS->m_setLayouts.GetSize();
		stInfo.pSetLayouts = pCS->m_setLayouts.GetData();
		m_hSets.Resize(pCS->m_setLayouts.GetSize());
		VkResult r = vkAllocateDescriptorSets(hDevice, &stInfo, m_hSets.GetData());	
		VULKAN_RESULT_PRINT(r, vkAllocateDescriptorSets);
	}

	m_pShader = pShader;
}

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, 0, pResource);
}

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

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

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

void CVkMaterial::PSSetTextureArray( uint32_t uSet, ITextureArray *pArray )
{
	SetShaderResource(0, uSet, pArray);
}

void CVkMaterial::PSSetTexture( uint32_t uRegister, IImage *pImage )
{
	SetShaderResource(uRegister, 0, pImage);
}

void CVkMaterial::PSSetSampler( uint32_t uRegister, ISampler *pImage )
{
	SetShaderResource(uRegister, 0, pImage);
}

void CVkMaterial::CSSetConstantsBuffer( uint32_t uRegister, uint32_t uSet, IBuffer *pBuffer )
{
	SetShaderResource(uRegister, uSet, pBuffer);

}

void CVkMaterial::CSSetTexture( uint32_t uRegister, uint32_t uSet, IImage *pImage )
{
	SetShaderResource(uRegister, uSet, pImage);

}

void CVkMaterial::CSSetSampler( uint32_t uRegister, uint32_t uSet, ISampler *pSampler )
{
	SetShaderResource(uRegister, uSet, pSampler);

}

void CVkMaterial::CSSetTextureArray( uint32_t uSet, ITextureArray *pArray )
{
	SetShaderResource(0, uSet, pArray);

}



void CVkMaterial::SetShaderResource( uint32_t uRegister, uint32_t uSet, IRenderingObject *pObject)
{
	VkDevice hDevice;
	CVkShader *pShader = dynamic_cast<CVkShader*>(m_pShader);
	CVkComputeShader *pCS = dynamic_cast<CVkComputeShader*>(m_pShader);
	if ( pShader )
	{
		if (pShader->m_setLayouts.GetSize() == 0 )
			return;
		hDevice = pShader->m_hDevice;
	}
	if (pCS)
	{
		if (pCS->m_setLayouts.GetSize() == 0 )
			return;
		hDevice = pCS->m_hDevice;
	}
	union {
		IRenderingObject *pVkObject;
		CVkBuffer *pBuffer;
		CVkTextureArray *pArray;
		CVkImage *pImage;
		CVkSampler *pSampler;
	}; 
	pVkObject = pObject;
	if (dynamic_cast<IBuffer*>(pObject))
	{
		VkWriteDescriptorSet write = {};
		VkDescriptorBufferInfo stInfo = {};
		write.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
		write.dstSet = m_hSets[uSet];
		write.dstBinding = uRegister;
		write.dstArrayElement = 0;
		write.descriptorType = pBuffer->m_eDescriptorType;
		write.descriptorCount = 1;
		write.pBufferInfo = &stInfo;
		stInfo.buffer = pBuffer->m_buffer;	
		stInfo.range = pBuffer->m_nSize;
		vkUpdateDescriptorSets(hDevice, 1, &write, 0, 0);
	}
	if (dynamic_cast<IImage*>(pObject))
	{
		VkWriteDescriptorSet write = {};
		VkDescriptorImageInfo stInfo = {};
		write.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
		write.dstSet = m_hSets[uSet];
		write.dstBinding = uRegister;
		write.dstArrayElement = 0;
		write.descriptorType = pImage->m_eDescriptorType;
		write.descriptorCount = 1;
		write.pImageInfo = &stInfo;
		stInfo.imageView = pImage->m_imageView;	
		stInfo.imageLayout = VK_IMAGE_LAYOUT_GENERAL;	
		vkUpdateDescriptorSets(hDevice, 1, &write, 0, 0);

	}
	if (dynamic_cast<CVkSampler*>(pObject))
	{
		VkWriteDescriptorSet write = {};
		VkDescriptorImageInfo stInfo = {};
		write.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
		write.dstSet = m_hSets[uSet];
		write.dstBinding = uRegister;
		write.dstArrayElement = 0;
		write.descriptorType = VK_DESCRIPTOR_TYPE_SAMPLER;
		write.descriptorCount = 1;
		write.pImageInfo = &stInfo;
		stInfo.sampler = pSampler->m_sampler;	
		vkUpdateDescriptorSets(hDevice, 1, &write, 0, 0);
	}
	if (dynamic_cast<ITextureArray*>(pObject))
	{
		VkWriteDescriptorSet writes[2] = {};
		VkDescriptorImageInfo stWrites[128];
		VkDescriptorImageInfo stSampler = {};

		for ( int i = 0; i < 128; i++ )
		{
			if (pArray->m_pImages[i] == NULL)
			{
				stWrites[i] = {
					.imageView = pArray->m_pImages[0]->m_imageView,
					.imageLayout = VK_IMAGE_LAYOUT_GENERAL,
				};	
				continue;
			}
			stWrites[i] = {
				.imageView = pArray->m_pImages[i]->m_imageView,
				.imageLayout = VK_IMAGE_LAYOUT_GENERAL,
			};	
		}

		writes[0] = {};
		writes[0].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
		writes[0].dstSet = m_hSets[uSet];
		writes[0].dstBinding = 0;
		writes[0].dstArrayElement = 0;
		writes[0].descriptorType = VK_DESCRIPTOR_TYPE_SAMPLER;
		writes[0].descriptorCount = 1;
		writes[0].pImageInfo = &stSampler;
		stSampler.sampler = pArray->m_hSampler;

		writes[1] = {};
		writes[1].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
		writes[1].dstSet = m_hSets[uSet];
		writes[1].dstBinding = 1;
		writes[1].dstArrayElement = 0;
		writes[1].descriptorType = VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE;
		writes[1].descriptorCount = 128;
		writes[1].pImageInfo = stWrites;
		vkUpdateDescriptorSets(hDevice, 2, writes, 0, 0);
	}
}