#include "materialsystem/imaterialsystem.h"
#include "tier0/platform.h"
#include "tier1/interface.h"
#include "tier1/utlvector.h"
#include "vulkan_state.h"

CUtlVector<VkCommandBuffer> g_vkCommandBuffers;
IVkCommandBufferManager *g_pCommandBufferManager;

struct VulkanCommandLastUsage_t
{
	uint32_t m_nLastUsedCommand;
	EDependencyMode m_eLastUsage;
	VulkanCommandDepenency_t m_dependency;
};

struct VulkanSwapchainCommandLastUsage_t
{
	uint32_t m_nLastUsedCommand;
	EDependencyMode m_eLastUsage;
	VulkanCommandSwapchainDepenency_t m_dependency;
};

class CVkCommandBuffer: public IVkCommandBuffer
{
public:
	virtual void AddCommand( CVkCommand *pCommand ) override;
	virtual void Reset() override;
	virtual void Submit( int iFrameIndex ) override;
	virtual void Render() override;
private:
	void SortDependencies();
	void TryBarrier( int iCurrent, int iCurrentBuffer );
	VulkanCommandDepenency_t *FindLastObjectDependency( IRenderingObject *pObject );

	CUtlVector<CVkCommand*> m_commands;
	CUtlVector<VkCommandBuffer> m_hBuffers = NULL;
	CUtlVector<VulkanCommandLastUsage_t> m_usedDependencies; 
	CUtlVector<CUtlVector<VulkanCommandLastUsage_t>> m_dependencies; 
	CUtlVector<VulkanSwapchainCommandLastUsage_t> m_usedSwapchainDependencies; 
	CUtlVector<CUtlVector<VulkanSwapchainCommandLastUsage_t>> m_swapchainDependencies; 
};

void CVkCommandBuffer::Reset()
{
	VkCommandBufferAllocateInfo commandBufferAllocInfo = {};
	int i = 0;

	m_hBuffers = {};
	m_hBuffers.Resize(g_vkCommandPools.GetSize());

	for ( auto pool: g_vkCommandPools)
	{
		commandBufferAllocInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
		commandBufferAllocInfo.commandPool = pool;
		commandBufferAllocInfo.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
		commandBufferAllocInfo.commandBufferCount = 1;
		vkAllocateCommandBuffers(g_vkDevice, &commandBufferAllocInfo, &m_hBuffers[i]);
		i++;
	}

	m_commands = {};
}

void CVkCommandBuffer::AddCommand( CVkCommand *pCommand )
{
	m_commands.AppendTail(pCommand);
}

void CVkCommandBuffer::Submit( int iFrameIndex )
{
	g_vkCommandBuffers.AppendTail(m_hBuffers[iFrameIndex]);
};

struct VulkanBarrierObjects_t
{
	VulkanCommandDepenency_t m_dependency;
};

void CVkCommandBuffer::Render()
{

	uint32_t nNumDependencies;
	int i = 0;
	int y = 0;

	SortDependencies();
	
	for (auto hBuffer: m_hBuffers)
	{
		y = 0;
		VkCommandBufferBeginInfo stCommandBufferBeginInfo = {};
		stCommandBufferBeginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
		vkBeginCommandBuffer(hBuffer, &stCommandBufferBeginInfo);

		for (auto &pCommand: m_commands)
		{
			TryBarrier(y,i);
			pCommand->Execute(hBuffer, i);
			y++;
		};
		vkEndCommandBuffer(hBuffer);
		
		i++;
	}
}

struct VulkanCommandBarrier_t
{
	EDependencyMode m_eOldDependency;
	EDependencyMode m_eNewDependency;
	IRenderingObject *m_pObject;
};

