#include "rendering.h"
#include "vk_helper.h"
#include "vk_video.h"
#include "vulkan/vulkan_core.h"

vk_shader_t post_agxShader = {};
vk_comppipeline_t post_agxPipeline = {};
VkDescriptorPool post_descriptorPool;
VkDescriptorSet post_descriptorSet;

void IPostProcessRenderer::Init()
{
	post_agxShader.Create("gfx/agx_comp.spv", VK_SHADER_STAGE_COMPUTE_BIT);
	CUtlVector<VkDescriptorSetLayoutBinding> bindings = {
		{
			.binding = 0,
			.descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
			.descriptorCount = 1,
			.stageFlags = VK_SHADER_STAGE_COMPUTE_BIT,
		},
		{
			.binding = 1,
			.descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
			.descriptorCount = 1,
			.stageFlags = VK_SHADER_STAGE_COMPUTE_BIT,
		}
	};
	post_agxPipeline.Create(post_agxShader, bindings, 0);
	
	CUtlVector<VkDescriptorPoolSize> pools;
	for (auto &binding: bindings)
	{
		VkDescriptorPoolSize dps = {};
		dps.type = binding.descriptorType;
		dps.descriptorCount = binding.descriptorCount;
		pools.AppendTail(dps);
	}

	VkDescriptorPoolCreateInfo poolInfo = {};
	poolInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;
	poolInfo.poolSizeCount = pools.GetSize();
	poolInfo.pPoolSizes = pools.GetData();
	poolInfo.maxSets = 1;
	vkCreateDescriptorPool(g_vkDevice, &poolInfo, NULL, &post_descriptorPool);

	VkDescriptorSetAllocateInfo allocInfo = {};
	allocInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
	allocInfo.descriptorPool = post_descriptorPool;
	allocInfo.descriptorSetCount = 1;
	allocInfo.pSetLayouts = &post_agxPipeline.m_descriptorSetLayout;
	vkAllocateDescriptorSets(g_vkDevice, &allocInfo, &post_descriptorSet);

}

void IPostProcessRenderer::Frame(float fDelta)
{
	
	CUtlVector<VkImageMemoryBarrier> barriers = {
		{
			.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER,
			.srcAccessMask = 0,
			.dstAccessMask = VK_ACCESS_SHADER_READ_BIT,
			.oldLayout = VK_IMAGE_LAYOUT_GENERAL,
			.newLayout = VK_IMAGE_LAYOUT_GENERAL,
			.image = g_meshColor.m_image,
			.subresourceRange = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1}
		},
		{
			.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER,
			.srcAccessMask = 0,
			.dstAccessMask = VK_ACCESS_SHADER_WRITE_BIT,
			.oldLayout = VK_IMAGE_LAYOUT_UNDEFINED,
			.newLayout = VK_IMAGE_LAYOUT_GENERAL,
			.image = g_swapchainImage, 
			.subresourceRange = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1}
		}
	};

	vkCmdPipelineBarrier(g_vkCommandBuffer, 
	VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT,
	VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
	0, 0, NULL, 0, NULL, barriers.GetSize(), barriers.GetData());

	CUtlVector<VkWriteDescriptorSet> writes(2);
	for (auto &write: writes)
	{
		write = {};
		write.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
		write.dstSet = post_descriptorSet;
		write.dstArrayElement = 0;
	}

	VkDescriptorImageInfo dii1 = {
		.imageView = g_meshColor.m_imageView,
		.imageLayout = VK_IMAGE_LAYOUT_GENERAL,
		.sampler = NULL,
	};
	VkDescriptorImageInfo dii2 = {
		.imageView = g_swapchainImageView,
		.imageLayout = VK_IMAGE_LAYOUT_GENERAL,
		.sampler = NULL,
	};

	writes[0].dstBinding = 0;
	writes[0].descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE;
	writes[0].descriptorCount = 1; 
	writes[0].pImageInfo = &dii1;
	writes[1].dstBinding = 1;
	writes[1].descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE;
	writes[1].descriptorCount = 1; 
	writes[1].pImageInfo = &dii2;
	vkUpdateDescriptorSets(g_vkDevice, writes.GetSize(), writes.GetData(), 0, NULL);

	vkCmdBindPipeline(g_vkCommandBuffer, VK_PIPELINE_BIND_POINT_COMPUTE, post_agxPipeline.m_pipeline);
	vkCmdBindDescriptorSets(g_vkCommandBuffer, VK_PIPELINE_BIND_POINT_COMPUTE, post_agxPipeline.m_layout, 0, 1, &post_descriptorSet, 0, 0);
	vkCmdDispatch(g_vkCommandBuffer, (g_nWindowWidth+31)/32, (g_nWindowHeight+31)/32, 1);

	barriers = {

		{
			.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER,
			.srcAccessMask = VK_ACCESS_SHADER_WRITE_BIT,
			.dstAccessMask = 0,
			.oldLayout = VK_IMAGE_LAYOUT_GENERAL,
			.newLayout = VK_IMAGE_LAYOUT_PRESENT_SRC_KHR,
			.image = g_swapchainImage, 
			.subresourceRange = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1}
		}
	};

	vkCmdPipelineBarrier(g_vkCommandBuffer, 
	VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
	VK_PIPELINE_STAGE_BOTTOM_OF_PIPE_BIT,
	0, 0, NULL, 0, NULL, barriers.GetSize(), barriers.GetData());

}