#include "VulkanPipeline.h"

#include <vulkan/vulkan_core.h>

#include "utilities/FileReader.h"
#include "utilities/Logger.h"
#include "Primitives.h"

void VulkanPipeline::Initialize(VkDevice InDevice)
{
	Device = InDevice;
}

void VulkanPipeline::Cleanup()
{
	// for (auto Framebuffer : SwapChainFrameBuffers)
	// {
	// 	vkDestroyFramebuffer(Device, Framebuffer, nullptr);
	// }

	vkDestroyPipeline(Device, GraphicsPipeline, nullptr);
	vkDestroyPipelineLayout(Device, PipelineLayout, nullptr);
}

void VulkanPipeline::CreateGraphicsPipeline(
	VkExtent2D			   SwapChainExtent,
	VkRenderPass		   RenderPass,
	VkDescriptorSetLayout& DescriptorSetLayout
)
{
	auto VertShaderCode = ReadFile("Shaders/vert.spv");
	auto FragShaderCode = ReadFile("Shaders/frag.spv");

	Log::Info("Vert buffer size: " + std::to_string(VertShaderCode.size()));
	Log::Info("Frag buffer size: " + std::to_string(FragShaderCode.size()));

	VkShaderModule VertShaderModule = CreateShaderModule(VertShaderCode);
	VkShaderModule FragShaderModule = CreateShaderModule(FragShaderCode);

	VkPipelineShaderStageCreateInfo VertShaderStageInfo{};
	VertShaderStageInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
	VertShaderStageInfo.stage = VK_SHADER_STAGE_VERTEX_BIT;
	VertShaderStageInfo.module = VertShaderModule;
	VertShaderStageInfo.pName = "main";

	VkPipelineShaderStageCreateInfo FragShaderStageInfo{};
	FragShaderStageInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
	FragShaderStageInfo.stage = VK_SHADER_STAGE_FRAGMENT_BIT;
	FragShaderStageInfo.module = FragShaderModule;
	FragShaderStageInfo.pName = "main";

	VkPipelineShaderStageCreateInfo ShaderStages[] = { VertShaderStageInfo, FragShaderStageInfo };

	VkPipelineVertexInputStateCreateInfo VertexInputInfo{};
	VertexInputInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO;

	auto BindingDescription = Vertex::GetBindingDescription();
	auto AttributeDescriptions = Vertex::GetAttributeDescriptions();

	VertexInputInfo.vertexBindingDescriptionCount = 1;
	VertexInputInfo.pVertexBindingDescriptions = &BindingDescription;
	VertexInputInfo.vertexAttributeDescriptionCount =
		static_cast<uint32_t>(AttributeDescriptions.size());
	VertexInputInfo.pVertexAttributeDescriptions = AttributeDescriptions.data();

	VkPipelineInputAssemblyStateCreateInfo InputAssembly{};
	InputAssembly.sType = VK_STRUCTURE_TYPE_PIPELINE_INPUT_ASSEMBLY_STATE_CREATE_INFO;
	InputAssembly.topology = VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST;
	InputAssembly.primitiveRestartEnable = VK_FALSE;

	// VkViewport Viewport{};
	// Viewport.x = 0.0f;
	// Viewport.x = 0.0f;
	// Viewport.width = SwapChainExtent.width;
	// Viewport.height = SwapChainExtent.height;
	// Viewport.minDepth = 0.0f;
	// Viewport.maxDepth = 1.0f;

	VkRect2D Scissor{};
	Scissor.offset = { 0, 0 };
	Scissor.extent = SwapChainExtent;

	std::vector<VkDynamicState> DynamicStates = { VK_DYNAMIC_STATE_VIEWPORT,
												  VK_DYNAMIC_STATE_SCISSOR };

	VkPipelineDynamicStateCreateInfo DynamicState{};
	DynamicState.sType = VK_STRUCTURE_TYPE_PIPELINE_DYNAMIC_STATE_CREATE_INFO;
	DynamicState.dynamicStateCount = static_cast<uint32_t>(DynamicStates.size());
	DynamicState.pDynamicStates = DynamicStates.data();

	VkPipelineViewportStateCreateInfo ViewportState{};
	ViewportState.sType = VK_STRUCTURE_TYPE_PIPELINE_VIEWPORT_STATE_CREATE_INFO;
	ViewportState.viewportCount = 1;
	ViewportState.scissorCount = 1;

	VkPipelineRasterizationStateCreateInfo Rasterizer{};
	Rasterizer.sType = VK_STRUCTURE_TYPE_PIPELINE_RASTERIZATION_STATE_CREATE_INFO;
	Rasterizer.depthClampEnable = VK_FALSE;
	Rasterizer.rasterizerDiscardEnable = VK_FALSE;
	Rasterizer.polygonMode = VK_POLYGON_MODE_FILL;
	Rasterizer.lineWidth = 1.0f;
	Rasterizer.cullMode = VK_CULL_MODE_BACK_BIT;
	Rasterizer.frontFace = VK_FRONT_FACE_COUNTER_CLOCKWISE;
	Rasterizer.depthBiasEnable = VK_FALSE;
	Rasterizer.depthBiasConstantFactor = 0.0f;
	Rasterizer.depthBiasClamp = 0.0f;
	Rasterizer.depthBiasSlopeFactor = 0.0f;

	VkPipelineMultisampleStateCreateInfo Multisampling{};
	Multisampling.sType = VK_STRUCTURE_TYPE_PIPELINE_MULTISAMPLE_STATE_CREATE_INFO;
	Multisampling.sampleShadingEnable = VK_FALSE;
	Multisampling.rasterizationSamples = VK_SAMPLE_COUNT_1_BIT;
	Multisampling.minSampleShading = 1.0f;
	Multisampling.pSampleMask = nullptr;
	Multisampling.alphaToCoverageEnable = VK_FALSE;
	Multisampling.alphaToOneEnable = VK_FALSE;

	VkPipelineColorBlendAttachmentState ColorBlendAttachement{};
	ColorBlendAttachement.colorWriteMask = VK_COLOR_COMPONENT_R_BIT | VK_COLOR_COMPONENT_G_BIT
		| VK_COLOR_COMPONENT_B_BIT | VK_COLOR_COMPONENT_A_BIT;
	ColorBlendAttachement.blendEnable = VK_FALSE;
	ColorBlendAttachement.srcColorBlendFactor = VK_BLEND_FACTOR_ONE;
	ColorBlendAttachement.dstColorBlendFactor = VK_BLEND_FACTOR_ZERO;
	ColorBlendAttachement.colorBlendOp = VK_BLEND_OP_ADD;
	ColorBlendAttachement.srcAlphaBlendFactor = VK_BLEND_FACTOR_ONE;
	ColorBlendAttachement.dstAlphaBlendFactor = VK_BLEND_FACTOR_ZERO;
	ColorBlendAttachement.alphaBlendOp = VK_BLEND_OP_ADD;

	VkPipelineColorBlendStateCreateInfo ColorBlending{};
	ColorBlending.sType = VK_STRUCTURE_TYPE_PIPELINE_COLOR_BLEND_STATE_CREATE_INFO;
	ColorBlending.logicOpEnable = VK_FALSE;
	ColorBlending.logicOp = VK_LOGIC_OP_COPY;
	ColorBlending.attachmentCount = 1;
	ColorBlending.pAttachments = &ColorBlendAttachement;
	ColorBlending.blendConstants[0] = 0.0f;
	ColorBlending.blendConstants[1] = 0.0f;
	ColorBlending.blendConstants[2] = 0.0f;
	ColorBlending.blendConstants[3] = 0.0f;

	VkPipelineLayoutCreateInfo PipelineLayoutInfo{};
	PipelineLayoutInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
	PipelineLayoutInfo.setLayoutCount = 1;
	PipelineLayoutInfo.pSetLayouts = &DescriptorSetLayout;
	PipelineLayoutInfo.pushConstantRangeCount = 0;
	PipelineLayoutInfo.pPushConstantRanges = nullptr;

	if (vkCreatePipelineLayout(Device, &PipelineLayoutInfo, nullptr, &PipelineLayout) != VK_SUCCESS)
	{
		Log::Error("Failed to create pipeline layout!");
	}
	else
	{
		Log::Info("Successfully created pipeline layout");
	}

	VkPipelineDepthStencilStateCreateInfo DepthStencil{};
	DepthStencil.sType = VK_STRUCTURE_TYPE_PIPELINE_DEPTH_STENCIL_STATE_CREATE_INFO;
	DepthStencil.depthTestEnable = VK_TRUE;
	DepthStencil.depthWriteEnable = VK_TRUE;
	DepthStencil.depthCompareOp = VK_COMPARE_OP_LESS;
	DepthStencil.depthBoundsTestEnable = VK_FALSE;
	DepthStencil.minDepthBounds = 0.0f;
	DepthStencil.maxDepthBounds = 1.0f;
	DepthStencil.stencilTestEnable = VK_FALSE;
	DepthStencil.front = {};
	DepthStencil.back = {};

	VkGraphicsPipelineCreateInfo PipelineInfo{};
	PipelineInfo.sType = VK_STRUCTURE_TYPE_GRAPHICS_PIPELINE_CREATE_INFO;
	PipelineInfo.stageCount = 2;
	PipelineInfo.pStages = ShaderStages;
	PipelineInfo.pVertexInputState = &VertexInputInfo;
	PipelineInfo.pInputAssemblyState = &InputAssembly;
	PipelineInfo.pViewportState = &ViewportState;
	PipelineInfo.pRasterizationState = &Rasterizer;
	PipelineInfo.pMultisampleState = &Multisampling;
	PipelineInfo.pDepthStencilState = &DepthStencil;
	PipelineInfo.pColorBlendState = &ColorBlending;
	PipelineInfo.pDynamicState = &DynamicState;
	PipelineInfo.layout = PipelineLayout;
	PipelineInfo.renderPass = RenderPass;
	PipelineInfo.subpass = 0;
	PipelineInfo.basePipelineHandle = VK_NULL_HANDLE;
	PipelineInfo.basePipelineIndex = -1;

	if (vkCreateGraphicsPipelines(
			Device,
			VK_NULL_HANDLE,
			1,
			&PipelineInfo,
			nullptr,
			&GraphicsPipeline
		)
		!= VK_SUCCESS)
	{
		Log::Error("Failed to create graphics pipeline!");
	}
	else
	{
		Log::Info("Successfully created graphics pipeline.");
	}

	vkDestroyShaderModule(Device, FragShaderModule, nullptr);
	vkDestroyShaderModule(Device, VertShaderModule, nullptr);
}

