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

#define CRASH_ON_MULTIPLE_MAPPINGS_TO_SAME_ALLOCATION

#include <memory>
#include <vulkan/vulkan.hpp>
#include <functional>

namespace openVulkanoCpp::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)
		{
		}

		~MemoryAllocation()
		{
			if (device && memory)
				device.free(memory);
		}

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

		void HandleChildMappingValidation() const;

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

		[[nodiscard]] 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
	{
		using Ptr = std::unique_ptr<ManagedBuffer, std::function<void(ManagedBuffer*)>>;

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

		~ManagedBuffer()
		{
			allocation->device.destroy(buffer);
		}

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

		[[nodiscard]] 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();
			}
		}
	};
}