#include "vulkan_state.h"
#include "commands.h"

CVkRenderCommandList::~CVkRenderCommandList()
{
	ResetRendering();
}


void CVkRenderCommandList::ResetRendering()
{
	for ( auto m: m_pCommandBuffers )
	{
		m_pCommandBufferManager->FreeCommandBufferWithCommands(m);
	}
	m_pPostRaster = NULL;
	m_materials = {};
	m_pCommandBuffers = {};
}

void CVkRenderCommandList::SetRenderTarget( uint32_t uIndex, IImage *pImage )
{
	SwitchRenderingStage(RENDERING_STAGE_SETUP_RASTER);
	VulkanRenderOutput_t *pOutput = FindOrCreateRenderOutput(uIndex);
	pOutput->m_stImage.m_eObjectType = FRAME_OBJECT_TYPE_SINGLE;
	pOutput->m_stImage.m_pSingle = pImage;
}

void CVkRenderCommandList::SetClearColor( uint32_t uIndex, float r, float g, float b, float a )
{
	SwitchRenderingStage(RENDERING_STAGE_SETUP_RASTER);
	VulkanRenderOutput_t *pOutput = FindOrCreateRenderOutput(uIndex);
	pOutput->m_fClearColor[0] = r;
	pOutput->m_fClearColor[1] = g;
	pOutput->m_fClearColor[2] = b;
	pOutput->m_fClearColor[3] = a;
}

void CVkRenderCommandList::SetLoadStoreModes( uint32_t uIndex, ELoadMode eLoadMode, EStoreMode eStoreMode )
{
	SwitchRenderingStage(RENDERING_STAGE_SETUP_RASTER);
	VulkanRenderOutput_t *pOutput = FindOrCreateRenderOutput(uIndex);
	pOutput->m_eLoadMode = eLoadMode;
	pOutput->m_eStoreMode = eStoreMode;

}

void CVkRenderCommandList::SetDepthTarget( IImage *pDepth )
{
	SwitchRenderingStage(RENDERING_STAGE_SETUP_RASTER);
	if (pDepth)
	{
		m_bDepthEnabled = true;
		m_depth.m_stImage.m_pSingle = pDepth;
	}
	else
	{
		m_bDepthEnabled = false;
	}

}
void CVkRenderCommandList::SetClearDepth( float fVal)
{
	SwitchRenderingStage(RENDERING_STAGE_SETUP_RASTER);
	m_depth.m_fClearDepth = fVal;
}

void CVkRenderCommandList::SetRenderResolution( uint32_t iWidth, uint32_t iHeight )
{
	SwitchRenderingStage(RENDERING_STAGE_SETUP_RASTER);
	m_uWidth = iWidth;
	m_uHeight = iHeight;
}


void CVkRenderCommandList::SetScissors( uint32_t uX, uint32_t uY, uint32_t uWidth, uint32_t uHeight )
{
	SwitchRenderingStage(RENDERING_STAGE_RASTER);
}

void CVkRenderCommandList::SetViewport( uint32_t uX, uint32_t uY, uint32_t uWidth, uint32_t uHeight, float fMinDepth, float fMaxDepth )
{
	SwitchRenderingStage(RENDERING_STAGE_RASTER);
}