VkShaderModule VulkanPipeline::CreateShaderModule(const std::vector<char>& Code)
{
	VkShaderModuleCreateInfo CreateInfo{};
	CreateInfo.sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO;
	CreateInfo.codeSize = Code.size();
	CreateInfo.pCode = reinterpret_cast<const uint32_t*>(Code.data());

	VkShaderModule ShaderModule;
	if (vkCreateShaderModule(Device, &CreateInfo, nullptr, &ShaderModule) != VK_SUCCESS)
	{
		Log::Error("Failed to create shader module.");
	}
	else
	{
		Log::Info("Successfully created shader module.");
	}

	return ShaderModule;
}

// void VulkanPipeline::CreateRenderPass(VkFormat SwapChainImageFormat)
// {
// 	VkAttachmentDescription ColorAttachment{};
// 	ColorAttachment.format = SwapChainImageFormat;
// 	ColorAttachment.samples = VK_SAMPLE_COUNT_1_BIT;
// 	ColorAttachment.loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR;
// 	ColorAttachment.storeOp = VK_ATTACHMENT_STORE_OP_STORE;
// 	ColorAttachment.stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE;
// 	ColorAttachment.stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE;
// 	ColorAttachment.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;
// 	ColorAttachment.finalLayout = VK_IMAGE_LAYOUT_PRESENT_SRC_KHR;
//
// 	VkAttachmentReference ColorAttachmentRef{};
// 	ColorAttachmentRef.attachment = 0;
// 	ColorAttachmentRef.layout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;
//
// 	VkSubpassDescription Subpass{};
// 	Subpass.pipelineBindPoint = VK_PIPELINE_BIND_POINT_GRAPHICS;
// 	Subpass.colorAttachmentCount = 1;
// 	Subpass.pColorAttachments = &ColorAttachmentRef;
//
// 	VkRenderPassCreateInfo RenderPassInfo{};
// 	RenderPassInfo.sType = VK_STRUCTURE_TYPE_RENDER_PASS_CREATE_INFO;
// 	RenderPassInfo.attachmentCount = 1;
// 	RenderPassInfo.pAttachments = &ColorAttachment;
// 	RenderPassInfo.subpassCount = 1;
// 	RenderPassInfo.pSubpasses = &Subpass;
//
// 	VkSubpassDependency Dependency{};
// 	Dependency.srcSubpass = VK_SUBPASS_EXTERNAL;
// 	Dependency.dstSubpass = 0;
// 	Dependency.srcStageMask = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT;
// 	Dependency.srcAccessMask = 0;
// 	Dependency.dstStageMask = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT;
// 	Dependency.dstAccessMask = VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT;
// 	RenderPassInfo.dependencyCount = 1;
// 	RenderPassInfo.pDependencies = &Dependency;
//
// 	if (vkCreateRenderPass(Device, &RenderPassInfo, nullptr, &RenderPass) != VK_SUCCESS)
// 	{
// 		Log::Error("Failed to create render pass!");
// 	}
// 	else
// 	{
// 		Log::Info("Successfully created render pass.");
// 	}
// }

// void VulkanPipeline::CreateFramebuffers(std::vector<VkImageView> SwapChainImageViews, VkExtent2D
// SwapChainExtent)
// {
// 	SwapChainFrameBuffers.resize(SwapChainImageViews.size());
//
// 	for (size_t i = 0; i < SwapChainImageViews.size(); i++)
// 	{
// 		VkImageView Attachments[] = {
// 			SwapChainImageViews[i]
// 		};
//
// 		VkFramebufferCreateInfo FramebufferInfo{};
// 		FramebufferInfo.sType = VK_STRUCTURE_TYPE_FRAMEBUFFER_CREATE_INFO;
// 		FramebufferInfo.renderPass = RenderPass;
// 		FramebufferInfo.attachmentCount = 1;
// 		FramebufferInfo.pAttachments = Attachments;
// 		FramebufferInfo.width = SwapChainExtent.width;
// 		FramebufferInfo.height = SwapChainExtent.height;
// 		FramebufferInfo.layers = 1;
//
// 		if (vkCreateFramebuffer(Device, &FramebufferInfo, nullptr, &SwapChainFrameBuffers[i]) !=
// VK_SUCCESS)
// 		{
// 			Log::Error("Failed to create frame buffer!");
// 		}
// 	}
// }
