#include "VulkanBuffers.h"

#include "imgui.h"
#include "imgui_impl_vulkan.h"
#include "utilities/Logger.h"
#include <cstddef>
#include <vulkan/vulkan_core.h>

void VulkanBuffers::Initialize(FBufferConfig InConfig)
// VkDevice	 InDevice,
// VkRenderPass InRenderPass)
{
	Config = InConfig;
}

void VulkanBuffers::Cleanup(size_t MAX_FRAMES_IN_FLIGHT)
{
	vkDestroyCommandPool(Config.Device, CommandPool, nullptr);

	for (size_t i = 0; i < MAX_FRAMES_IN_FLIGHT; i++)
	{
		vkDestroyBuffer(Config.Device, UniformBuffers[i], nullptr);
		vkFreeMemory(Config.Device, UniformBuffersMemory[i], nullptr);
	}

	vkDestroyBuffer(Config.Device, IndexBuffer, nullptr);
	vkFreeMemory(Config.Device, IndexBufferMemory, nullptr);

	vkDestroyBuffer(Config.Device, VertexBuffer, nullptr);
	vkFreeMemory(Config.Device, VertexBufferMemory, nullptr);
}

void VulkanBuffers::CreateCommandPool(std::optional<uint32_t> GraphicsFamily)
{
	VkCommandPoolCreateInfo PoolInfo{};
	PoolInfo.sType = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO;
	PoolInfo.flags = VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT;
	PoolInfo.queueFamilyIndex = GraphicsFamily.value();

	if (vkCreateCommandPool(Config.Device, &PoolInfo, nullptr, &CommandPool) != VK_SUCCESS)
	{
		Log::Error("Failed to create command pool!");
	}
	else
	{
		Log::Info("Successfully created command pool");
	}
}

void VulkanBuffers::CreateCommandBuffers(int FramesInFlight)
{
	CommandBuffers.resize(FramesInFlight);
	VkCommandBufferAllocateInfo AllocateInfo{};
	AllocateInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
	AllocateInfo.commandPool = CommandPool;
	AllocateInfo.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
	AllocateInfo.commandBufferCount = (uint32_t)CommandBuffers.size();

	if (vkAllocateCommandBuffers(Config.Device, &AllocateInfo, CommandBuffers.data()) != VK_SUCCESS)
	{
		Log::Error("Failed to allocate command buffers!");
	}
	else
	{
		Log::Info("Successfully allocated command buffers.");
	}
}

