/*
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/.
 */

#include "ResourceManager.hpp"
#include "Scene/Vertex.hpp"
#include "Scene/Geometry.hpp"
#include "Scene/Material.hpp"
#include "Vulkan/Context.hpp"
#include "Vulkan/Scene/VulkanShader.hpp"
#include "Vulkan/Scene/VulkanGeometry.hpp"
#include "Vulkan/Scene/VulkanNode.hpp"
#include "Vulkan/Scene/VulkanTexture.hpp"
#include "Vulkan/Image.hpp"
#include "Vulkan/Scene/VulkanCamera.hpp"

namespace openVulkanoCpp::Vulkan
{
	ResourceManager* ResourceManager::INSTANCE;

	ResourceManager::ResourceManager()
	{
		static_assert(sizeof(DescriptorSetLayoutBinding) == sizeof(vk::DescriptorSetLayoutBinding));
		freeFunction = [this](ManagedBuffer* buffer) { this->FreeBuffer(buffer); };
	}

	ResourceManager::~ResourceManager() noexcept
	{
		if (device) ResourceManager::Close();
	}

	void ResourceManager::Init(Context* context, int buffers)
	{
		this->context = context;
		this->device = context->device->device;
		this->buffers = buffers;

		uniformBufferAlignment = context->device->properties.limits.minUniformBufferOffsetAlignment;

		cmdPools = new vk::CommandPool[buffers];
		cmdBuffers = new vk::CommandBuffer[buffers];
		semaphores = new vk::Semaphore[buffers];
		for (int i = 0; i < buffers; i++)
		{
			cmdPools[i] = this->device.createCommandPool({ {}, context->device->queueIndices.transfer });
			cmdBuffers[i] = this->device.allocateCommandBuffers({ cmdPools[i], vk::CommandBufferLevel::ePrimary, 1 })[0];
			semaphores[i] = this->device.createSemaphore({});
		}
		toFree.resize(buffers);

		transferQueue = this->device.getQueue(context->device->queueIndices.transfer, 0);

		// Setup descriptor pool
		constexpr vk::DescriptorPoolSize sizeInfo[] = {
				{ vk::DescriptorType::eSampler, 100000 },
				{ vk::DescriptorType::eCombinedImageSampler, 100000 },
				{ vk::DescriptorType::eSampledImage, 100000 },
				{ vk::DescriptorType::eStorageImage, 100000 },
				{ vk::DescriptorType::eUniformTexelBuffer, 100000 },
				{ vk::DescriptorType::eStorageTexelBuffer, 100000 },
				{ vk::DescriptorType::eUniformBuffer, 100000 },
				{ vk::DescriptorType::eStorageBuffer, 100000 },
				{ vk::DescriptorType::eUniformBufferDynamic, 100000 },
				{ vk::DescriptorType::eStorageBufferDynamic, 100000 },
				{ vk::DescriptorType::eInputAttachment, 100000 }
		};
		const vk::DescriptorPoolCreateInfo poolCreateInfo(vk::DescriptorPoolCreateFlagBits::eFreeDescriptorSet, 100000, std::size(sizeInfo), sizeInfo);
		descriptorPool = device.createDescriptorPool(poolCreateInfo);

		INSTANCE = this;
	}

	void ResourceManager::Close()
	{
		transferQueue.waitIdle();
		for (int i = 0; i < buffers; i++)
		{
			device.freeCommandBuffers(cmdPools[i], 1, &cmdBuffers[i]);
			device.destroyCommandPool(cmdPools[i]);
		}
		device.destroyDescriptorPool(descriptorPool);
		shaders.clear();
		cmdBuffers = nullptr;
		cmdPools = nullptr;
		device = nullptr;
	}

	void ResourceManager::StartFrame(uint64_t frameId)
	{
		currentBuffer = frameId;
		FreeBuffers();
		device.resetCommandPool(cmdPools[currentBuffer], {});
		cmdBuffers[currentBuffer].begin({ vk::CommandBufferUsageFlagBits::eOneTimeSubmit });
	}

	vk::Semaphore ResourceManager::EndFrame()
	{
		cmdBuffers[currentBuffer].end();
		vk::SubmitInfo si = { 0, nullptr, nullptr, 1, &cmdBuffers[currentBuffer], 1, &semaphores[currentBuffer] };
		transferQueue.submit(1, &si, vk::Fence());
		return semaphores[currentBuffer];
	}

	void ResourceManager::Resize()
	{
		for (auto& shader : shaders)
		{
			shader->Resize();
		}
	}

	VulkanGeometry* ResourceManager::PrepareGeometry(Scene::Geometry* geometry)
	{
		const std::unique_lock lock(mutex);
		if(!geometry->renderGeo)
		{
			ManagedBuffer::Ptr vertexBuffer(
					CreateDeviceOnlyBufferWithData(sizeof(Vertex) * geometry->GetVertexCount(), vk::BufferUsageFlagBits::eVertexBuffer, geometry->GetVertices()),
					freeFunction);
			ManagedBuffer::Ptr indexBuffer(
					CreateDeviceOnlyBufferWithData(Utils::EnumAsInt(geometry->indexType) * geometry->GetIndexCount(), vk::BufferUsageFlagBits::eIndexBuffer, geometry->GetIndices()),
					freeFunction);
			VulkanGeometry* vkGeo = new VulkanGeometry(geometry, vertexBuffer, indexBuffer);
			geometry->renderGeo = vkGeo;
			return vkGeo;
		}
		return dynamic_cast<VulkanGeometry*>(geometry->renderGeo);
	}