void CVkRenderCommandList::SetMaterial( IMaterial *pMaterial )
{
	CVkMaterial *pVkMat = (CVkMaterial*)pMaterial;
	CVkComputeShader *pCS = dynamic_cast<CVkComputeShader*>(pVkMat->m_pShader);

	if (pCS)
	{
		SwitchRenderingStage(RENDERING_STAGE_POST_RASTER);

		CVkSetShaderCommand *pSetShader = CREATE_COMMAND(m_pCommandBufferManager, SetShader);
		pSetShader->pShader = pCS;
		m_pPostRaster->AddCommand(pSetShader);

		CVkSetShaderDataCommand *pSetShaderData = CREATE_COMMAND(m_pCommandBufferManager, SetShaderData);
		pSetShaderData->pShaderData = pMaterial;
		m_pPostRaster->AddCommand(pSetShaderData);
	}
	else
	{
		SwitchRenderingStage(RENDERING_STAGE_SETUP_RASTER);
		bool bWasCreated = false;
		m_pCurrentMaterialBuffer = FindOrCreateMaterialCommandBuffer(pMaterial, &bWasCreated);
		if (bWasCreated)
		{
			m_pCurrentMaterialBuffer->Reset();
			CVkBeginCommand *pBeginCommand = CREATE_COMMAND(m_pCommandBufferManager, Begin);
			pBeginCommand->images = m_pOutput;
			pBeginCommand->nResolutionX = m_uWidth;
			pBeginCommand->nResolutionY = m_uHeight;
			for ( auto &i: pBeginCommand->images)
			{
				pBeginCommand->AddDependency(i.m_stImage.m_pSingle, DEPENDENCY_MODE_DRAWCALL_MIXED_IMAGE);
			}
			if ( m_bDepthEnabled )
			{
				pBeginCommand->AddDependency(m_depth.m_stImage.m_pSingle, DEPENDENCY_MODE_DRAWCALL_OUTPUT_DEPTH_IMAGE);
				pBeginCommand->bDepthEnabled = m_bDepthEnabled;
				pBeginCommand->stDepthImage = m_depth;

			}
			m_pCurrentMaterialBuffer->AddCommand(pBeginCommand);

			CVkSetShaderCommand *pSetShader = CREATE_COMMAND(m_pCommandBufferManager, SetShader);
			pSetShader->pShader = ((CVkMaterial*)pMaterial)->m_pShader;
			m_pCurrentMaterialBuffer->AddCommand(pSetShader);

			CVkSetShaderDataCommand *pSetShaderData = CREATE_COMMAND(m_pCommandBufferManager, SetShaderData);
			pSetShaderData->pShaderData = pMaterial;
			m_pCurrentMaterialBuffer->AddCommand(pSetShaderData);

			CVkSetScissorsCommand *pScissorsCommand = CREATE_COMMAND(m_pCommandBufferManager, SetScissors);
			pScissorsCommand->uWidth = m_uWidth;
			pScissorsCommand->uHeight = m_uHeight;
			m_pCurrentMaterialBuffer->AddCommand(pScissorsCommand);

			CVkSetViewportCommand *pViewportCommand = CREATE_COMMAND(m_pCommandBufferManager, SetViewport);
			pViewportCommand->fX = 0;
			pViewportCommand->fY = 0;
			pViewportCommand->fWidth = m_uWidth;
			pViewportCommand->fHeight = m_uHeight;
			pViewportCommand->fDepthMin = 0;
			pViewportCommand->fDepthMax = 1;
			m_pCurrentMaterialBuffer->AddCommand(pViewportCommand);
		}

	}
}

void CVkRenderCommandList::SetVertexBuffer( uint32_t uBinding, IVertexBuffer *pBuffer )
{
	SwitchRenderingStage(RENDERING_STAGE_RASTER);
	CVkSetVertexBufferCommand *pCmd = CREATE_COMMAND(m_pCommandBufferManager, SetVertexBuffer);
	pCmd->uBinding = uBinding;
	pCmd->pBuffer = pBuffer;
	pCmd->AddDependency(pBuffer, DEPENDENCY_MODE_BUFFER_SOURCE);
	m_pCurrentMaterialBuffer->AddCommand(pCmd);
}

void CVkRenderCommandList::SetIndexBuffer( IVertexBuffer *pBuffer )
{
	SwitchRenderingStage(RENDERING_STAGE_RASTER);
}

void CVkRenderCommandList::DrawPrimitives( uint32_t nVertexCount, uint32_t nFirstVertex, uint32_t nInstanceCount, uint32_t nFirstInstance )
{
	SwitchRenderingStage(RENDERING_STAGE_RASTER);

	CVkDrawPrimitivesCommand *pCmd = CREATE_COMMAND(m_pCommandBufferManager, DrawPrimitives);
	pCmd->nVertexCount = nVertexCount;
	pCmd->nFirstVertex = nFirstVertex;
	pCmd->nInstanceCount = nInstanceCount;
	pCmd->nFirstInstance = nFirstInstance;
	FlushBarriers(pCmd);

	m_pCurrentMaterialBuffer->AddCommand(pCmd);
}

void CVkRenderCommandList::DrawPrimitivesIndexed( uint32_t nIndexCount, uint32_t nFirstIndex, uint32_t nVertexOffset, uint32_t nInstanceCount, uint32_t nFirstInstance )
{
	SwitchRenderingStage(RENDERING_STAGE_RASTER);

}


