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


DECLARE_VULKAN_COMMAND(ClearColor)
{

	CVkImage *pImg;

	VkClearColorValue color = {.float32 = {r,g,b,a}};
	VkImageSubresourceRange range = {
		.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT,
		.levelCount = 1,
		.layerCount = 1,
	};
	vkCmdClearColorImage(hCommandBuffer, ((CVkImage*)VulkanGetObject(stImage, iCurrentFrame))->m_image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, &color, 1, &range);
}

DECLARE_VULKAN_COMMAND(Begin)
{
	if (vkCmdBeginRendering || vkCmdBeginRenderingKHR)
	{
		V_printf("USING DYNAMIC RENDERING\n");
		V_printf("%p\n", vkCmdBeginRendering);
		V_printf("%p\n", vkCmdBeginRenderingKHR);
		VkRenderingInfo stRenderingInfo = {};
		CUtlVector<VkRenderingAttachmentInfo> attachments = {};
		VkRenderingAttachmentInfo depthAttachment = {};
		for (auto i: images)
		{
			VkRenderingAttachmentInfo a = {};
			a.sType = VK_STRUCTURE_TYPE_RENDERING_ATTACHMENT_INFO;
			a.imageLayout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;
			switch( i.m_eLoadMode )
			{
			case LOAD_MODE_DONT_CARE:
				a.loadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE;
				break;
			case LOAD_MODE_LOAD:
				a.loadOp = VK_ATTACHMENT_LOAD_OP_LOAD;
				break;
			case LOAD_MODE_CLEAR:
				a.loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR;
				a.clearValue.color.float32[0] = i.m_fClearColor[0];
				a.clearValue.color.float32[1] = i.m_fClearColor[1];
				a.clearValue.color.float32[2] = i.m_fClearColor[2];
				a.clearValue.color.float32[3] = i.m_fClearColor[3];
				break;
			}
			a.storeOp = VK_ATTACHMENT_STORE_OP_STORE;
			a.imageView = ((CVkImage*)VulkanGetObject(i.m_stImage, iCurrentFrame))->m_imageView;
			attachments.AppendTail(a);
		}
		if (bDepthEnabled)
		{
			depthAttachment.sType = VK_STRUCTURE_TYPE_RENDERING_ATTACHMENT_INFO;
			depthAttachment.imageLayout = VK_IMAGE_LAYOUT_DEPTH_ATTACHMENT_OPTIMAL;
			depthAttachment.loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR;
			depthAttachment.clearValue.depthStencil.depth = stDepthImage.m_fClearDepth;
			depthAttachment.storeOp = VK_ATTACHMENT_STORE_OP_STORE;
			depthAttachment.imageView = ((CVkImage*)VulkanGetObject(stDepthImage.m_stImage, iCurrentFrame))->m_imageView;
			stRenderingInfo.pDepthAttachment = &depthAttachment;
		}

		stRenderingInfo.sType = VK_STRUCTURE_TYPE_RENDERING_INFO;
		stRenderingInfo.layerCount = 1;
		stRenderingInfo.colorAttachmentCount = attachments.GetSize();
		stRenderingInfo.pColorAttachments = attachments.GetData();
		stRenderingInfo.renderArea = (VkRect2D){0,0,nResolutionX,nResolutionY};
		
		if (vkCmdBeginRendering)
			vkCmdBeginRendering(hCommandBuffer, &stRenderingInfo);
		else
			vkCmdBeginRenderingKHR(hCommandBuffer, &stRenderingInfo);
	}
	else
	{
		V_printf("USING RENDERPASS");
		CUtlVector<VkAttachmentDescription> attachments = {};
		CUtlVector<VkAttachmentReference> input_refs = {};
		CUtlVector<VkImageView> image_views = {};
		CUtlVector<VkClearValue> colors = {};
		attachments.Resize(images.GetSize());
		input_refs.Resize(images.GetSize());
		image_views.Resize(images.GetSize());
		colors.Resize(images.GetSize());
		uint32_t i = 0;

		for (auto f: images)
		{
			VkAttachmentDescription a = {};
			a.initialLayout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;
			a.finalLayout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;
			a.format = CVkImage::GetImageFormat(
					((CVkImage*)VulkanGetObject(f.m_stImage, iCurrentFrame))->m_eFormat);
			a.samples = CVkImage::GetMultisampling(
					((CVkImage*)VulkanGetObject(f.m_stImage, iCurrentFrame))->m_eMultisampleType);
			switch( f.m_eLoadMode )
			{
			case LOAD_MODE_DONT_CARE:
				a.loadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE;
				break;
			case LOAD_MODE_LOAD:
				a.loadOp = VK_ATTACHMENT_LOAD_OP_LOAD;
				break;
			case LOAD_MODE_CLEAR:
				a.loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR;
				colors[i].color.float32[0] = f.m_fClearColor[0];
				colors[i].color.float32[1] = f.m_fClearColor[1];
				colors[i].color.float32[2] = f.m_fClearColor[2];
				colors[i].color.float32[3] = f.m_fClearColor[3];
				break;
			}
			a.storeOp = VK_ATTACHMENT_STORE_OP_STORE;
			attachments[i] = a;

			VkAttachmentReference ref = {};
			ref.layout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;
			ref.attachment = i;
			input_refs[i] = ref;
			
			image_views[i] = ((CVkImage*)VulkanGetObject(f.m_stImage, iCurrentFrame))->m_imageView;
			i++;

		}
		if (bDepthEnabled)
		{
			CVkImage *depth = (CVkImage*)VulkanGetObject(stDepthImage.m_stImage, iCurrentFrame); 
			VkAttachmentDescription desc = {};
			desc.format = CVkImage::GetImageFormat(depth->m_eFormat);
			desc.samples = CVkImage::GetMultisampling(depth->m_eMultisampleType);
			desc.initialLayout = VK_IMAGE_LAYOUT_DEPTH_ATTACHMENT_OPTIMAL;
			desc.finalLayout = VK_IMAGE_LAYOUT_DEPTH_ATTACHMENT_OPTIMAL;
			desc.loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR;
			desc.storeOp = VK_ATTACHMENT_STORE_OP_STORE;
			attachments.AppendTail(desc);
			VkClearValue val = {};
			val.depthStencil.depth = stDepthImage.m_fClearDepth;
			colors.AppendTail(val);
			image_views.AppendTail(depth->m_imageView);
		}
		VkAttachmentReference depthRef;
		depthRef.attachment = attachments.GetSize() - 1;
		depthRef.layout = VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL;
		VkSubpassDescription subpass = {};
		subpass.pipelineBindPoint = VK_PIPELINE_BIND_POINT_GRAPHICS;
		subpass.colorAttachmentCount = input_refs.GetSize();
		subpass.pColorAttachments = input_refs.GetData();
		if (bDepthEnabled)
		{
			subpass.pDepthStencilAttachment = &depthRef;
		}
		VkRenderPassCreateInfo rpci = {};
		rpci.sType = VK_STRUCTURE_TYPE_RENDER_PASS_CREATE_INFO;
		rpci.attachmentCount = attachments.GetSize();
		rpci.pAttachments = attachments.GetData();
		rpci.subpassCount = 1;
		rpci.pSubpasses = &subpass;
		vkCreateRenderPass(hDevice, &rpci, NULL, &hRenderPass);

		VkFramebufferCreateInfo fbci = {};
		fbci.sType = VK_STRUCTURE_TYPE_FRAMEBUFFER_CREATE_INFO;
		fbci.renderPass = hRenderPass;
		fbci.attachmentCount = image_views.GetSize();
		fbci.pAttachments = image_views.GetData();
		fbci.width = nResolutionX;
		fbci.height = nResolutionY;
		fbci.layers = 1;
		vkCreateFramebuffer(hDevice, &fbci, NULL, &hFrameBuffer);
		VkRenderPassBeginInfo rpbi = {};
		rpbi.sType = VK_STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO;
		rpbi.renderPass = hRenderPass;
		rpbi.framebuffer = hFrameBuffer;
		rpbi.renderArea = (VkRect2D){0,0,nResolutionX,nResolutionY};
		rpbi.clearValueCount = colors.GetSize();
		rpbi.pClearValues = colors.GetData();
		vkCmdBeginRenderPass(hCommandBuffer, &rpbi, VK_SUBPASS_CONTENTS_INLINE);
	}
}