	void ResourceManager::PrepareMaterial(Scene::Material* material)
	{
		const std::unique_lock lock(mutex);
		if(!material->shader->renderShader)
		{
			material->shader->renderShader = CreateShader(material->shader);
		}
		if (material->texture && !material->texture->renderTexture)
		{
			material->texture->renderTexture = PrepareTexture(material->texture);
		}
	}

	void ResourceManager::PrepareNode(Scene::Node* node)
	{
		const std::unique_lock lock(mutex);
		if (!node->renderNode)
		{
			UniformBuffer* uBuffer = new UniformBuffer();
			ManagedBuffer* buffer;
			VulkanNode* vkNode;
			const vk::DeviceSize allocSize = Utils::Align(sizeof(Math::Matrix4f), uniformBufferAlignment);
			vk::DeviceSize frameSize = 0;
			if (node->GetUpdateFrequency() != Scene::UpdateFrequency::Never)
			{
				frameSize = allocSize;
				vkNode = new VulkanNodeDynamic();
				const uint32_t imgs = context->swapChain.GetImageCount();
				buffer = CreateBuffer(imgs * allocSize, vk::BufferUsageFlagBits::eUniformBuffer, vk::MemoryPropertyFlagBits::eHostCoherent | vk::MemoryPropertyFlagBits::eHostVisible);
				buffer->Map();
			}
			else
			{
				vkNode = new VulkanNode();
				buffer = CreateDeviceOnlyBufferWithData(sizeof(Math::Matrix4f), vk::BufferUsageFlagBits::eUniformBuffer, &node->worldMat);
			}

			uBuffer->Init(buffer, frameSize, allocSize, GetDescriptorLayoutSet(NODE_LAYOUT_BINDING), NODE_LAYOUT_BINDING, 0);
			vkNode->Init(node, uBuffer);
			node->renderNode = vkNode;
		}
	}

	void ResourceManager::PrepareCamera(Scene::Camera* camera)
	{
		const std::unique_lock lock(mutex);
		if (!camera->renderCamera)
		{
			const vk::DeviceSize allocSize = Utils::Align(Scene::Camera::SIZE, uniformBufferAlignment);
			const uint32_t imgs = context->swapChain.GetImageCount();
			ManagedBuffer* buffer = CreateBuffer(imgs * allocSize, vk::BufferUsageFlagBits::eUniformBuffer, vk::MemoryPropertyFlagBits::eHostCoherent | vk::MemoryPropertyFlagBits::eHostVisible);
			buffer->Map();
			UniformBuffer* uBuffer = new UniformBuffer();
			uBuffer->Init(buffer, allocSize, allocSize, GetDescriptorLayoutSet(CAM_LAYOUT_BINDING), CAM_LAYOUT_BINDING, 1);
			VulkanCamera* vkCam = new VulkanCamera();
			vkCam->Init(camera, uBuffer);
			camera->renderCamera = vkCam;
		}
	}

	vk::DescriptorSetLayout* ResourceManager::GetDescriptorLayoutSet(const DescriptorSetLayoutBinding& descriptorSetLayoutBinding)
	{
		auto& layout = descriptorSetLayoutCache[descriptorSetLayoutBinding];
		if (!layout)
		{
			vk::DescriptorSetLayoutCreateInfo createInfo({}, 1, &reinterpret_cast<const vk::DescriptorSetLayoutBinding&>(descriptorSetLayoutBinding));
			layout = device.createDescriptorSetLayout(createInfo);
		}
		return &layout;
	}

	ManagedBuffer* ResourceManager::CreateSharedMemoryBuffer(const size_t size)
	{
		const std::unique_lock lock(mutex);
		if (!recycleBuffers.empty())
		{
			for(auto buff : recycleBuffers)
			{
				if (buff->size == size)
				{
					Logger::DATA->info("Recycle Buffer");
					Utils::Remove(recycleBuffers, buff);
					return buff;
				}
			}
		}
		ManagedBuffer* buffer = CreateBuffer(size, vk::BufferUsageFlagBits::eVertexBuffer, vk::MemoryPropertyFlagBits::eHostCoherent/* | vk::MemoryPropertyFlagBits::eDeviceLocal*/);
		return buffer;
	}

	void ResourceManager::RemoveShader(VulkanShader* shader)
	{
		const std::unique_lock lock(mutex);
		std::vector<std::unique_ptr<VulkanShader>>::iterator object =
				find_if(shaders.begin(), shaders.end(),
						[&](auto& obj){ return obj.get() == shader; }
				);
		shaders.erase(object);
	}