void CVkRenderCommandList::CopyImageToImage( IImage *pSrc, IImage *pDst )
{
	SwitchRenderingStage(RENDERING_STAGE_POST_RASTER);
	
}
void CVkRenderCommandList::BlitImageToImage( IImage *pSrc, ImageSector_t src, IImage *pDst, ImageSector_t dst )
{
	SwitchRenderingStage(RENDERING_STAGE_POST_RASTER);
	CVkBlitCommand *pCmd = CREATE_COMMAND(m_pCommandBufferManager, Blit);
	pCmd->stInputImage.m_eObjectType = FRAME_OBJECT_TYPE_SINGLE;
	pCmd->stInputImage.m_pSingle = pSrc;
	pCmd->stOutputImage.m_eObjectType = FRAME_OBJECT_TYPE_SINGLE;
	pCmd->stOutputImage.m_pSingle = pDst;
	pCmd->AddDependency(pSrc, DEPENDENCY_MODE_BLIT_IMAGE_SOURCE);
	pCmd->AddDependency(pDst, DEPENDENCY_MODE_BLIT_IMAGE_DESTINATION);
	pCmd->iSrcMin[0] = src.m_iX;
	pCmd->iSrcMin[1] = src.m_iY;
	pCmd->iSrcMin[2] = 0;
	pCmd->iDstMin[0] = dst.m_iX;
	pCmd->iDstMin[1] = dst.m_iY;
	pCmd->iDstMin[2] = 0;
	pCmd->iSrcMax[0] = src.m_iWidth;
	pCmd->iSrcMax[1] = src.m_iHeight;
	pCmd->iSrcMax[2] = 1;
	pCmd->iDstMax[0] = dst.m_iWidth;
	pCmd->iDstMax[1] = dst.m_iHeight;
	pCmd->iDstMax[2] = 1;
	m_pPostRaster->AddCommand(pCmd);

}

void CVkRenderCommandList::ClearImage( IImage *pImage, float fR, float fG, float fB, float fA )
{
	SwitchRenderingStage(RENDERING_STAGE_POST_RASTER);
	CVkClearColorCommand *pCmd = CREATE_COMMAND(m_pCommandBufferManager, ClearColor);
	pCmd->AddDependency(pImage, DEPENDENCY_MODE_COLOR_CLEAR_DESTINATION);
	pCmd->r = fR;
	pCmd->g = fG;
	pCmd->b = fB;
	pCmd->a = fA;
	pCmd->stImage.m_eObjectType = FRAME_OBJECT_TYPE_SINGLE;
	pCmd->stImage.m_pSingle = pImage;
	m_pPostRaster->AddCommand(pCmd);

}

void CVkRenderCommandList::ClearDepth( IImage *pImage, float fVal )
{

}

void CVkRenderCommandList::ResolveImage( IImage *pOriginal, IImage *pResolved )
{
	SwitchRenderingStage(RENDERING_STAGE_POST_RASTER);

	CVkResolveImageCommand *pCmd = CREATE_COMMAND(m_pCommandBufferManager, ResolveImage);
	pCmd->stInputImage.m_eObjectType = FRAME_OBJECT_TYPE_SINGLE;
	pCmd->stInputImage.m_pSingle = pOriginal;
	pCmd->stOutputImage.m_eObjectType = FRAME_OBJECT_TYPE_SINGLE;
	pCmd->stOutputImage.m_pSingle = pResolved;
	pCmd->AddDependency(pOriginal, DEPENDENCY_MODE_IMAGE_RESOLVE_SOURCE);
	pCmd->AddDependency(pResolved, DEPENDENCY_MODE_IMAGE_RESOLVE_DESTINATION);
	m_pPostRaster->AddCommand(pCmd);
}


void CVkRenderCommandList::StartRecording()
{
	ResetRendering();
}

void CVkRenderCommandList::EndRecording()
{
	SwitchRenderingStage(RENDERING_STAGE_FINISHED);
}

