/*
 * 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 "Device.hpp"
#include "Base/Utils.hpp"
#include "Shader/ShaderRegistry.hpp"

namespace OpenVulkano::Vulkan
{
	namespace
	{
		class DeviceQueueCreateInfoBuilder
		{
			std::vector<vk::DeviceQueueCreateInfo> createInfos;
			std::vector<std::vector<float>> prioritiesVector;
		public:
			DeviceQueueCreateInfoBuilder() = default;
			~DeviceQueueCreateInfoBuilder() = default;

			void AddQueueFamily(const uint32_t queueFamilyIndex, const std::vector<float>& priorities)
			{
				prioritiesVector.push_back(priorities);
				createInfos.emplace_back(vk::DeviceQueueCreateFlags(), queueFamilyIndex, priorities.size(), prioritiesVector[prioritiesVector.size()-1].data());
			}

			void AddQueueFamily(uint32_t queueFamilyIndex, uint32_t count = 1)
			{
				std::vector<float> priorities;
				priorities.resize(count);
				std::fill(priorities.begin(), priorities.end(), 0.0f);
				AddQueueFamily(queueFamilyIndex, priorities);
			}

			std::vector<vk::DeviceQueueCreateInfo>& GetDeviceQueueCreateInfos()
			{
				return createInfos;
			}
		};
	}

	Device::Device(vk::PhysicalDevice& physicalDevice)
	{
		this->physicalDevice = physicalDevice;
		useDebugMarkers = false;
		QueryDevice();
	}

	void Device::PrepareDevice(const vk::ArrayProxy<const std::string>& requestedExtensions, const vk::SurfaceKHR& surface)
	{
		queueIndices.graphics = FindBestQueue(vk::QueueFlagBits::eGraphics, surface); // Make sure that the graphics queue supports the surface
		BuildDevice(requestedExtensions);
		//TODO setup debug marker
		pipelineCache = device.createPipelineCache(vk::PipelineCacheCreateInfo());

		graphicsQueue = device.getQueue(queueIndices.graphics, 0);
		graphicsCommandPool = device.createCommandPool({ vk::CommandPoolCreateFlagBits::eResetCommandBuffer, queueIndices.graphics, });
	}

	void Device::WaitIdle() const
	{ //TODO wait all queues idle
		graphicsQueue.waitIdle();
		device.waitIdle();
	}

	vk::CommandBuffer Device::CreateCommandBuffer(vk::CommandBufferLevel level) const
	{
		const vk::CommandBufferAllocateInfo cmdBufferAllocInfo(graphicsCommandPool, level, 1);
		return device.allocateCommandBuffers(cmdBufferAllocInfo)[0];
	}

	void Device::FlushCommandBuffer(vk::CommandBuffer& cmdBuffer) const
	{
		graphicsQueue.submit(vk::SubmitInfo{ 0, nullptr, nullptr, 1, &cmdBuffer }, vk::Fence());
		WaitIdle();
	}

	void Device::ExecuteNow(const std::function<void(const vk::CommandBuffer&)>& function) const
	{
		vk::CommandBuffer commandBuffer = CreateCommandBuffer();
		commandBuffer.begin(vk::CommandBufferBeginInfo{ vk::CommandBufferUsageFlagBits::eOneTimeSubmit });
		function(commandBuffer);
		commandBuffer.end();
		FlushCommandBuffer(commandBuffer);
		device.freeCommandBuffers(graphicsCommandPool, commandBuffer);
	}

	vk::ShaderModule Device::CreateShaderModule(const std::string& filename) const
	{
		Array<char> buffer;
		const void* data = nullptr;
		size_t size = 0;
		if (Utils::StartsWith(filename, "Shader/"))
		{
			auto shader = ShaderRegistry::GetInstance().GetShader(filename.substr(7));
			data = shader.first;
			size = shader.second;
		}
		if (!data)
		{
			buffer = Utils::ReadFile(filename);
			size = buffer.Size();
			data = buffer.Data();
		}
		vk::ShaderModuleCreateInfo smci = {{}, size, static_cast<const uint32_t*>(data) };
		return CreateShaderModule(smci);
	}

	vk::ShaderModule Device::CreateShaderModule(vk::ShaderModuleCreateInfo& createInfo) const
	{
		return device.createShaderModule(createInfo);
	}

	void Device::QueryDevice()
	{
		// Query device features
		queueFamilyProperties = physicalDevice.getQueueFamilyProperties();
		properties = physicalDevice.getProperties();
		features = physicalDevice.getFeatures();
		for (auto& ext : physicalDevice.enumerateDeviceExtensionProperties()) { supportedExtensions.insert(ext.extensionName.data()); }

		// Query device memory properties
		memoryProperties = physicalDevice.getMemoryProperties();
		queueIndices.graphics = FindBestQueue(vk::QueueFlagBits::eGraphics);
		queueIndices.compute = FindBestQueue(vk::QueueFlagBits::eCompute);
		queueIndices.transfer = FindBestQueue(vk::QueueFlagBits::eTransfer);
	}

	void Device::BuildDevice(const vk::ArrayProxy<const std::string>& requestedExtensions)
	{
		vk::DeviceCreateInfo deviceCreateInfo;
		deviceCreateInfo.pEnabledFeatures = &features; //TODO add option to disable not needed features


		DeviceQueueCreateInfoBuilder deviceQueueCreateInfoBuilder;
		deviceQueueCreateInfoBuilder.AddQueueFamily(queueIndices.GetGraphics(), queueFamilyProperties[queueIndices.GetGraphics()].queueCount);
		if (queueIndices.GetCompute() != VK_QUEUE_FAMILY_IGNORED)
			deviceQueueCreateInfoBuilder.AddQueueFamily(queueIndices.GetCompute(), queueFamilyProperties[queueIndices.GetCompute()].queueCount);
		if (queueIndices.GetTransfer() != VK_QUEUE_FAMILY_IGNORED)
			deviceQueueCreateInfoBuilder.AddQueueFamily(queueIndices.GetTransfer(), queueFamilyProperties[queueIndices.GetTransfer()].queueCount);
		const std::vector<vk::DeviceQueueCreateInfo> deviceQueueCreateInfos = deviceQueueCreateInfoBuilder.GetDeviceQueueCreateInfos();

		deviceCreateInfo.queueCreateInfoCount = static_cast<uint32_t>(deviceQueueCreateInfos.size());
		deviceCreateInfo.pQueueCreateInfos = deviceQueueCreateInfos.data();

		std::vector<const char*> enabledExtensions;
		for (const auto& extension : requestedExtensions)
		{
			enabledExtensions.push_back(extension.c_str());
		}
#ifdef DEBUG
		if (IsExtensionAvailable({ VK_EXT_DEBUG_MARKER_EXTENSION_NAME }))
				{ // Enable debug marker extension if available
					enabledExtensions.push_back(VK_EXT_DEBUG_MARKER_EXTENSION_NAME);
					useDebugMarkers = true;
				}
#endif
		deviceCreateInfo.enabledExtensionCount = static_cast<uint32_t>(enabledExtensions.size());
		deviceCreateInfo.ppEnabledExtensionNames = enabledExtensions.data();

		device = physicalDevice.createDevice(deviceCreateInfo);
	}

	uint32_t Device::FindBestQueue(const vk::QueueFlags& desiredFlags, const vk::SurfaceKHR& surface) const
	{
		uint32_t best = VK_QUEUE_FAMILY_IGNORED;
		VkQueueFlags bestExtraFlagsCount = VK_QUEUE_FLAG_BITS_MAX_ENUM;
		for (size_t i = 0; i < queueFamilyProperties.size(); i++)
		{
			vk::QueueFlags flags = queueFamilyProperties[i].queueFlags;
			if (!(flags & desiredFlags)) continue; // Skip queue without desired flags
			if (surface && VK_FALSE == physicalDevice.getSurfaceSupportKHR(i, surface)) continue;
			const VkQueueFlags currentExtraFlags = (flags & ~desiredFlags).operator VkQueueFlags();

			if (0 == currentExtraFlags) return i; // return exact match

			if (best == VK_QUEUE_FAMILY_IGNORED || currentExtraFlags < bestExtraFlagsCount)
			{
				best = i;
				bestExtraFlagsCount = currentExtraFlags;
			}
		}
		return best;
	}

	vk::Format Device::GetSupportedDepthFormat(const std::vector<vk::Format>& depthFormats) const
	{
		for (auto& format : depthFormats)
		{
			vk::FormatProperties formatProps;
			formatProps = physicalDevice.getFormatProperties(format);
			if (formatProps.optimalTilingFeatures & vk::FormatFeatureFlagBits::eDepthStencilAttachment)
			{
				return format;
			}
		}

		throw std::runtime_error("No supported depth format");
	}

	bool Device::GetMemoryType(uint32_t typeBits, const vk::MemoryPropertyFlags& properties, uint32_t* typeIndex) const
	{
		for (uint32_t i = 0; i < 32; i++)
		{
			if ((typeBits & 1) == 1)
			{
				if ((memoryProperties.memoryTypes[i].propertyFlags & properties) == properties)
				{
					*typeIndex = i;
					return true;
				}
			}
			typeBits >>= 1;
		}
		return false;
	}

	uint32_t Device::GetMemoryType(uint32_t typeBits, const vk::MemoryPropertyFlags& properties) const
	{
		uint32_t result = 0;
		if (!GetMemoryType(typeBits, properties, &result))
		{
			throw std::runtime_error("Unable to find memory type " + to_string(properties));
		}
		return result;
	}

	void Device::Close()
	{
		device.destroyCommandPool(graphicsCommandPool);
		//TODO fill
	}
}