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

// Workaround for libc++
#ifndef __cpp_lib_three_way_comparison
#define __cpp_lib_three_way_comparison 201907
#endif

#include "IShaderOwner.hpp"
#include "MemoryPool.hpp"
#include "Base/Wrapper.hpp"
#include "Base/Event.hpp"
#include "Base/Render/IResourceManager.hpp"
#include "Vulkan/Image.hpp"
#include "Scene/Shader/DescriptorInputDescription.hpp"
#include <vulkan/vulkan.hpp>
#include <map>
#include <memory>
#include <mutex>

namespace OpenVulkano
{
	namespace Scene
	{
		class Node;
		class Geometry;
		class Material;
		class Texture;
		class Shader;
		class Camera;
		class UniformBuffer;
		class VertexBuffer;
	}

	namespace Vulkan
	{
		class Context;
		class VulkanGeometry;
		class VulkanTexture;
		class VulkanCamera;
		class VulkanNode;
		class VulkanShader;
		class VulkanUniformBuffer;
		class VulkanVertexBuffer;
		class UniformBuffer;
		class ManagedBuffer;
		class MemoryAllocation;

		class ResourceManager final : public IShaderOwner, public IResourceManager
		{
			friend UniformBuffer;
			friend VulkanTexture;

			struct FrameResources;

			Context* context;
			vk::Device device = nullptr;
			vk::Queue transferQueue = nullptr;
			Array<FrameResources> frameResources;
			MemoryPool memPool;
			std::vector<Unique<VulkanShader>> shaders;
			std::vector<Unique<VulkanGeometry>> geometries;
			std::vector<Unique<VulkanVertexBuffer>> vertexBuffers;
			std::vector<Unique<VulkanNode>> nodes;
			std::vector<Unique<VulkanTexture>> textures;
			std::vector<Unique<VulkanCamera>> cameras;
			std::vector<Unique<VulkanUniformBuffer>> uniforms;
			std::mutex mutex;
			vk::DeviceSize uniformBufferAlignment;
			vk::DescriptorPool descriptorPool;
			std::map<DescriptorSetLayoutBinding, vk::DescriptorSetLayout> descriptorSetLayoutCache;
			std::map<vk::SamplerCreateInfo, vk::Sampler> samplerCache;

			int currentBuffer = -1;

			vk::CommandBuffer& GetCmdBuffer();

		public:
			static ResourceManager* INSTANCE;

			ResourceManager();

			virtual ~ResourceManager() noexcept;

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

			void Close();

			void StartFrame(uint64_t frameId);

			vk::Semaphore EndFrame();

			void Resize();

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

			void PrepareMaterial(Scene::Material* material);

			VulkanUniformBuffer* PrepareUniformBuffer(Scene::UniformBuffer* buffer);

			VulkanVertexBuffer* PrepareVertexBuffer(Scene::VertexBuffer* vbo);

			VulkanNode* PrepareNode(Scene::Node* node);

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

			VulkanCamera* PrepareCamera(Scene::Camera* camera);

			void RemoveShader(VulkanShader* shader) override;

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

			MemoryPool::ManagedBufferPtr CreateSharedMemoryBuffer(size_t size);

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

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

			UniformBuffer* CreateUniformBuffer(const DescriptorSetLayoutBinding& binding, size_t size, void* data, uint32_t setId = 2, bool hostVis = false);

			vk::Sampler CreateSampler(const vk::SamplerCreateInfo& samplerConfig);

			[[nodiscard]] Unique<Scene::Texture> CreateSharedMemoryTexture(const Math::Vector3ui& resolution, DataFormat format) override;

		protected: // Allocation management
			MemoryPool::ManagedBufferPtr CreateDeviceOnlyBufferWithData(vk::DeviceSize size, vk::BufferUsageFlagBits usage, const void* data);

			inline void RecordCopy(vk::Buffer src, vk::Buffer dest, vk::DeviceSize size)
			{
				vk::BufferCopy copyRegion = { 0, 0, size };
				GetCmdBuffer().copyBuffer(src, dest, 1, &copyRegion);
			}

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

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

			Event<ResourceManager*> OnShutdown;
		};
	}
}