DECLARE_VULKAN_COMMAND(End)
{
	if (vkCmdEndRendering)
		vkCmdEndRendering(hCommandBuffer);
	else if (vkCmdEndRenderingKHR)
		vkCmdEndRenderingKHR(hCommandBuffer);
	else
	{
		vkCmdEndRenderPass(hCommandBuffer);

	}
}

DECLARE_VULKAN_COMMAND(SetShader)
{
	CVkShader *pS = dynamic_cast<CVkShader*>(pShader);
	CVkComputeShader *pCS = dynamic_cast<CVkComputeShader*>(pShader);
	if (pS)
		vkCmdBindPipeline(hCommandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, pS->m_hPipeline);
	if (pCS)
		vkCmdBindPipeline(hCommandBuffer, VK_PIPELINE_BIND_POINT_COMPUTE, pCS->m_hPipeline);
}
DECLARE_VULKAN_COMMAND(SetShaderData)
{
	CVkMaterial *pMat = (CVkMaterial*)pShaderData;
	CVkShader *pS = dynamic_cast<CVkShader*>(pMat->m_pShader);
	CVkComputeShader *pCS = dynamic_cast<CVkComputeShader*>(pMat->m_pShader);
	if (pS)
		vkCmdBindDescriptorSets(hCommandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, pS->m_hPipelineLayout, 0, pMat->m_hSets.GetSize(), pMat->m_hSets.GetData(), 0, 0 );
	if (pCS)
		vkCmdBindDescriptorSets(hCommandBuffer, VK_PIPELINE_BIND_POINT_COMPUTE, pCS->m_hPipelineLayout, 0, pMat->m_hSets.GetSize(), pMat->m_hSets.GetData(), 0, 0 );
}