void VulkanBuffers::RecordCommandBuffer(FRecordBuffersParams& Params)
{
	VkCommandBufferBeginInfo BeginInfo{};
	BeginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
	BeginInfo.flags = 0;
	BeginInfo.pInheritanceInfo = nullptr;

	if (vkBeginCommandBuffer(Params.InCommandBuffer, &BeginInfo) != VK_SUCCESS)
	{
		Log::Error("Failed to begin recording command buffer!");
	}
	else
	{
		// Log::Info("Successfully began recording command buffers.");
	}

	VkRenderPassBeginInfo RenderPassInfo{};
	RenderPassInfo.sType = VK_STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO;
	RenderPassInfo.renderPass = Config.RenderPass;
	RenderPassInfo.framebuffer = Params.SwapChainFramebuffers[Params.ImageIndex];
	RenderPassInfo.renderArea.offset = { 0, 0 };
	RenderPassInfo.renderArea.extent = { Params.SwapChainExtent };

	std::array<VkClearValue, 2> ClearValues{};
	ClearValues[0].color = { { 0.0f, 0.0f, 0.0f, 1.0f } };
	ClearValues[1].depthStencil = { 1.0f, 0 };
	RenderPassInfo.clearValueCount = static_cast<uint32_t>(ClearValues.size());
	RenderPassInfo.pClearValues = ClearValues.data();

	vkCmdBeginRenderPass(Params.InCommandBuffer, &RenderPassInfo, VK_SUBPASS_CONTENTS_INLINE);

	vkCmdBindPipeline(
		Params.InCommandBuffer,
		VK_PIPELINE_BIND_POINT_GRAPHICS,
		Params.GraphicsPipeline
	);

	VkViewport Viewport{};
	Viewport.x = 0.0f;
	Viewport.y = 0.0f;
	Viewport.width = static_cast<float>(Params.SwapChainExtent.width);
	Viewport.height = static_cast<float>(Params.SwapChainExtent.height);
	Viewport.minDepth = 0.0f;
	Viewport.maxDepth = 1.0f;
	vkCmdSetViewport(Params.InCommandBuffer, 0, 1, &Viewport);

	VkRect2D Scissor{};
	Scissor.offset = { 0, 0 };
	Scissor.extent = Params.SwapChainExtent;
	vkCmdSetScissor(Params.InCommandBuffer, 0, 1, &Scissor);

	VkBuffer	 VertexBuffers[] = { Params.InVertexBuffer };
	VkDeviceSize Offsets[] = { 0 };
	vkCmdBindVertexBuffers(Params.InCommandBuffer, 0, 1, VertexBuffers, Offsets);
	vkCmdBindIndexBuffer(Params.InCommandBuffer, Params.InIndexBuffer, 0, VK_INDEX_TYPE_UINT16);

	vkCmdBindDescriptorSets(
		Params.InCommandBuffer,
		VK_PIPELINE_BIND_POINT_GRAPHICS,
		Params.PipelineLayout,
		0,
		1,
		&Params.DescriptorSet,
		0,
		nullptr
	);

	// vkCmdDraw(Params.InCommandBuffer, static_cast<uint32_t>(Params.InVerticesSize), 1, 0, 0);
	vkCmdDrawIndexed(Params.InCommandBuffer, static_cast<uint32_t>(Params.InIndexSize), 1, 0, 0, 0);

	if (Params.bDrawImGui)
	{
		ImGui::Render();
		ImGui_ImplVulkan_RenderDrawData(ImGui::GetDrawData(), Params.InCommandBuffer);
	}

	vkCmdEndRenderPass(Params.InCommandBuffer);

	if (vkEndCommandBuffer(Params.InCommandBuffer) != VK_SUCCESS)
	{
		Log::Error("Failed to record command buffer!");
	}
	else
	{
		// Log::Info("Successfully recorded command buffer.");
	}
}

VkCommandBuffer VulkanBuffers::BeginSingleTimeCommands()
{
	VkCommandBufferAllocateInfo AllocateInfo{};
	AllocateInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
	AllocateInfo.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
	AllocateInfo.commandPool = CommandPool;
	AllocateInfo.commandBufferCount = 1;

	VkCommandBuffer CommandBuffer;
	vkAllocateCommandBuffers(Config.Device, &AllocateInfo, &CommandBuffer);

	VkCommandBufferBeginInfo BeginInfo{};
	BeginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
	BeginInfo.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT;

	vkBeginCommandBuffer(CommandBuffer, &BeginInfo);

	return CommandBuffer;
}

void VulkanBuffers::EndSingleTimeCommands(VkCommandBuffer InCommandBuffer, VkQueue GraphicsQueue)
{
	vkEndCommandBuffer(InCommandBuffer);

	VkSubmitInfo SubmitInfo{};
	SubmitInfo.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;
	SubmitInfo.commandBufferCount = 1;
	SubmitInfo.pCommandBuffers = &InCommandBuffer;

	vkQueueSubmit(GraphicsQueue, 1, &SubmitInfo, VK_NULL_HANDLE);
	vkQueueWaitIdle(GraphicsQueue);

	vkFreeCommandBuffers(Config.Device, CommandPool, 1, &InCommandBuffer);
}

