#include "cglm/mat4.h"
#include "filesystem.h"
#include "rendering.h"
#include "tier1/utlvector.h"
#include "vk_helper.h"
#include "vulkan/vulkan_core.h"


vk_tripipeline_t g_meshPipeline = {};

VkDescriptorPool g_meshDescriptorPool;
VkDescriptorSet g_meshDescriptorSet;

VkSampler g_meshSampler;


abstract_class CMesh: public IMesh
{
public:
	CMesh();
	void SetPosition( vec3 position ) override;
	void SetRotationEuler( vec3 angle ) override;
	void SetRotationQuat( vec4 quaternion) override;
	void SetMatrix( mat4 matrix ) override;
	void SetScale( vec3 scale ) override;

	void SetVertexBuffer( IVertexBuffer *pBuffer ) override;
	void SetIndexBuffer( IIndexBuffer *pBuffer ) override;
	void Draw() override;

	Material_t m_material;
	CVertexBuffer *m_pVertexBuffer = NULL;
	CIndexBuffer *m_pIndexBuffer = NULL;
	mat4 m_matrix;
};

CMesh::CMesh()
{
	glm_mat4_identity(m_matrix);
};

void CMesh::SetPosition( vec3 position )
{
	m_matrix[0][3] = position[0];
	m_matrix[1][3] = position[1];
	m_matrix[2][3] = position[2];
}

void CMesh::SetRotationEuler( vec3 angle )
{

}

void CMesh::SetRotationQuat( vec4 quaternion)
{

}

void CMesh::SetMatrix( mat4 matrix )
{
	memcpy(m_matrix,matrix,64);
}

void CMesh::SetScale( vec3 scale )
{

}

void CMesh::SetVertexBuffer( IVertexBuffer *pBuffer )
{
	m_pVertexBuffer = (CVertexBuffer*)pBuffer;
}

void CMesh::SetIndexBuffer( IIndexBuffer *pBuffer )
{
	m_pIndexBuffer = (CIndexBuffer*)pBuffer;
}

CUtlVector<CMesh> g_drawnMeshes;

void CMesh::Draw()
{
	if (!g_pCurrentMaterial)
		m_material = {};
	else
		m_material = g_pCurrentMaterial->m;
	g_drawnMeshes.AppendTail(*this);
}

void IMeshRenderer::Init()
{	
	CUtlVector<vk_shader_t> shaders(2);
	for (auto &shader: shaders)
	{
		shader.m_shaderModule = NULL;
	}
	shaders[0].Create("gfx/mesh_vert.spv", VK_SHADER_STAGE_VERTEX_BIT);
	shaders[1].Create("gfx/mesh_frag.spv", VK_SHADER_STAGE_FRAGMENT_BIT);
	CUtlVector<VkDescriptorSetLayoutBinding> bindings(2);
	bindings[0] = {};
	bindings[0].binding = 0;
	bindings[0].stageFlags = VK_SHADER_STAGE_VERTEX_BIT;
	bindings[0].descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
	bindings[0].descriptorCount = 1;
	bindings[1] = {};
	bindings[1].binding = 1;
	bindings[1].stageFlags = VK_SHADER_STAGE_FRAGMENT_BIT;
	bindings[1].descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
	bindings[1].descriptorCount = 1024;
	CUtlVector<VkFormat> formats(1);
	formats[0] = g_meshColor.m_format;
	g_meshPipeline.Create(shaders, bindings, 76, formats);
	shaders[1].Destroy();
	shaders[0].Destroy();

	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, &g_meshDescriptorPool);

	VkDescriptorSetAllocateInfo allocInfo = {};
	allocInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
	allocInfo.descriptorPool = g_meshDescriptorPool;
	allocInfo.descriptorSetCount = 1;
	allocInfo.pSetLayouts = &g_meshPipeline.m_descriptorSetLayout;
	vkAllocateDescriptorSets(g_vkDevice, &allocInfo, &g_meshDescriptorSet);


	VkPhysicalDeviceProperties properties{};
	vkGetPhysicalDeviceProperties(g_vkPhysicalDevice, &properties);
	VkSamplerCreateInfo samplerInfo{};
	samplerInfo.sType = VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO;
	samplerInfo.magFilter = VK_FILTER_LINEAR;
	samplerInfo.minFilter = VK_FILTER_LINEAR;
	samplerInfo.addressModeU = VK_SAMPLER_ADDRESS_MODE_REPEAT;
	samplerInfo.addressModeV = VK_SAMPLER_ADDRESS_MODE_REPEAT;
	samplerInfo.addressModeW = VK_SAMPLER_ADDRESS_MODE_REPEAT;
	samplerInfo.anisotropyEnable = VK_FALSE;
	samplerInfo.maxAnisotropy = properties.limits.maxSamplerAnisotropy;
	samplerInfo.borderColor = VK_BORDER_COLOR_INT_OPAQUE_BLACK;
	samplerInfo.unnormalizedCoordinates = VK_FALSE;
	samplerInfo.compareEnable = VK_FALSE;
	samplerInfo.compareOp = VK_COMPARE_OP_ALWAYS;
	samplerInfo.mipmapMode = VK_SAMPLER_MIPMAP_MODE_LINEAR;
	samplerInfo.mipLodBias = 0.0f;
	samplerInfo.minLod = 0.0f;
	samplerInfo.maxLod = 0.0f;
	vkCreateSampler(g_vkDevice, &samplerInfo, nullptr, &g_meshSampler);
}