DECLARE_VULKAN_COMMAND(SetVertexBuffer)
{
	CVkBuffer *pVkBuffer = (CVkBuffer*)pBuffer;
	VkDeviceSize uOffset = 0;
	vkCmdBindVertexBuffers(hCommandBuffer, uBinding, 1, &pVkBuffer->m_buffer, &uOffset);
}

DECLARE_VULKAN_COMMAND(DrawPrimitives)
{
	vkCmdDraw(hCommandBuffer, nVertexCount, nInstanceCount, nFirstVertex, nFirstInstance);
}

DECLARE_VULKAN_COMMAND(SetScissors)
{
	VkRect2D r = {
		.offset = {
			iX,
			iY,
		},
		.extent = {
			uWidth,
			uHeight,
		},
	};
	if (vkCmdSetScissorWithCount)
		vkCmdSetScissorWithCount(hCommandBuffer, 1, &r);
	else
		vkCmdSetScissor(hCommandBuffer, 0, 1, &r);
}

DECLARE_VULKAN_COMMAND(SetViewport)
{
	VkViewport v = {
		fX,
		fY+fHeight,
		fWidth,
		-fHeight,
		fDepthMin,
		fDepthMax
	};
	if (vkCmdSetScissorWithCount)
		vkCmdSetViewportWithCount(hCommandBuffer, 1, &v);
	else
		vkCmdSetViewport(hCommandBuffer, 0, 1, &v);
}
DECLARE_VULKAN_COMMAND(ResolveImage)
{
	CVkImage *pInput = (CVkImage*)VulkanGetObject(stInputImage, iCurrentFrame);
	CVkImage *pOutput = (CVkImage*)VulkanGetObject(stOutputImage, iCurrentFrame);
	VkImageResolve2 r = {
		.sType = VK_STRUCTURE_TYPE_IMAGE_RESOLVE_2,
		.srcSubresource = {
			.aspectMask = pInput->m_range.aspectMask,
			.layerCount = 1,
		},
		.srcOffset = {0,0,0},
		.dstSubresource = {
			.aspectMask = pOutput->m_range.aspectMask,
			.layerCount = 1,
		},
		.dstOffset = {0,0,0},
		.extent = {pInput->m_nWidth, pInput->m_nHeight, 1}
	};
	VkResolveImageInfo2 i = {
		.sType = VK_STRUCTURE_TYPE_RESOLVE_IMAGE_INFO_2,
		.srcImage = pInput->m_image,
		.srcImageLayout = VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL,
		.dstImage = pOutput->m_image,
		.dstImageLayout = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
		.regionCount = 1,
		.pRegions = &r,
	};
	vkCmdResolveImage2(hCommandBuffer, &i);
}

