#include "VulkanVertexBuffer.h"
#include "utilities/Logger.h"
#include <cstdint>
#include <vulkan/vulkan_core.h>
#include <cstring>

void VulkanVertexBuffer::Initialize(FVertexBufferConfig InConfig)
{
	Config = InConfig;
}

void VulkanVertexBuffer::Cleanup(const uint32_t MAX_FRAMES_IN_FLIGHT)
{
	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 VulkanVertexBuffer::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)
	{
		Log::Error("Failed to create buffer!");
	}
	else
	{
		Log::Info("Successfully created buffer.");
	}

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

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

	if (vkAllocateMemory(Config.Device, &AllocateInfo, nullptr, &BufferMemory) != VK_SUCCESS)
	{
		Log::Error("Failed to allocate vertex buffer memory!");
	}
	else
	{
		Log::Info("Successfully allocated vertex buffer memory.");
	}

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

void VulkanVertexBuffer::CopyBuffer(VkBuffer SrcBuffer, VkBuffer DstBuffer, VkDeviceSize Size)
{
	VkCommandBufferAllocateInfo AllocateInfo{};
	AllocateInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
	AllocateInfo.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
	AllocateInfo.commandPool = Config.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);

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

	vkEndCommandBuffer(CommandBuffer);

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

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

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

void VulkanVertexBuffer::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(StagingBuffer, VertexBuffer, BufferSize);

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

void VulkanVertexBuffer::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(StagingBuffer, IndexBuffer, BufferSize);

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

void VulkanVertexBuffer::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 VulkanVertexBuffer::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;
}