void IMeshRenderer_PrepassNoMSAA()
{

};

void IMeshRenderer_Prepass()
{

};

void IMeshRenderer_EdgeDetection()
{

};

void IMeshRenderer_Light()
{

};

void IMeshRenderer_Combine()
{

};

void IMeshRenderer::Frame( float fDelta )
{
	CUtlVector<VkWriteDescriptorSet> writes(2);
	for (auto &write: writes)
	{
		write = {};
		write.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
		write.dstSet = g_meshDescriptorSet;
		write.dstArrayElement = 0;
	}

	VkDescriptorBufferInfo bufferInfo = {};
	bufferInfo.buffer = g_cameraProperties.m_buffer;
	bufferInfo.offset = 0;
	bufferInfo.range = g_cameraProperties.m_nSize;
	writes[0].dstBinding = 0;
	writes[0].descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;	
	writes[0].descriptorCount = 1;
	writes[0].pBufferInfo = &bufferInfo;

	CUtlVector<VkDescriptorImageInfo> textures;
	textures.Reserve(g_textures.GetSize());
	for (ITexture *t: g_textures)
	{
		CTexture *texture = (CTexture*)t;
		VkDescriptorImageInfo image = {};
		image.imageLayout = VK_IMAGE_LAYOUT_GENERAL;
		image.imageView = texture->image.m_imageView;
		image.sampler = g_meshSampler;
		textures.AppendTail(image);
	};
	textures[0].sampler = g_invalidTextureSampler;
	writes[1].dstBinding = 1;
	writes[1].descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
	writes[1].descriptorCount = textures.GetSize();
	writes[1].pImageInfo = textures.GetData();
	vkUpdateDescriptorSets(g_vkDevice, writes.GetSize(), writes.GetData(), 0, NULL);


	VkImageMemoryBarrier barriers[4] = {
	{
		.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER,
		.srcAccessMask = 0,
		.dstAccessMask = VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT,
		.oldLayout = VK_IMAGE_LAYOUT_UNDEFINED,
		.newLayout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL,
		.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_COLOR_ATTACHMENT_WRITE_BIT,
		.oldLayout = VK_IMAGE_LAYOUT_UNDEFINED,
		.newLayout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL,
		.image = g_meshColorMSAA.m_image,
		.subresourceRange = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1}
	},
	{
		.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER,
		.srcAccessMask = 0,
		.dstAccessMask = VK_ACCESS_DEPTH_STENCIL_ATTACHMENT_WRITE_BIT,
		.oldLayout = VK_IMAGE_LAYOUT_UNDEFINED,
		.newLayout = VK_IMAGE_LAYOUT_DEPTH_ATTACHMENT_OPTIMAL,
		.image = g_meshDepth.m_image,
		.subresourceRange = {VK_IMAGE_ASPECT_DEPTH_BIT, 0, 1, 0, 1}
	},
	{
		.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER,
		.srcAccessMask = 0,
		.dstAccessMask = VK_ACCESS_DEPTH_STENCIL_ATTACHMENT_WRITE_BIT,
		.oldLayout = VK_IMAGE_LAYOUT_UNDEFINED,
		.newLayout = VK_IMAGE_LAYOUT_DEPTH_ATTACHMENT_OPTIMAL,
		.image = g_meshDepthMSAA.m_image,
		.subresourceRange = {VK_IMAGE_ASPECT_DEPTH_BIT, 0, 1, 0, 1}
	},
	};
	vkCmdPipelineBarrier(g_vkCommandBuffer, 
	VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT, 
	VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT | VK_PIPELINE_STAGE_EARLY_FRAGMENT_TESTS_BIT,
	0, 0, NULL, 0, NULL, 4, barriers);

	VkRenderingAttachmentInfo colorAttachment = {
		.sType = VK_STRUCTURE_TYPE_RENDERING_ATTACHMENT_INFO,
		.imageView = g_meshColorMSAA.m_imageView,
		.imageLayout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL,
		.resolveMode = VK_RESOLVE_MODE_AVERAGE_BIT,
		.resolveImageView = g_meshColor.m_imageView,
		.resolveImageLayout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL,
		.loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR,
		.storeOp = VK_ATTACHMENT_STORE_OP_STORE,
		.clearValue = {.color = {0.0f, 0.0f, 0.0f, 1.0f}},
	};
	VkRenderingAttachmentInfo depthAttachment = {
		.sType = VK_STRUCTURE_TYPE_RENDERING_ATTACHMENT_INFO,
		.imageView = g_meshDepthMSAA.m_imageView,
		.imageLayout = VK_IMAGE_LAYOUT_DEPTH_ATTACHMENT_OPTIMAL,
		.resolveMode = VK_RESOLVE_MODE_AVERAGE_BIT,
		.resolveImageView = g_meshDepth.m_imageView,
		.resolveImageLayout = VK_IMAGE_LAYOUT_DEPTH_ATTACHMENT_OPTIMAL,
		.loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR,
		.storeOp = VK_ATTACHMENT_STORE_OP_STORE,
		.clearValue = {.depthStencil = {.depth = 1}},
	};

	VkRenderingInfo renderInfo = {
		.sType = VK_STRUCTURE_TYPE_RENDERING_INFO,
		.renderArea = {{0, 0}, {g_nWindowWidth, g_nWindowHeight}},
		.layerCount = 1,
		.colorAttachmentCount = 1,
		.pColorAttachments = &colorAttachment,
		.pDepthAttachment = &depthAttachment,
	};
	vkCmdBeginRendering(g_vkCommandBuffer, &renderInfo);
	
	vkCmdSetRasterizerDiscardEnable(g_vkCommandBuffer, VK_FALSE);
	vkCmdSetDepthBiasEnable(g_vkCommandBuffer, VK_FALSE);
	_vkCmdSetPolygonModeEXT(g_vkCommandBuffer, VK_POLYGON_MODE_FILL);
	vkCmdSetCullMode(g_vkCommandBuffer, VK_CULL_MODE_BACK_BIT);
	vkCmdSetFrontFace(g_vkCommandBuffer, VK_FRONT_FACE_COUNTER_CLOCKWISE);

	vkCmdSetDepthTestEnable(g_vkCommandBuffer, VK_TRUE);
	vkCmdSetDepthWriteEnable(g_vkCommandBuffer, VK_TRUE);
	vkCmdSetDepthCompareOp(g_vkCommandBuffer, VK_COMPARE_OP_LESS);
	vkCmdSetStencilTestEnable(g_vkCommandBuffer, VK_FALSE);

	_vkCmdSetRasterizationSamplesEXT(g_vkCommandBuffer, VK_SAMPLE_COUNT_4_BIT);
	VkSampleMask sampleMask = 0xFFFFFFFF;
	_vkCmdSetSampleMaskEXT(g_vkCommandBuffer, VK_SAMPLE_COUNT_4_BIT, &sampleMask);
	_vkCmdSetAlphaToCoverageEnableEXT(g_vkCommandBuffer, VK_FALSE);

	VkViewport viewport = {0, 0, (float)g_nWindowWidth, (float)g_nWindowHeight, 0.0f, 1.0f};
	VkRect2D scissor = {{0, 0}, {g_nWindowWidth, g_nWindowHeight}};
	vkCmdSetViewportWithCount(g_vkCommandBuffer, 1, &viewport);
	vkCmdSetScissorWithCount(g_vkCommandBuffer, 1, &scissor);

	vkCmdSetPrimitiveTopology(g_vkCommandBuffer, VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST);
	vkCmdSetPrimitiveRestartEnable(g_vkCommandBuffer, VK_FALSE);

	VkVertexInputBindingDescription2EXT binding = {
		.sType = VK_STRUCTURE_TYPE_VERTEX_INPUT_BINDING_DESCRIPTION_2_EXT,
		.binding = 0,
		.stride = 20,
		.inputRate = VK_VERTEX_INPUT_RATE_VERTEX,
		.divisor = 1,
	};
	VkVertexInputAttributeDescription2EXT attributes[2] = {
		{
			VK_STRUCTURE_TYPE_VERTEX_INPUT_ATTRIBUTE_DESCRIPTION_2_EXT, 
			NULL,
			0, 0, 
			VK_FORMAT_R32G32B32_SFLOAT, 
			0
		},
		{
			VK_STRUCTURE_TYPE_VERTEX_INPUT_ATTRIBUTE_DESCRIPTION_2_EXT, 
			NULL,
			1, 0,
			VK_FORMAT_R32G32_SFLOAT, 
			12
		}
	};
	_vkCmdSetVertexInputEXT(g_vkCommandBuffer, 1, &binding, 2, attributes);

	VkBool32 blendEnable = VK_FALSE;
	VkColorBlendEquationEXT blendEquation = {
	VK_BLEND_FACTOR_SRC_ALPHA, VK_BLEND_FACTOR_ONE_MINUS_SRC_ALPHA, VK_BLEND_OP_ADD,
	VK_BLEND_FACTOR_ONE, VK_BLEND_FACTOR_ZERO, VK_BLEND_OP_ADD
	};
	VkColorComponentFlags writeMask = VK_COLOR_COMPONENT_R_BIT | VK_COLOR_COMPONENT_G_BIT |
				      VK_COLOR_COMPONENT_B_BIT | VK_COLOR_COMPONENT_A_BIT;

	_vkCmdSetColorBlendEnableEXT(g_vkCommandBuffer, 0, 1, &blendEnable);
	_vkCmdSetColorBlendEquationEXT(g_vkCommandBuffer, 0, 1, &blendEquation);
	_vkCmdSetColorWriteMaskEXT(g_vkCommandBuffer, 0, 1, &writeMask);
	
	vkCmdBindPipeline(g_vkCommandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, g_meshPipeline.m_pipeline);
	vkCmdBindDescriptorSets(g_vkCommandBuffer,VK_PIPELINE_BIND_POINT_GRAPHICS, g_meshPipeline.m_layout, 0, 1, &g_meshDescriptorSet, 0, NULL);
	for (auto &mesh: g_drawnMeshes)
	{
		VkDeviceSize offset = 0;
		uint32_t textureID = 0;
		vkCmdPushConstants(g_vkCommandBuffer, g_meshPipeline.m_layout, VK_SHADER_STAGE_ALL, 0, 64, &mesh.m_matrix);
		vkCmdPushConstants(g_vkCommandBuffer, g_meshPipeline.m_layout, VK_SHADER_STAGE_ALL, 64, 4, &textureID);
		vkCmdBindVertexBuffers(g_vkCommandBuffer, 0, 1, &mesh.m_pVertexBuffer->m_buffer.m_buffer, &offset);
		if (mesh.m_pIndexBuffer)
		{
			vkCmdBindIndexBuffer(
					g_vkCommandBuffer, 
					mesh.m_pIndexBuffer->m_buffer.m_buffer, 
					0, 
					VK_INDEX_TYPE_UINT32
					);
			vkCmdDrawIndexed(g_vkCommandBuffer, mesh.m_pIndexBuffer->m_buffer.m_nSize/4, 1, 0, 0, 0);
		}
		else 
		{
			vkCmdDraw(g_vkCommandBuffer, mesh.m_pVertexBuffer->m_buffer.m_nSize/12,1,0,0);
		}
	}
	vkCmdEndRendering(g_vkCommandBuffer);

	VkImageMemoryBarrier barrier = {
		.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER,
		.srcAccessMask = VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT,
		.dstAccessMask = 0,
		.oldLayout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL,
		.newLayout = VK_IMAGE_LAYOUT_PRESENT_SRC_KHR,
		.image = g_meshColor.m_image,
		.subresourceRange = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1}
	};

	vkCmdPipelineBarrier(g_vkCommandBuffer, 
	VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT, 
	VK_PIPELINE_STAGE_BOTTOM_OF_PIPE_BIT,
	0, 0, NULL, 0, NULL, 1, &barrier);
	
	barrier = {
		.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER,
		.srcAccessMask = VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT,
		.dstAccessMask = 0,
		.oldLayout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL,
		.newLayout = VK_IMAGE_LAYOUT_PRESENT_SRC_KHR,
		.image = g_meshColor.m_image,
		.subresourceRange = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1}
	};

	vkCmdPipelineBarrier(g_vkCommandBuffer, 
	VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT, 
	VK_PIPELINE_STAGE_BOTTOM_OF_PIPE_BIT,
	0, 0, NULL, 0, NULL, 1, &barrier);


	g_drawnMeshes = CUtlVector<CMesh>();

}


IMesh *IMeshRenderer::CreateMesh()
{
	CMesh *mesh = new CMesh;
	return mesh;
}

void IMeshRenderer::Destroy( IMesh *pModel )
{

}