static VkPipelineStageFlags2 GetStage( EStageHazard stage )
{
	switch (stage)
	{
	case STAGE_HAZARD_COMPUTE:
		return VK_PIPELINE_STAGE_2_COMPUTE_SHADER_BIT;
	case STAGE_HAZARD_TRANSFER:
		return VK_PIPELINE_STAGE_2_TRANSFER_BIT;
	case STAGE_HAZARD_RASTER_COLOR_OUT:
		return VK_PIPELINE_STAGE_2_COLOR_ATTACHMENT_OUTPUT_BIT;
	case STAGE_HAZARD_NONE:
		return VK_PIPELINE_STAGE_2_NONE;
	case STAGE_HAZARD_RASTER_EARLY_DEPTH_OUT:
		return VK_PIPELINE_STAGE_2_EARLY_FRAGMENT_TESTS_BIT | VK_PIPELINE_STAGE_2_LATE_FRAGMENT_TESTS_BIT;
	case STAGE_HAZARD_RASTER_LATE_DEPTH_OUT:
		return VK_PIPELINE_STAGE_2_LATE_FRAGMENT_TESTS_BIT;
	}

}
static VkAccessFlags2 GetAccessSrc( uint32_t hazard, EStageHazard stage )
{
	VkAccessFlags2 f = 0;
	if (hazard & BARRIER_HAZARD_DRAW_ARGUMENTS)
		f |= VK_ACCESS_2_SHADER_WRITE_BIT;
	if (hazard & BARRIER_HAZARD_DESCRIPTOR)
		f |= VK_ACCESS_2_SHADER_WRITE_BIT;
	if (stage == STAGE_HAZARD_RASTER_EARLY_DEPTH_OUT)
	{
		if (hazard & BARRIER_HAZARD_DEPTH_STENCIL)
			f |= VK_ACCESS_2_DEPTH_STENCIL_ATTACHMENT_WRITE_BIT;
	}
	if (stage == STAGE_HAZARD_RASTER_COLOR_OUT)
	{
		if (hazard & BARRIER_HAZARD_RASTER_COLOR)
			f |= VK_ACCESS_2_COLOR_ATTACHMENT_READ_BIT;

	}
	if (hazard == 0)
		f |= VK_ACCESS_2_MEMORY_READ_BIT | VK_ACCESS_2_MEMORY_WRITE_BIT;
	if (stage == STAGE_HAZARD_COMPUTE)
		f |= VK_ACCESS_2_SHADER_STORAGE_READ_BIT | VK_ACCESS_2_SHADER_STORAGE_WRITE_BIT;
	return f;

}
static VkAccessFlags2 GetAccessDst( uint32_t hazard, EStageHazard stage )
{
	VkAccessFlags2 f = 0;
	if (hazard & BARRIER_HAZARD_DRAW_ARGUMENTS)
		f |= VK_ACCESS_2_INDIRECT_COMMAND_READ_BIT;
	if (hazard & BARRIER_HAZARD_DESCRIPTOR)
		f |= VK_ACCESS_2_SHADER_READ_BIT;
	if (hazard & BARRIER_HAZARD_DEPTH_STENCIL)
		f |= VK_ACCESS_2_DEPTH_STENCIL_ATTACHMENT_READ_BIT | VK_ACCESS_2_DEPTH_STENCIL_ATTACHMENT_WRITE_BIT;
	if (hazard & BARRIER_HAZARD_RASTER_COLOR)
		f |= VK_ACCESS_2_COLOR_ATTACHMENT_READ_BIT | VK_ACCESS_2_COLOR_ATTACHMENT_WRITE_BIT;
	if (hazard == 0)
		f |= VK_ACCESS_2_MEMORY_READ_BIT | VK_ACCESS_2_MEMORY_WRITE_BIT;
	if (stage == STAGE_HAZARD_COMPUTE)
		f |= VK_ACCESS_2_SHADER_STORAGE_READ_BIT | VK_ACCESS_2_SHADER_STORAGE_WRITE_BIT;
	return f;
}
static VkImageLayout GetLayout( EImageLayoutHazard hazard )
{
	switch (hazard)
	{
	case IMAGE_LAYOUT_COLOR:
		return VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;
	case IMAGE_LAYOUT_DEPTH_STENCIL:
		return VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL;
	case IMAGE_LAYOUT_GENERAL:
		return VK_IMAGE_LAYOUT_GENERAL;
	case IMAGE_LAYOUT_PRESENT:
		return VK_IMAGE_LAYOUT_PRESENT_SRC_KHR;
	case IMAGE_LAYOUT_UNDEFINED:
		return VK_IMAGE_LAYOUT_UNDEFINED;
	case IMAGE_LAYOUT_TRANSFER_SRC:
		return VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL;
	case IMAGE_LAYOUT_TRANSFER_DST:
		return VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL;
	}

}