void VulkanBuffers::TransitionImageLayout(
	VkImage		  Image,
	VkFormat	  Format,
	VkImageLayout OldLayout,
	VkImageLayout NewLayout,
	bool		  bHasStencilComponent
)
{
	VkCommandBuffer CommandBuffer = BeginSingleTimeCommands();

	VkImageMemoryBarrier Barrier{};
	Barrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
	Barrier.oldLayout = OldLayout;
	Barrier.newLayout = NewLayout;
	Barrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
	Barrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
	Barrier.image = Image;
	Barrier.subresourceRange.baseMipLevel = 0;
	Barrier.subresourceRange.levelCount = 1;
	Barrier.subresourceRange.baseArrayLayer = 0;
	Barrier.subresourceRange.layerCount = 1;
	Barrier.srcAccessMask = 0;
	Barrier.dstAccessMask = 0;

	VkPipelineStageFlags SourceStage;
	VkPipelineStageFlags DestinationStage;

	if (NewLayout == VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL)
	{
		Barrier.subresourceRange.aspectMask = VK_IMAGE_ASPECT_DEPTH_BIT;

		if (bHasStencilComponent)
		{
			Barrier.subresourceRange.aspectMask |= VK_IMAGE_ASPECT_STENCIL_BIT;
		}
	}
	else
	{
		Barrier.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
	}

	if (OldLayout == VK_IMAGE_LAYOUT_UNDEFINED && NewLayout == VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL)
	{
		Barrier.srcAccessMask = 0;
		Barrier.dstAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;

		SourceStage = VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT;
		DestinationStage = VK_PIPELINE_STAGE_TRANSFER_BIT;
	}
	else if (OldLayout == VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL
			 && NewLayout == VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
	{
		Barrier.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;
		Barrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT;

		SourceStage = VK_PIPELINE_STAGE_TRANSFER_BIT;
		DestinationStage = VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT;
	}
	else if (OldLayout == VK_IMAGE_LAYOUT_UNDEFINED
			 && NewLayout == VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL)
	{
		Barrier.srcAccessMask = 0;
		Barrier.dstAccessMask = VK_ACCESS_DEPTH_STENCIL_ATTACHMENT_READ_BIT
			| VK_ACCESS_DEPTH_STENCIL_ATTACHMENT_WRITE_BIT;

		SourceStage = VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT;
		DestinationStage = VK_PIPELINE_STAGE_EARLY_FRAGMENT_TESTS_BIT;
	}
	else
	{
		Log::Error("Unsupported layout transition!");
	}

	vkCmdPipelineBarrier(
		CommandBuffer,
		SourceStage,
		DestinationStage,
		0,
		0,
		nullptr,
		0,
		nullptr,
		1,
		&Barrier
	);

	EndSingleTimeCommands(CommandBuffer, Config.GraphicsQueue);
}

void VulkanBuffers::CopyBufferToImage(
	VkBuffer Buffer,
	VkImage	 Image,
	uint32_t Width,
	uint32_t Height
)
{
	VkCommandBuffer CommandBuffer = BeginSingleTimeCommands();

	VkBufferImageCopy Region{};
	Region.bufferOffset = 0;
	Region.bufferRowLength = 0;
	Region.bufferImageHeight = 0;
	Region.imageSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
	Region.imageSubresource.mipLevel = 0;
	Region.imageSubresource.baseArrayLayer = 0;
	Region.imageSubresource.layerCount = 1;

	Region.imageOffset = { 0, 0, 0 };
	Region.imageExtent = { Width, Height, 1 };

	vkCmdCopyBufferToImage(
		CommandBuffer,
		Buffer,
		Image,
		VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
		1,
		&Region
	);

	EndSingleTimeCommands(CommandBuffer, Config.GraphicsQueue);
}

void VulkanBuffers::CopyBuffer(
	VkQueue		 GraphicsQueue,
	VkBuffer	 SrcBuffer,
	VkBuffer	 DstBuffer,
	VkDeviceSize Size
)
{
	VkCommandBuffer CommandBuffer = BeginSingleTimeCommands();

	VkBufferCopy CopyRegion{};
	CopyRegion.size = Size;
	vkCmdCopyBuffer(CommandBuffer, SrcBuffer, DstBuffer, 1, &CopyRegion);

	EndSingleTimeCommands(CommandBuffer, GraphicsQueue);
}

void VulkanBuffers::CreateBuffer(
	VkDeviceSize		  Size,
	VkBufferUsageFlags	  Usage,
	VkMemoryPropertyFlags Properties,
	VkBuffer&			  Buffer,
	VkDeviceMemory&		  BufferMemory
)
{
	VkBufferCreateInfo BufferInfo{};
	BufferInfo.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO;
	BufferInfo.size = Size;
	BufferInfo.usage = Usage;
	BufferInfo.sharingMode = VK_SHARING_MODE_EXCLUSIVE;

	if (vkCreateBuffer(Config.Device, &BufferInfo, nullptr, &Buffer) != VK_SUCCESS)
	{
		throw std::runtime_error("failed to create buffer!");
	}

	VkMemoryRequirements MemRequirements;
	vkGetBufferMemoryRequirements(Config.Device, Buffer, &MemRequirements);

	VkMemoryAllocateInfo allocInfo{};
	allocInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
	allocInfo.allocationSize = MemRequirements.size;
	allocInfo.memoryTypeIndex = FindMemoryType(MemRequirements.memoryTypeBits, Properties);

	if (vkAllocateMemory(Config.Device, &allocInfo, nullptr, &BufferMemory) != VK_SUCCESS)
	{
		throw std::runtime_error("failed to allocate buffer memory!");
	}

	vkBindBufferMemory(Config.Device, Buffer, BufferMemory, 0);
}

void VulkanBuffers::CreateVertexBuffer(const std::vector<Vertex>& InVertices)
{
	VkDeviceSize BufferSize = sizeof(InVertices[0]) * InVertices.size();

	VkBuffer	   StagingBuffer;
	VkDeviceMemory StagingBufferMemory;
	CreateBuffer(
		BufferSize,
		VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
		VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,
		StagingBuffer,
		StagingBufferMemory
	);
	Log::Info("Starting buffer copy...");

	void* Data;
	vkMapMemory(Config.Device, StagingBufferMemory, 0, BufferSize, 0, &Data);
	memcpy(Data, InVertices.data(), (size_t)BufferSize);
	vkUnmapMemory(Config.Device, StagingBufferMemory);

	Log::Info("Finished buffer copy...");
	CreateBuffer(
		BufferSize,
		VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_VERTEX_BUFFER_BIT,
		VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT,
		VertexBuffer,
		VertexBufferMemory
	);

	CopyBuffer(Config.GraphicsQueue, StagingBuffer, VertexBuffer, BufferSize);

	vkDestroyBuffer(Config.Device, StagingBuffer, nullptr);
	vkFreeMemory(Config.Device, StagingBufferMemory, nullptr);
}

void VulkanBuffers::CreateIndexBuffer(const std::vector<uint16_t>& InIndices)
{
	VkDeviceSize BufferSize = sizeof(InIndices[0]) * InIndices.size();

	VkBuffer	   StagingBuffer;
	VkDeviceMemory StagingBufferMemory;
	CreateBuffer(
		BufferSize,
		VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
		VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,
		StagingBuffer,
		StagingBufferMemory
	);
	Log::Info("Starting buffer copy...");

	void* Data;
	vkMapMemory(Config.Device, StagingBufferMemory, 0, BufferSize, 0, &Data);
	memcpy(Data, InIndices.data(), (size_t)BufferSize);
	vkUnmapMemory(Config.Device, StagingBufferMemory);

	Log::Info("Finished buffer copy...");
	CreateBuffer(
		BufferSize,
		VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_INDEX_BUFFER_BIT,
		VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT,
		IndexBuffer,
		IndexBufferMemory
	);

	CopyBuffer(Config.GraphicsQueue, StagingBuffer, IndexBuffer, BufferSize);

	vkDestroyBuffer(Config.Device, StagingBuffer, nullptr);
	vkFreeMemory(Config.Device, StagingBufferMemory, nullptr);
}

void VulkanBuffers::CreateUniformBuffers(const uint32_t MAX_FRAMES_IN_FLIGHT)
{
	VkDeviceSize BufferSize = sizeof(UniformBufferObject);

	UniformBuffers.resize(MAX_FRAMES_IN_FLIGHT);
	UniformBuffersMemory.resize(MAX_FRAMES_IN_FLIGHT);
	UniformBuffersMapped.resize(MAX_FRAMES_IN_FLIGHT);

	for (size_t i = 0; i < MAX_FRAMES_IN_FLIGHT; i++)
	{
		CreateBuffer(
			BufferSize,
			VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT,
			VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,
			UniformBuffers[i],
			UniformBuffersMemory[i]
		);
		vkMapMemory(
			Config.Device,
			UniformBuffersMemory[i],
			0,
			BufferSize,
			0,
			&UniformBuffersMapped[i]
		);
	}
}

uint32_t VulkanBuffers::FindMemoryType(uint32_t TypeFilter, VkMemoryPropertyFlags Properties)
{
	VkPhysicalDeviceMemoryProperties MemoryProperties;
	vkGetPhysicalDeviceMemoryProperties(Config.PhysicalDevice, &MemoryProperties);

	for (uint32_t i = 0; i < MemoryProperties.memoryTypeCount; i++)
	{
		if ((TypeFilter & (1 << i))
			&& (MemoryProperties.memoryTypes[i].propertyFlags & Properties) == Properties)
		{
			return i;
		}
	}

	Log::Error("Failed to find suitable memory type!");
	return 0;
}