void CVkCommandBuffer::SortDependencies()
{
	m_dependencies = {};
	m_swapchainDependencies = {};
	m_dependencies.Resize(m_commands.GetSize());
	m_swapchainDependencies.Resize(m_commands.GetSize());
	m_usedDependencies = {};
	m_usedSwapchainDependencies = {};
	
	int i = 0;
	for ( auto dependency: m_dependencies)
	{
		dependency = NULL;
	}

	for (auto &pCommand: m_commands)
	{
		for ( auto dependency: pCommand->m_dependencies)
		{
			VulkanCommandLastUsage_t *pLastUsage = NULL;

			for ( auto &dep: m_usedDependencies )
			{
				if (dep.m_dependency.m_pObject == dependency.m_pObject)
					pLastUsage = &dep;
			}

			VulkanCommandLastUsage_t stLastUsage = {};
			stLastUsage.m_nLastUsedCommand = i;
			stLastUsage.m_dependency = dependency;
			stLastUsage.m_eLastUsage = DEPENDENCY_MODE_FROM_PREVIOUS;

			if (pLastUsage)
				stLastUsage.m_eLastUsage = pLastUsage->m_dependency.m_eDependencyMode;

			if (pLastUsage)
				m_dependencies[pLastUsage->m_nLastUsedCommand+1].AppendTail(stLastUsage);
			else
				m_dependencies[0].AppendTail(stLastUsage);

			m_usedDependencies.AppendTail(stLastUsage);
		}
		for ( auto dependency: pCommand->m_swapchainDependencies)
		{
			VulkanSwapchainCommandLastUsage_t *pLastUsage = NULL;

			for ( auto &dep: m_usedSwapchainDependencies )
			{
				if (dep.m_dependency.m_ppObjects == dependency.m_ppObjects)
					pLastUsage = &dep;
			}

			VulkanSwapchainCommandLastUsage_t stLastUsage = {};
			stLastUsage.m_nLastUsedCommand = i;
			stLastUsage.m_dependency = dependency;
			stLastUsage.m_eLastUsage = DEPENDENCY_MODE_FROM_PREVIOUS;

			if (pLastUsage)
				stLastUsage.m_eLastUsage = pLastUsage->m_dependency.m_eDependencyMode;

			if (pLastUsage)
				m_swapchainDependencies[pLastUsage->m_nLastUsedCommand+1].AppendTail(stLastUsage);
			else
				m_swapchainDependencies[0].AppendTail(stLastUsage);

			m_usedSwapchainDependencies.AppendTail(stLastUsage);
		}
		i++;
	}
}


void CVkCommandBuffer::TryBarrier( int iCurrent, int iCurrentBuffer )
{
	int x = 0;
	CUtlVector<VulkanCommandLastUsage_t> dependency = m_dependencies[iCurrent];
	CUtlVector<VulkanSwapchainCommandLastUsage_t> swapchainDependency = m_swapchainDependencies[iCurrent];
	CUtlVector<VulkanCommandBarrier_t> barriers;

	VkDependencyInfo stDependencyInfo = {};
	stDependencyInfo.sType = VK_STRUCTURE_TYPE_DEPENDENCY_INFO;
	CUtlVector<VkMemoryBarrier2> memoryBarriers = {};
	CUtlVector<VkImageMemoryBarrier2> imageBarriers = {};
	CUtlVector<VkBufferMemoryBarrier2> bufferBarriers = {};

	for (auto barrier: dependency)
	{
		VulkanCommandBarrier_t stBarrier = {};
		stBarrier.m_eOldDependency = barrier.m_eLastUsage;
		stBarrier.m_eNewDependency = barrier.m_dependency.m_eDependencyMode;
		stBarrier.m_pObject = barrier.m_dependency.m_pObject;
		barriers.AppendTail(stBarrier);
		x++;
	};
	x = 0;
	for (auto barrier: swapchainDependency)
	{
		VulkanCommandBarrier_t stBarrier = {};
		stBarrier.m_eOldDependency = barrier.m_eLastUsage;
		stBarrier.m_eNewDependency = barrier.m_dependency.m_eDependencyMode;
		stBarrier.m_pObject = barrier.m_dependency.m_ppObjects[iCurrentBuffer];
		barriers.AppendTail(stBarrier);
		x++;
	};
	for ( auto barrier: barriers )
	{
		union {
			IRenderingObject *pObject;
			CVkImage *pImage;
			CVkBuffer *pBuffer;
		};

		pObject = barrier.m_pObject;

		if (dynamic_cast<CVkImage*>(pObject))
		{
			
			VkImageMemoryBarrier2 imageMemoryBarrier = {};
			imageMemoryBarrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER_2;
			imageMemoryBarrier.image = pImage->m_image;
			imageMemoryBarrier.subresourceRange = pImage->m_range;
			imageMemoryBarrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
			imageMemoryBarrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
			imageMemoryBarrier.newLayout = VulkanGetImageLayout(barrier.m_eNewDependency);
			imageMemoryBarrier.dstStageMask = VulkanGetStageFlags(barrier.m_eNewDependency);
			imageMemoryBarrier.dstAccessMask = VulkanGetAccessFlags(barrier.m_eNewDependency);
			if (barrier.m_eOldDependency == DEPENDENCY_MODE_FROM_PREVIOUS)
			{	
				imageMemoryBarrier.oldLayout = VulkanGetImageLayout(DEPENDENCY_MODE_ALL_COMMANDS);
				imageMemoryBarrier.srcStageMask = VulkanGetStageFlags(DEPENDENCY_MODE_ALL_COMMANDS);
				imageMemoryBarrier.srcAccessMask = VulkanGetAccessFlags(DEPENDENCY_MODE_ALL_COMMANDS);
			} else {
				imageMemoryBarrier.oldLayout = VulkanGetImageLayout(barrier.m_eOldDependency);
				imageMemoryBarrier.srcStageMask = VulkanGetStageFlags(barrier.m_eOldDependency);
				imageMemoryBarrier.srcAccessMask = VulkanGetAccessFlags(barrier.m_eOldDependency);
			}
			imageBarriers.AppendTail(imageMemoryBarrier);
			
			pImage->m_eLastUsage = barrier.m_eNewDependency;
		}
		if (dynamic_cast<CVkBuffer*>(pObject))
		{
		}
	}

	if (imageBarriers.GetSize() == 0 && bufferBarriers.GetSize() == 0 )
		return;

	stDependencyInfo.imageMemoryBarrierCount = imageBarriers.GetSize();
	stDependencyInfo.bufferMemoryBarrierCount = bufferBarriers.GetSize();
	stDependencyInfo.pImageMemoryBarriers = imageBarriers.GetData();
	stDependencyInfo.pBufferMemoryBarriers = bufferBarriers.GetData();
	vkCmdPipelineBarrier2(m_hBuffers[iCurrentBuffer], &stDependencyInfo);
};