DECLARE_VULKAN_COMMAND(Barrier)
{
	VkMemoryBarrier2 barrier = {};
	barrier.sType = VK_STRUCTURE_TYPE_MEMORY_BARRIER_2;
	barrier.srcStageMask = GetStage(eStageIn);
	barrier.srcAccessMask = GetAccessSrc(eHazard, eStageIn);
	barrier.dstStageMask = GetStage(eStageOut);
	barrier.dstAccessMask = GetAccessDst(eHazard, eStageOut);
	VkDependencyInfo dependency = {};
	dependency.sType = VK_STRUCTURE_TYPE_DEPENDENCY_INFO;
	dependency.memoryBarrierCount = 1;
	dependency.pMemoryBarriers = &barrier;
	if (vkCmdPipelineBarrier2)
		vkCmdPipelineBarrier2(hCommandBuffer, &dependency);
	else if (vkCmdPipelineBarrier2KHR)
		vkCmdPipelineBarrier2KHR(hCommandBuffer, &dependency);
	else
	{
		VkMemoryBarrier barrier = {};
		barrier.sType = VK_STRUCTURE_TYPE_MEMORY_BARRIER;
		barrier.srcAccessMask = GetAccessSrc(eHazard, eStageIn);
		barrier.dstAccessMask = GetAccessDst(eHazard, eStageOut);
		vkCmdPipelineBarrier(hCommandBuffer, 
				GetStage(eStageIn),
				GetStage(eStageOut),
				0, 1,&barrier, 0, NULL, 0, NULL);

	}

}

DECLARE_VULKAN_COMMAND(LayoutBarrier)
{
	VkImageMemoryBarrier2 barrier = {};
	barrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER_2;
	barrier.srcStageMask = GetStage(eStageIn);
	barrier.srcAccessMask = GetAccessSrc(eHazard, eStageIn);
	barrier.dstStageMask = GetStage(eStageOut);
	barrier.dstAccessMask = GetAccessDst(eHazard, eStageOut);
	barrier.image = ((CVkImage*)pImage)->m_image;
	barrier.subresourceRange = ((CVkImage*)pImage)->m_range;
	barrier.oldLayout = GetLayout(eLayoutIn);
	barrier.newLayout = GetLayout(eLayoutOut);
	VkDependencyInfo dependency = {};
	dependency.sType = VK_STRUCTURE_TYPE_DEPENDENCY_INFO;
	dependency.imageMemoryBarrierCount = 1;
	dependency.pImageMemoryBarriers = &barrier;
	if (vkCmdPipelineBarrier2)
		vkCmdPipelineBarrier2(hCommandBuffer, &dependency);
	else if (vkCmdPipelineBarrier2KHR)
		vkCmdPipelineBarrier2KHR(hCommandBuffer, &dependency);
	else
	{
		VkImageMemoryBarrier barrier = {};
		barrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
		barrier.srcAccessMask = GetAccessSrc(eHazard, eStageIn);
		barrier.dstAccessMask = GetAccessDst(eHazard, eStageOut);
		barrier.image = ((CVkImage*)pImage)->m_image;
		barrier.subresourceRange = ((CVkImage*)pImage)->m_range;
		barrier.oldLayout = GetLayout(eLayoutIn);
		barrier.newLayout = GetLayout(eLayoutOut);
		vkCmdPipelineBarrier(hCommandBuffer, 
				GetStage(eStageIn),
				GetStage(eStageOut),
				0, 0, NULL, 0, NULL, 1, &barrier);

	}
}