#pragma once

#define CRASH_ON_MULTIPLE_MAPPINGS_TO_SAME_ALLOCATION

#include <vulkan/vulkan.hpp>

namespace openVulkanoCpp
{
	namespace Vulkan
	{
		struct MemoryAllocation
		{
			vk::DeviceMemory memory;
			size_t used;
			const size_t size;
			const uint32_t type;
			const vk::Device device;
			void* mapped;

		private:
			uint32_t mappedCount;
			static constexpr uint32_t CHILD_MAPPED_FLAG = 1u << 31;

		public:
			MemoryAllocation(size_t size, uint32_t type, vk::Device device):
				memory(nullptr), used(0), size(size), type(type), device(device), mapped(nullptr), mappedCount(0)
			{
			}

			size_t FreeSpace() const
			{
				return size - used;
			}

			void HandleChildMappingValidation() const
			{
				if (IsChildMapped())
				{
					Logger::RENDER->error("A single memory allocation should only be mapped once! Mapping a single allocation multiple times might not work or might not work reliable with all driver implementations.");
#ifdef CRASH_ON_MULTIPLE_MAPPINGS_TO_SAME_ALLOCATION
					throw std::runtime_error("A single memory allocation should only be mapped once!");
#endif
				}
			}
			
			void* Map()
			{
				HandleChildMappingValidation();
				if (!mapped)
				{
					mapped = device.mapMemory(memory, 0, size, vk::MemoryMapFlags());
				}
				mappedCount++;
				return mapped;
			}

			void UnMap()
			{
				if (mappedCount > 0)
				{
					mappedCount--;
					if (mappedCount == 0)
					{
						device.unmapMemory(memory);
					}
				}
			}

			bool IsChildMapped() const
			{
				return mappedCount & CHILD_MAPPED_FLAG;
			}

			void* MapChild(size_t offset, vk::DeviceSize size)
			{
				HandleChildMappingValidation();
				mappedCount |= CHILD_MAPPED_FLAG;
				mappedCount++;
				return device.mapMemory(memory, offset, size, vk::MemoryMapFlags());
			}

			void UnMapChild()
			{
				mappedCount &= ~CHILD_MAPPED_FLAG;
				mappedCount--;
				if (mappedCount == 0)
				{
					device.unmapMemory(memory);
				}
			}
		};

		struct ManagedBuffer
		{
			MemoryAllocation* allocation;
			vk::DeviceSize offset, size;
			vk::Buffer buffer;
			vk::BufferUsageFlags usage;
			vk::MemoryPropertyFlags properties;
			void* mapped = nullptr;

			bool IsLast() const
			{
				return (offset + size == allocation->used);
			}

			bool IsMapped() const
			{
				return mapped || allocation->mapped;
			}

			/**
			 * \brief Maps the buffer into the memory of the host.
			 * \tparam T The type of the buffers data.
			 * \param offset The offset from where to map the buffer.
			 * \param size The size to be mapped. VK_WHOLE_SIZE to map the whole buffer.
			 * \pparam longTermMapping If the mapping is intended to be held long term. Short term mappings must be freed before mapping a different region in the same memory allocation for them to work reliable with all drivers.
			 * \return The pointer to the mapped buffer.
			 */
			template <typename T = void>
			T* Map(size_t offset = 0, vk::DeviceSize size = VK_WHOLE_SIZE, bool longTermMapping = true)
			{
				if (!mapped)
				{
					if (allocation->mapped || longTermMapping)
					{
						mapped = static_cast<uint8_t*>(allocation->Map()) + offset + this->offset;
					}
					else
					{
						if (size == VK_WHOLE_SIZE) size = this->size;
						mapped = allocation->MapChild(this->offset + offset, size);
					}
					
					return static_cast<T*>(mapped);
				}
			}

			/**
			 * \brief Un-maps the buffer from the host.
			 */
			void UnMap()
			{
				if (mapped)
				{
					if (allocation->mapped)
					{
						allocation->UnMap();
					}
					else
					{
						allocation->UnMapChild();
					}
					mapped = nullptr;
				}
			}

			void Copy(void* data)
			{
				if (mapped)
				{
					memcpy(mapped, data, size);
				}
				else
				{
					void* dataMapped = Map(0, VK_WHOLE_SIZE, false);
					memcpy(dataMapped, data, size);
					UnMap();
				}
			}

			void Copy(void* data, uint32_t size, uint32_t offset)
			{
				if(mapped) memcpy(static_cast<char*>(mapped) + offset, data, size);
				else
				{
					void* dataMapped = Map(offset, size, false);
					memcpy(dataMapped, data, size);
					UnMap();
				}
			}
		};
	}
}