void CVkRenderCommandList::Barrier( IRenderingObject *pObject, bool bIsRead, bool bIsWrite )
{
	m_barriers.AppendTail({pObject, bIsRead, bIsWrite});
}
void CVkRenderCommandList::DispatchCompute( uint32_t uX, uint32_t uY, uint32_t uZ )
{	
	SwitchRenderingStage(RENDERING_STAGE_POST_RASTER);

	CVkDispatchCommand *pCmd = CREATE_COMMAND(m_pCommandBufferManager, Dispatch);
	pCmd->uX = uX;
	pCmd->uY = uY;
	pCmd->uZ = uZ;
	FlushBarriers(pCmd);
	m_pPostRaster->AddCommand(pCmd);
	
}

void CVkRenderCommandList::Submit()
{
	for ( auto m: m_pCommandBuffers)
	{
		m->Submit();
	}
}

void CVkRenderCommandList::FlushBarriers( CVkCommand *pCmd )
{	
	for ( auto &b: m_barriers )
	{
		if (dynamic_cast<CVkBuffer*>(b.pObject))
		{
			if (b.m_bIsRead)
				if (b.m_bIsWrite)
					pCmd->AddDependency(b.pObject, DEPENDENCY_MODE_SHADER_BUFFER_READ_WRITE);
				else
					pCmd->AddDependency(b.pObject, DEPENDENCY_MODE_SHADER_BUFFER_READ);
			else
				if (b.m_bIsWrite)
					pCmd->AddDependency(b.pObject, DEPENDENCY_MODE_SHADER_BUFFER_WRITE);

		}
		if (dynamic_cast<CVkImage*>(b.pObject))
		{
			if (b.m_bIsRead)
				if (b.m_bIsWrite)
					pCmd->AddDependency(b.pObject, DEPENDENCY_MODE_SHADER_IMAGE_READ_WRITE);
				else
					pCmd->AddDependency(b.pObject, DEPENDENCY_MODE_SHADER_IMAGE_READ);
			else
				if (b.m_bIsWrite)
					pCmd->AddDependency(b.pObject, DEPENDENCY_MODE_SHADER_IMAGE_WRITE);

		}
	}
	m_barriers = {};
}

void CVkRenderCommandList::SwitchRenderingStage( EVulkanRenderingStage eStage )
{
	if (eStage != RENDERING_STAGE_RASTER)
	{
		for ( auto m: m_materials)
		{
			CVkEndCommand *pEndCommand = CREATE_COMMAND(m_pCommandBufferManager, End);
			m.m_pCommandBuffer->AddCommand(pEndCommand);
			m.m_pCommandBuffer->Render();
			m_pCommandBuffers.AppendTail(m.m_pCommandBuffer);
		}
		m_materials = {};
	};
	if (eStage == RENDERING_STAGE_POST_RASTER )
	{
		if (m_eCurrentStage != RENDERING_STAGE_POST_RASTER)
		{
			m_pPostRaster = m_pCommandBufferManager->CreateCommandBuffer();
			m_pPostRaster->Reset();
		}
	}
	else 
	{
		if (m_pPostRaster)
		{
			m_pPostRaster->Render();
			m_pCommandBuffers.AppendTail(m_pPostRaster);
			m_pPostRaster = NULL;
		}
	}
	m_eCurrentStage = eStage;
}

VulkanRenderOutput_t *CVkRenderCommandList::FindOrCreateRenderOutput( uint32_t uIndex )
{
	for ( auto &v: m_pOutput )
	{
		if (v.m_uIndex == uIndex)
			return &v;
	}

	VulkanRenderOutput_t output = {};
	output.m_uIndex = uIndex;
	m_pOutput.AppendTail(output);
	return &m_pOutput[m_pOutput.GetSize()-1];
}

IVkCommandBuffer *CVkRenderCommandList::FindOrCreateMaterialCommandBuffer( IMaterial *pMaterial, bool *pbWasCreated )
{
	for (auto &m: m_materials)
	{
		if (m.m_pMaterial == pMaterial)
		{
			if (pbWasCreated)
				*pbWasCreated = false;
			return m.m_pCommandBuffer;
		}
	}

	VulkanMaterialCommandBuffer_t mat;
	mat.m_pCommandBuffer = m_pCommandBufferManager->CreateCommandBuffer();
	mat.m_pMaterial = pMaterial;
	m_materials.AppendTail(mat);
	if (pbWasCreated)
		*pbWasCreated = true;
	return m_materials[m_materials.GetSize()-1].m_pCommandBuffer;
}