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

extern VkSampler g_invalidTextureSampler;
abstract_class CBrush: public IBrush
{
public:
	void SetPosition( vec3 position ) override;
	void SetRotationEuler( vec3 angle ) override;
	void SetRotationQuat( vec4 quaternion) override;
	void SetMatrix( mat3 matrix ) override;
	void SetScale( vec3 scale ) override;

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

	IMaterial *m_pMaterial = NULL;
	CVertexBuffer *m_pVertexBuffer = NULL;
	CIndexBuffer *m_pIndexBuffer = NULL;
};
void CBrush::SetPosition( vec3 position )
{

}

void CBrush::SetRotationEuler( vec3 angle )
{

}

void CBrush::SetRotationQuat( vec4 quaternion)
{

}

void CBrush::SetMatrix( mat3 matrix )
{

}

void CBrush::SetScale( vec3 scale )
{

}


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

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

CUtlVector<CBrush> g_drawnMeshes;
IMaterial *g_pDefaultMaterial;
IMaterial *g_pCurrentMaterial;

void CBrush::Draw()
{
	g_drawnMeshes.AppendTail(*this);
}

abstract_class CMaterial: public IMaterial
{
};

extern CUtlVector<ITexture*> g_textures;

vk_tripipeline_t g_brushPipeline = {};
vk_image2d_t meshdepth;
vk_image2d_t meshcolor;
extern bool g_bConfigNotify;

VkDescriptorPool g_brushDescriptorPool;
VkDescriptorSet g_brushDescriptorSet;
vk_buffer_t g_brushProjection;
struct MeshProjection {
	mat4 projection;
} *g_brushProject;

VkSampler g_brushSampler;

void IBrushRenderer::Init()
{
	CUtlVector<vk_shader_t> shaders(2);
	for (auto &shader: shaders)
	{
		shader.m_shaderModule = NULL;
	}
	shaders[0].Create("gfx/brush_vert.spv", VK_SHADER_STAGE_VERTEX_BIT);
	shaders[1].Create("gfx/brush_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;
	g_brushPipeline.Create(shaders, bindings, 4);
	shaders[1].Destroy();
	shaders[0].Destroy();

	meshdepth.Create(1280, 720, VK_FORMAT_D32_SFLOAT, VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT);
	meshcolor.Create(1280, 720, VK_FORMAT_R8G8B8A8_UNORM, VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT);

	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_brushDescriptorPool);

	VkDescriptorSetAllocateInfo allocInfo = {};
	allocInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
	allocInfo.descriptorPool = g_brushDescriptorPool;
	allocInfo.descriptorSetCount = 1;
	allocInfo.pSetLayouts = &g_brushPipeline.m_descriptorSetLayout;
	vkAllocateDescriptorSets(g_vkDevice, &allocInfo, &g_brushDescriptorSet);

	g_brushProjection.Create(64, VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT);
	g_brushProject = (MeshProjection*)g_brushProjection.Map(0, 64);

	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_brushSampler);
}

void IBrushRenderer::Frame( float fDelta )
{
	glm_mat4_identity(g_brushProject->projection);
	glm_perspective(glm_rad(90),(float)g_nWindowWidth/g_nWindowHeight, 0.01, 100, g_brushProject->projection);
	glm_rotate(g_brushProject->projection, glm_rad(90), (vec4){1,0,0,0});
	glm_scale(g_brushProject->projection, (vec4){1,-1,1,1});
	glm_rotate(g_brushProject->projection, glm_rad(-90), (vec4){0,0,1,0});
	if (g_bConfigNotify)
	{
		meshdepth.Destroy();
		meshdepth.Create(g_nWindowWidth, g_nWindowHeight, VK_FORMAT_D32_SFLOAT, VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT);
	}


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

	VkDescriptorBufferInfo bufferInfo = {};
	bufferInfo.buffer = g_brushProjection.m_buffer;
	bufferInfo.offset = 0;
	bufferInfo.range = g_brushProjection.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_SHADER_READ_ONLY_OPTIMAL;
		image.imageView = texture->image.m_imageView;
		image.sampler = g_brushSampler;
		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 barrier = {
		.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_swapchainImage,
		.subresourceRange = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1}
	};
	vkCmdPipelineBarrier(g_vkCommandBuffer, 
	VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT, 
	VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT,
	0, 0, NULL, 0, NULL, 1, &barrier);

	VkRenderingAttachmentInfo colorAttachment = {
		.sType = VK_STRUCTURE_TYPE_RENDERING_ATTACHMENT_INFO,
		.imageView = g_swapchainImageView,
		.imageLayout = 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 = meshdepth.m_imageView,
		.imageLayout = 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_1_BIT);
	VkSampleMask sampleMask = 0xFFFFFFFF;
	_vkCmdSetSampleMaskEXT(g_vkCommandBuffer, VK_SAMPLE_COUNT_1_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_brushPipeline.m_pipeline);
	vkCmdBindDescriptorSets(g_vkCommandBuffer,VK_PIPELINE_BIND_POINT_GRAPHICS, g_brushPipeline.m_layout, 0, 1, &g_brushDescriptorSet, 0, NULL);
	for (auto &mesh: g_drawnMeshes)
	{
		VkDeviceSize offset = 0;
		uint32_t textureID = 0;
		if (g_pCurrentMaterial == 0)
			textureID = 0;
		else
			textureID = ((CMaterial*)g_pCurrentMaterial)->m.albedo;
		vkCmdPushConstants(g_vkCommandBuffer, g_brushPipeline.m_layout, VK_SHADER_STAGE_ALL, 0, 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);

	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_swapchainImage,
		.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<CBrush>();
}


IVertexBuffer *IBrushRenderer::CreateVertexBuffer( uint32_t uSize )
{
	CVertexBuffer *pBuffer = new CVertexBuffer();
	pBuffer->m_buffer.Create(uSize, VK_BUFFER_USAGE_VERTEX_BUFFER_BIT | VK_BUFFER_USAGE_STORAGE_BUFFER_BIT);
	return pBuffer;
}

IIndexBuffer *IBrushRenderer::CreateIndexBuffer( uint32_t uSize )
{
	CIndexBuffer *pBuffer = new CIndexBuffer();
	pBuffer->m_buffer.Create(uSize, VK_BUFFER_USAGE_INDEX_BUFFER_BIT | VK_BUFFER_USAGE_STORAGE_BUFFER_BIT);
	return pBuffer;
}

IBrush *IBrushRenderer::CreateMesh()
{
	CBrush *mesh = new CBrush;
	return mesh;
}

void IBrushRenderer::Destroy( IBrush *pModel )
{

}

IMaterial *IBrushRenderer::LoadMaterial( const char *szMaterial )
{
	FileHandle_t file = IFileSystem::Open(szMaterial, IFILE_READ);
	CMaterial *pMaterial = new CMaterial;
	if (!file)
	{
		return g_pDefaultMaterial;
	}

	IFileSystem::Close(file);
}

void IBrushRenderer::SetMaterial( IMaterial *pMaterial )
{
	g_pCurrentMaterial = pMaterial;
}