/*
 * 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 "ManagedBuffer.hpp"
#include "MemoryAllocation.hpp"
#include "Scene/Vertex.hpp"
#include "Scene/Geometry.hpp"
#include "Scene/Material.hpp"
#include "Scene/UniformBuffer.hpp"
#include "Scene/Camera.hpp"
#include "Math/ByteSize.hpp"
#include "Vulkan/Context.hpp"
#include "Vulkan/Image.hpp"
#include "Vulkan/Scene/VulkanShader.hpp"
#include "Vulkan/Scene/VulkanGeometry.hpp"
#include "Vulkan/Scene/VulkanNode.hpp"
#include "Vulkan/Scene/VulkanTexture.hpp"
#include "Vulkan/Scene/VulkanCamera.hpp"
#include "Vulkan/Scene/VulkanUniformBuffer.hpp"

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

	struct ResourceManager::FrameResources
	{
		vk::Device device;
		vk::CommandPool cmdPool;
		vk::CommandBuffer cmdBuffer;
		vk::Semaphore semaphore;
		
		FrameResources(Device* device)
		{
			this->device = device->device;
			cmdPool = this->device.createCommandPool({ {}, device->queueIndices.transfer });
			cmdBuffer = this->device.allocateCommandBuffers({ cmdPool, vk::CommandBufferLevel::ePrimary, 1 })[0];
			semaphore = this->device.createSemaphore({});
		}
		
		~FrameResources()
		{
			device.freeCommandBuffers(cmdPool, 1, &cmdBuffer);
			device.destroyCommandPool(cmdPool);
			device.destroy(semaphore);
		}
	};
	
	
	ResourceManager::ResourceManager()
	{
		static_assert(sizeof(DescriptorSetLayoutBinding) == sizeof(vk::DescriptorSetLayoutBinding));
	}

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

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

		uniformBufferAlignment = context->device->properties.limits.minUniformBufferOffsetAlignment;
		
		frameResources = Array<FrameResources>(static_cast<size_t>(buffers), context->device.get());

		memPool.Init(context->device.get(), 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();
		transferQueue = nullptr;
		geometries.clear();
		nodes.clear();
		textures.clear();
		cameras.clear();
		uniforms.clear();
		device.destroyDescriptorPool(descriptorPool);
		for (auto& sampler : samplerCache)
		{
			device.destroy(sampler.second);
		}
		samplerCache.clear();
		for(auto& layout : descriptorSetLayoutCache)
		{
			device.destroy(layout.second);
		}
		descriptorSetLayoutCache.clear();
		shaders.clear();
		frameResources = {};
		memPool.Close();
		device = nullptr;
	}

	vk::CommandBuffer& ResourceManager::GetCmdBuffer() { return frameResources[currentBuffer].cmdBuffer; }
	
	void ResourceManager::StartFrame(uint64_t frameId)
	{
		currentBuffer = frameId;
		memPool.StartFrame(frameId);
		device.resetCommandPool(frameResources[currentBuffer].cmdPool, {});
		GetCmdBuffer().begin({ vk::CommandBufferUsageFlagBits::eOneTimeSubmit });
	}

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

	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());
			ManagedBuffer::Ptr indexBuffer;
			if (geometry->GetIndexCount())
				indexBuffer = CreateDeviceOnlyBufferWithData(Utils::EnumAsInt(geometry->indexType) * geometry->GetIndexCount(), vk::BufferUsageFlagBits::eIndexBuffer, geometry->GetIndices());
			VulkanGeometry* vkGeo = new VulkanGeometry(geometry, vertexBuffer, indexBuffer);
			geometries.emplace_back(vkGeo);
			geometry->renderGeo = vkGeo;
			return vkGeo;
		}
		return dynamic_cast<VulkanGeometry*>(geometry->renderGeo);
	}

	void ResourceManager::PrepareMaterial(Scene::Material* material)
	{
		if (material->texture && !material->texture->renderTexture)
		{
			PrepareTexture(material->texture);
		}
	}

	VulkanNode* ResourceManager::PrepareNode(Scene::Node* node)
	{
		const std::unique_lock lock(mutex);
		if (!node->renderNode)
		{
			UniformBuffer* uBuffer = new UniformBuffer();
			ManagedBuffer::Ptr 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 = memPool.CreateBuffer(imgs * allocSize, vk::BufferUsageFlagBits::eUniformBuffer, vk::MemoryPropertyFlagBits::eHostCoherent | vk::MemoryPropertyFlagBits::eHostVisible);
				buffer->Map();
			}
			else
			{
				vkNode = new VulkanNode();
				buffer = CreateDeviceOnlyBufferWithData(Scene::Node::SIZE, vk::BufferUsageFlagBits::eUniformBuffer, &node->worldMat);
			}

			uBuffer->Init(std::move(buffer), frameSize, allocSize, GetDescriptorLayoutSet(Scene::Node::DESCRIPTOR_SET_LAYOUT_BINDING), Scene::Node::DESCRIPTOR_SET_LAYOUT_BINDING, 0);
			vkNode->Init(node, uBuffer);
			node->renderNode = vkNode;
			nodes.emplace_back(vkNode);
		}
		return static_cast<VulkanNode*>(node->renderNode);
	}

	VulkanCamera* 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::Ptr buffer = memPool.CreateBuffer(imgs * allocSize, vk::BufferUsageFlagBits::eUniformBuffer, vk::MemoryPropertyFlagBits::eHostCoherent | vk::MemoryPropertyFlagBits::eHostVisible);
			buffer->Map();
			UniformBuffer* uBuffer = new UniformBuffer();
			uBuffer->Init(std::move(buffer), allocSize, allocSize, GetDescriptorLayoutSet(Scene::Camera::DESCRIPTOR_SET_LAYOUT_BINDING), Scene::Camera::DESCRIPTOR_SET_LAYOUT_BINDING, 1);
			VulkanCamera* vkCam = new VulkanCamera();
			vkCam->Init(camera, uBuffer);
			cameras.emplace_back(vkCam);
			camera->renderCamera = vkCam;
		}
		return static_cast<VulkanCamera*>(camera->renderCamera);
	}

	UniformBuffer* ResourceManager::CreateUniformBuffer(const DescriptorSetLayoutBinding& binding, size_t size, void* data, uint32_t setId)
	{
		const std::unique_lock lock(mutex);
		const vk::DeviceSize allocSize = Utils::Align(size, uniformBufferAlignment);

		auto buffer = CreateDeviceOnlyBufferWithData(allocSize, vk::BufferUsageFlagBits::eUniformBuffer, data);

		UniformBuffer* uBuffer = new UniformBuffer();
		uBuffer->Init(std::move(buffer), 0, allocSize, GetDescriptorLayoutSet(binding), binding, setId);
		return uBuffer;
	}

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

	MemoryPool::ManagedBufferPtr ResourceManager::CreateSharedMemoryBuffer(const size_t size)
	{
		const std::unique_lock lock(mutex);
		return memPool.CreateSharedMemoryBuffer(size);
	}

	void ResourceManager::RemoveShader(VulkanShader* shader)
	{
		if (!transferQueue) return;
		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::CopyDataToImage(vk::DeviceSize size, void* data, Image* image)
	{
		ManagedBuffer::Ptr uploadBuffer = memPool.CreateBuffer(size, vk::BufferUsageFlagBits::eTransferSrc, vk::MemoryPropertyFlagBits::eHostCoherent | vk::MemoryPropertyFlagBits::eHostVisible);
		uploadBuffer->Copy(data, size, 0);

		image->SetLayout(GetCmdBuffer(), vk::ImageAspectFlagBits::eColor, vk::ImageLayout::eTransferDstOptimal);

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

		image->SetLayout(GetCmdBuffer(), vk::ImageAspectFlagBits::eColor, vk::ImageLayout::eShaderReadOnlyOptimal, vk::ImageLayout::eTransferDstOptimal);

		// TODO set access masks for mip and array layers
		//GetCmdBuffer().pipelineBarrier(vk::PipelineStageFlagBits::eTopOfPipe, vk::PipelineStageFlagBits::eTransfer, {}, 0, nullptr, 0, nullptr, 1, &barrier );
	}

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

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

	VulkanTexture* ResourceManager::PrepareTexture(Scene::Texture* texture)
	{
		const std::unique_lock lock(mutex);
		if (texture->renderTexture) return static_cast<VulkanTexture*>(texture->renderTexture);
		VulkanTexture* vkTexture;
		if (texture->updateFrequency == Scene::UpdateFrequency::Never)
			vkTexture = new VulkanTexture();
		else
			vkTexture = new VulkanTextureDynamic();

		vkTexture->Init(this, texture, GetDescriptorLayoutSet(Scene::Texture::DESCRIPTOR_SET_LAYOUT_BINDING), Scene::Texture::DESCRIPTOR_SET_LAYOUT_BINDING);

		textures.emplace_back(vkTexture);

		return vkTexture;
	}

	VulkanUniformBuffer* ResourceManager::PrepareUniformBuffer(Scene::UniformBuffer* buffer)
	{
		const std::unique_lock lock(mutex);
		if (buffer->renderBuffer) return static_cast<VulkanUniformBuffer*>(buffer->renderBuffer);

		VulkanUniformBuffer* vkBuffer;
		ManagedBuffer::Ptr mBuffer;
		const vk::DeviceSize allocSize = Utils::Align(buffer->size, uniformBufferAlignment);
		vk::DeviceSize frameSize = 0;
		if (buffer->GetUpdateFrequency() != Scene::UpdateFrequency::Never)
		{
			frameSize = allocSize;
			vkBuffer = new VulkanUniformBufferDynamic();
			const uint32_t imgs = context->swapChain.GetImageCount();
			mBuffer = memPool.CreateBuffer(imgs * allocSize, vk::BufferUsageFlagBits::eUniformBuffer, vk::MemoryPropertyFlagBits::eHostCoherent | vk::MemoryPropertyFlagBits::eHostVisible);
			mBuffer->Map();
		}
		else
		{
			vkBuffer = new VulkanUniformBuffer();
			mBuffer = CreateDeviceOnlyBufferWithData(Scene::Node::SIZE, vk::BufferUsageFlagBits::eUniformBuffer, buffer->data);
			buffer->updated = false;
		}


		UniformBuffer* uBuffer = new UniformBuffer();
		const uint64_t s = mBuffer->size;
		uBuffer->Init(std::move(mBuffer), 0, s, GetDescriptorLayoutSet(buffer->binding), buffer->binding, buffer->setId);

		vkBuffer->Init(buffer, uBuffer);
		uniforms.emplace_back(vkBuffer);
		return vkBuffer;
	}

	vk::Sampler ResourceManager::CreateSampler(const vk::SamplerCreateInfo& samplerConfig)
	{
		auto& sampler = samplerCache[samplerConfig];
		if (!sampler) sampler = device.createSampler(samplerConfig);
		return sampler;
	}
}