	void ResourceManager::FreeBuffer(ManagedBuffer* buffer)
	{
		if (buffer)
			toFree[currentBuffer].push_back(buffer);
	}

	void ResourceManager::DoFreeBuffer(ManagedBuffer* buffer)
	{
		if (buffer->IsLast())
		{
			device.destroyBuffer(buffer->buffer);
			buffer->allocation->used -= buffer->size;
		}
		else
		{
			recycleBuffers.push_back(buffer);
		}
	}

	void ResourceManager::FreeBuffers()
	{
		for (auto& i : toFree[currentBuffer])
		{
			DoFreeBuffer(i);
		}
		toFree[currentBuffer].clear();
	}

	void ResourceManager::CopyDataToImage(vk::DeviceSize size, void* data, Image* image)
	{
		ManagedBuffer* uploadBuffer = CreateBuffer(size, vk::BufferUsageFlagBits::eTransferSrc, vk::MemoryPropertyFlagBits::eHostCoherent | vk::MemoryPropertyFlagBits::eHostVisible);
		uploadBuffer->Copy(data, size, 0);

		vk::BufferImageCopy region(0, 0, 0, { vk::ImageAspectFlagBits::eColor, 0, 0, 1 }, { 0, 0, 0 }, image->extent);
		cmdBuffers[currentBuffer].copyBufferToImage(uploadBuffer->buffer, image->image, vk::ImageLayout::eTransferDstOptimal, 1, &region);

		FreeBuffer(uploadBuffer);
	}

	ManagedBuffer* ResourceManager::CreateDeviceOnlyBufferWithData(vk::DeviceSize size, vk::BufferUsageFlagBits usage, void* data)
	{
		ManagedBuffer* target = CreateBuffer(size, usage | vk::BufferUsageFlagBits::eTransferDst, vk::MemoryPropertyFlagBits::eDeviceLocal);
		ManagedBuffer* uploadBuffer = CreateBuffer(size, vk::BufferUsageFlagBits::eTransferSrc, vk::MemoryPropertyFlagBits::eHostCoherent | vk::MemoryPropertyFlagBits::eHostVisible);
		uploadBuffer->Copy(data, size, 0);
		RecordCopy(uploadBuffer->buffer, target->buffer, size);
		FreeBuffer(uploadBuffer);
		return target;
	}

	ManagedBuffer* ResourceManager::CreateBuffer(vk::DeviceSize size, const vk::BufferUsageFlags& usage, const vk::MemoryPropertyFlags& properties)
	{
		size = Utils::Align(size, 16);
		const vk::BufferCreateInfo bufferCreateInfo = { {}, size, usage, vk::SharingMode::eExclusive };
		vk::Buffer buffer = device.createBuffer(bufferCreateInfo);
		const vk::MemoryRequirements memoryRequirements = device.getBufferMemoryRequirements(buffer);
		uint32_t memtype = context->device->GetMemoryType(memoryRequirements.memoryTypeBits, properties);
		if (memoryRequirements.size != size) Logger::DATA->warn("Memory Requirement Size ({0}) != Size ({1})", memoryRequirements.size, size);
		MemoryAllocation* allocation = GetFreeMemoryAllocation(memoryRequirements.size, memtype);
		uint32_t offset = allocation->used;
		device.bindBufferMemory(buffer, allocation->memory, offset);
		allocation->used += memoryRequirements.size;
		return new ManagedBuffer{ allocation, offset, size, buffer, usage, properties, nullptr };
	}

	MemoryAllocation* ResourceManager::CreateMemoryAllocation(size_t size, uint32_t type, bool addToCache)
	{
		MemoryAllocation* alloc = new MemoryAllocation(size, type, device);
		const vk::MemoryAllocateInfo allocInfo = { size, type };
		alloc->memory = device.allocateMemory(allocInfo);
		if (addToCache) allocations.push_back(alloc);
		return alloc;
	}

	MemoryAllocation* ResourceManager::GetFreeMemoryAllocation(size_t size, uint32_t type, bool createIfAllFull)
	{
		MemoryAllocation* alloc = nullptr;
		for (MemoryAllocation* allocation : allocations)
		{
			if (allocation->type == type && allocation->FreeSpace() >= size)
			{
				alloc = allocation;
				break;
			}
		}
		if(!alloc && createIfAllFull) alloc = CreateMemoryAllocation(64 * 1024 * 1024, type, true);
		if(alloc) lastAllocation = alloc;
		return alloc;
	}

	VulkanShader* ResourceManager::CreateShader(Scene::Shader* shader)
	{
		const std::unique_lock lock(mutex);
		VulkanShader* vkShader = new VulkanShader();
		vkShader->Init(context, shader, this);
		shaders.emplace_back(vkShader);
		return vkShader;
	}

	VulkanTexture* ResourceManager::PrepareTexture(Scene::Texture* texture)
	{
		VulkanTexture* vkTexture = new VulkanTexture();

		vkTexture->Init(this, texture);
		//vkTexture->

		return vkTexture;
	}
}