class CVkCommandBufferManager: public IVkCommandBufferManager
{
public:
	virtual void Init() override;
	virtual void Shutdown() override;
	virtual IVkCommandBuffer *CreateCommandBuffer() override;
	virtual CVkCommand *CreateCommand( const char *szName ) override;
};

EXPOSE_INTERFACE(CVkCommandBufferManager, IVkCommandBufferManager, VULKAN_COMMAND_BUFFER_MANAGER_INTERFACE_NAME);

void CVkCommandBufferManager::Init()
{

}

void CVkCommandBufferManager::Shutdown()
{

}

IVkCommandBuffer *CVkCommandBufferManager::CreateCommandBuffer()
{
	return new CVkCommandBuffer;
}


struct VulkanCommandRegistry_t
{
	const char *m_szName;
	fnCreateVulkanCommand_t m_pfnCreate;
	struct VulkanCommandRegistry_t *m_pNext;
} *s_pVulkanCommands = NULL;

CVkCommand *CVkCommandBufferManager::CreateCommand( const char *szName )
{
	VulkanCommandRegistry_t *pCommand;
	for ( pCommand = s_pVulkanCommands; pCommand; pCommand = pCommand->m_pNext )
	{
		if (!V_strcmp(pCommand->m_szName, szName))
			return pCommand->m_pfnCreate();
	}
	Plat_FatalErrorFunc("Command %s not found\n", szName);
	return NULL;
}

CVkCommandRegistry::CVkCommandRegistry( const char *szName, fnCreateVulkanCommand_t pfnCreate )
{
	VulkanCommandRegistry_t *pCommand = new VulkanCommandRegistry_t;
	pCommand->m_szName = szName;
	pCommand->m_pfnCreate = pfnCreate;
	pCommand->m_pNext = s_pVulkanCommands;
	s_pVulkanCommands = pCommand;
}


void CVkCommand::AddDependency( IRenderingObject *pObject, EDependencyMode eDependencyMode )
{
	m_dependencies.AppendTail({ pObject, eDependencyMode });
}
void CVkCommand::AddSwapchainDependency( IRenderingObject **ppObject, EDependencyMode eDependencyMode )
{
	m_swapchainDependencies.AppendTail({ ppObject, eDependencyMode });
}