/*
 * 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/.
 */

#pragma once

#include "vulkan/vulkan.hpp"
#include "Base/ICloseable.hpp"
#include "IShaderOwner.hpp"
#include "ManagedResource.hpp"
#include "Vulkan/Image.hpp"
#include "Scene/Shader/DescriptorInputDescription.hpp"
#include "Scene/Camera.hpp"
#include <mutex>
#include <functional>
#include <map>

namespace openVulkanoCpp
{
	namespace Scene
	{
		class Node;
		class Geometry;
		class Material;
		class Texture;
		class Shader;
	}

	namespace Vulkan
	{
		class Context;
		class VulkanGeometry;
		class VulkanTexture;
		class UniformBuffer;

		class ResourceManager : public ICloseable, public IShaderOwner
		{
			friend UniformBuffer;

			Context* context;
			vk::Device device = nullptr;
			vk::Queue transferQueue = nullptr;
			vk::CommandPool* cmdPools = nullptr;
			vk::CommandBuffer* cmdBuffers = nullptr;
			vk::Semaphore* semaphores = nullptr;
			std::vector<MemoryAllocation*> allocations;
			std::vector<std::unique_ptr<VulkanShader>> shaders;
			MemoryAllocation* lastAllocation = nullptr;
			std::mutex mutex;
			vk::DeviceSize uniformBufferAlignment;
			std::vector<std::vector<ManagedBuffer*>> toFree;
			std::vector<ManagedBuffer*> recycleBuffers;
			std::function<void(ManagedBuffer*)> freeFunction;
			vk::DescriptorPool descriptorPool;
			std::map<DescriptorSetLayoutBinding, vk::DescriptorSetLayout> descriptorSetLayoutCache;

			int buffers = -1, currentBuffer = -1;

		public:
			static ResourceManager* INSTANCE;

			ResourceManager();

			virtual ~ResourceManager() noexcept;

			void Init(Context* context, int buffers = 2);

			void Close() override;

			void StartFrame(uint64_t frameId);

			vk::Semaphore EndFrame();

			void Resize();

			VulkanGeometry* PrepareGeometry(Scene::Geometry* geometry);

			void PrepareMaterial(Scene::Material* material);

			void PrepareNode(Scene::Node* node);

			VulkanTexture* PrepareTexture(Scene::Texture* texture);

			void PrepareCamera(Scene::Camera* camera);

			void RemoveShader(VulkanShader* shader) override;

			void CopyDataToImage(vk::DeviceSize size, void* data, openVulkanoCpp::Vulkan::Image* image);

			ManagedBuffer* CreateSharedMemoryBuffer(size_t size);

			void FreeBuffer(ManagedBuffer* buffer);

			[[nodiscard]] Context* GetContext() const { return context; }

			[[nodiscard]] vk::Device GetDevice() const { return device; }

		protected: // Allocation management
			void DoFreeBuffer(ManagedBuffer* buffer);

			void FreeBuffers();

			ManagedBuffer* CreateDeviceOnlyBufferWithData(vk::DeviceSize size, vk::BufferUsageFlagBits usage, void* data);

			inline void RecordCopy(vk::Buffer src, vk::Buffer dest, vk::DeviceSize size) const
			{
				vk::BufferCopy copyRegion = { 0, 0, size };
				cmdBuffers[currentBuffer].copyBuffer(src, dest, 1, &copyRegion);
			}
			
			ManagedBuffer* CreateBuffer(vk::DeviceSize size, const vk::BufferUsageFlags& usage, const vk::MemoryPropertyFlags& properties);
			
			MemoryAllocation* CreateMemoryAllocation(size_t size, uint32_t type, bool addToCache = true);

			MemoryAllocation* GetFreeMemoryAllocation(size_t size, uint32_t type, bool createIfAllFull = true);

			vk::DescriptorSetLayout* GetDescriptorLayoutSet(const DescriptorSetLayoutBinding& descriptorSetLayoutBinding);

		public:
			VulkanShader* CreateShader(Scene::Shader* shader);
		};
	}
}