/*
 * 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 "SwapChain.hpp"
#include "Base/Logger.hpp"
#include "Base/Utils.hpp"
#include "Base/UI/IVulkanWindow.hpp"
#include "Base/EngineConfiguration.hpp"
#include "Scene/DataFormat.hpp"

namespace OpenVulkano::Vulkan
{
	void SwapChain::Init(Device* device, vk::SurfaceKHR surface, IVulkanWindow* window)
	{
		if (!device) throw std::runtime_error("The device must not be null");
		if (!window) throw std::runtime_error("The window must not be null");
		this->device = device;
		this->surface = surface;
		this->window = window;

		CreateSwapChain({window->GetWidth(), window->GetHeight() });

		FrameBuffer::Init(device, vk::Extent3D(size, 1));

		Logger::RENDER->trace("Initialized swapchain");
	}

	void SwapChain::Close()
	{
		FrameBuffer::Close();
		DestroySwapChain();
		device = nullptr;
		window = nullptr;
		Logger::RENDER->trace("Closed swapchain");
	}

	void SwapChain::Resize(const uint32_t newWidth, const uint32_t newHeight)
	{
		if(newWidth == 0 || newHeight == 0) return; // Swap chain size of 0 pixel is not allowed

		CreateSwapChain({ newWidth, newHeight });
		FrameBuffer::Resize(vk::Extent3D(size, 1));
		Logger::RENDER->trace("Resized swapchain");
	}

	uint32_t SwapChain::AcquireNextImage(const vk::Fence& fence)
	{
		currentSemaphoreId = (currentSemaphoreId + 1) % imageAvailableSemaphores.size();
		const auto resultValue = device->device.acquireNextImageKHR(swapChain, UINT64_MAX, imageAvailableSemaphores[currentSemaphoreId], fence);
		const vk::Result result = resultValue.result;
		if (result != vk::Result::eSuccess && result != vk::Result::eSuboptimalKHR) throw std::error_code(result);
		SetCurrentFrameId(resultValue.value);

		vk::Fence submitFence = GetCurrentSubmitFence();
		device->device.waitForFences(1, &submitFence, true, -1);
		device->device.resetFences(1, &submitFence);

		return currentFrameBufferId;
	}

	void SwapChain::CreateSwapChain(vk::Extent2D size)
	{
		Logger::RENDER->debug("Creating swap chain for window {0} ...", window->GetWindowId());
		surfaceFormat = ChoseSurfaceFormat();
		presentMode = ChosePresentMode();
		const vk::SurfaceCapabilitiesKHR surfaceCapabilities = device->physicalDevice.getSurfaceCapabilitiesKHR(surface);
		if(surfaceCapabilities.currentExtent.width != ~static_cast<uint32_t>(0))
		{ // The surface does provide it's size to the vulkan driver
			if (surfaceCapabilities.currentExtent != size)
			{
				Logger::RENDER->warn("Surface resolution ({}, {}) does not match given render resolution ({}, {}).", size.width, size.height, surfaceCapabilities.currentExtent.width, surfaceCapabilities.currentExtent.height);
			}
			size = surfaceCapabilities.currentExtent;
		}

		vk::SurfaceTransformFlagBitsKHR preTransform; //TODO add option to allow rotation and other modifications
		if (surfaceCapabilities.supportedTransforms & vk::SurfaceTransformFlagBitsKHR::eIdentity)
		{ preTransform = vk::SurfaceTransformFlagBitsKHR::eIdentity; }
		else { preTransform = surfaceCapabilities.currentTransform; }

		uint32_t usingImages = std::max(preferredImageCount, surfaceCapabilities.minImageCount);
		if (surfaceCapabilities.maxImageCount > 0) //GPU has limit of swap chain images
			usingImages = std::min(usingImages, surfaceCapabilities.maxImageCount);
		Logger::RENDER->debug("GPU supports {0} to {1} swap chain images. Preferred: {2}, Using: {3}", surfaceCapabilities.minImageCount, surfaceCapabilities.maxImageCount, preferredImageCount, usingImages);

		const vk::SwapchainCreateInfoKHR createInfo({}, surface, usingImages, surfaceFormat.format,
													surfaceFormat.colorSpace, size, 1,
													vk::ImageUsageFlagBits::eColorAttachment | vk::ImageUsageFlagBits::eTransferDst,
													vk::SharingMode::eExclusive, 0, nullptr, preTransform,
													vk::CompositeAlphaFlagBitsKHR::eOpaque, presentMode, VK_TRUE, swapChain);
		const vk::SwapchainKHR newSwapChain = device->device.createSwapchainKHR(createInfo);

		DestroySwapChain();
		swapChain = newSwapChain;
		this->size = size;

		CreateSwapChainImages();

		// We use one extra semaphore to prevent reusing them before they have been consumed
		for (size_t i = 0; i < images.size() + 1; i++)
			imageAvailableSemaphores.emplace_back(device->device.createSemaphore({}));

		fullscreenViewport = vk::Viewport{ 0, 0, (float)size.width, (float)size.height, 0, 1 };
		fullscreenScissor = vk::Rect2D{ {0,0}, size };
		Logger::RENDER->debug("Swap chain for window {0} created", window->GetWindowId());
	}

	void SwapChain::CreateSwapChainImages()
	{
		vk::ImageViewCreateInfo imgViewCreateInfo;
		imgViewCreateInfo.format = surfaceFormat.format;
		imgViewCreateInfo.subresourceRange.aspectMask = vk::ImageAspectFlagBits::eColor;
		imgViewCreateInfo.subresourceRange.levelCount = 1;
		imgViewCreateInfo.subresourceRange.layerCount = 1;
		imgViewCreateInfo.viewType = vk::ImageViewType::e2D;

		auto swapChainImages = device->device.getSwapchainImagesKHR(swapChain);

		images.resize(swapChainImages.size());
		device->ExecuteNow([&](auto cmdBuffer) {
			for (uint32_t i = 0; i < swapChainImages.size(); i++)
			{
				images[i].image = swapChainImages[i];
				imgViewCreateInfo.image = swapChainImages[i];
				images[i].view = device->device.createImageView(imgViewCreateInfo);
				images[i].fence = device->device.createFence({vk::FenceCreateFlags(vk::FenceCreateFlagBits::eSignaled)});

				const vk::ImageMemoryBarrier imgMemBarrier({}, vk::AccessFlagBits::eTransferWrite, vk::ImageLayout::eUndefined, vk::ImageLayout::ePresentSrcKHR, 0, 0, images[i].image, vk::ImageSubresourceRange(vk::ImageAspectFlagBits::eColor, 0, 1, 0, 1));
				cmdBuffer.pipelineBarrier(vk::PipelineStageFlagBits::eTopOfPipe, vk::PipelineStageFlagBits::eTransfer, {}, nullptr, nullptr, imgMemBarrier);
			}
		});
	}

	void SwapChain::DestroySwapChain()
	{
		for(auto& image : images)
		{
			device->device.destroyImageView(image.view);
			device->device.destroyFence(image.fence);
		}
		device->device.destroySwapchainKHR(swapChain);
		for (auto& semaphore : imageAvailableSemaphores)
			device->device.destroySemaphore(semaphore);
		imageAvailableSemaphores.clear();
		Logger::RENDER->trace("Destroyed swapchain");
	}

	vk::PresentModeKHR SwapChain::ChosePresentMode()
	{
		std::vector<vk::PresentModeKHR> presentModes = device->physicalDevice.getSurfacePresentModesKHR(surface);
#ifdef DEBUG
		std::string availableModes = "";
		for (const auto& presentMode : presentModes)
		{
			if (availableModes.length() > 0) availableModes += ", ";
			availableModes += vk::to_string(presentMode);
		}
		Logger::RENDER->debug("Available swap chain present modes {0}. Searching best mode for: vsync={1}", availableModes, useVsync);
#endif
		vk::PresentModeKHR mode = vk::PresentModeKHR::eFifo;
		if (useVsync)
		{
			if (Utils::Contains(presentModes, vk::PresentModeKHR::eMailbox)) mode = vk::PresentModeKHR::eMailbox;
		}
		else
		{
			if (Utils::Contains(presentModes, vk::PresentModeKHR::eImmediate)) mode = vk::PresentModeKHR::eImmediate;
			else if (Utils::Contains(presentModes, vk::PresentModeKHR::eFifoRelaxed)) mode = vk::PresentModeKHR::eFifoRelaxed;
		}
		Logger::RENDER->debug("Using swap chain present mode {0}", vk::to_string(mode));
		return mode;
	}

	vk::SurfaceFormatKHR SwapChain::ChoseSurfaceFormat()
	{ //TODO allow to chose output format
		std::vector<vk::SurfaceFormatKHR> surfaceFormats = device->physicalDevice.getSurfaceFormatsKHR(surface);
#ifdef DEBUG
		std::string sfList;
		surfaceFormats.reserve(260);
		for (const auto& surfaceFormat : surfaceFormats)
		{
			if (!sfList.empty()) sfList.append(", ");
			DataFormat format(static_cast<DataFormat::Format>(surfaceFormat.format));
			sfList.append("{ ").append(format.GetName()).append(", ");
			sfList.append(magic_enum::enum_name(surfaceFormat.colorSpace)).append(" }");
		}
		Logger::RENDER->info("Available surface formats: {}", sfList);
#endif
		if(surfaceFormats.size() == 1 && surfaceFormats[0].format == vk::Format::eUndefined)
		{ // GPU does not have a preferred surface format
			return vk::SurfaceFormatKHR{ vk::Format::eB8G8R8A8Unorm, surfaceFormats[0].colorSpace };
		}
		else
		{ //TODO chose best fitting
			for (const auto& format : surfaceFormats)
			{
				if (DataFormat(static_cast<DataFormat::Format>(format.format)).IsSRGB() == EngineConfiguration::GetEngineConfiguration()->GetPreferFramebufferFormatSRGB())
				{
					return format;
				}
			}
			return surfaceFormats[0];
		}
	}

	std::vector<IImage*> SwapChain::GetImages()
	{
		std::vector<IImage*> imgs;
		imgs.reserve(images.size());
		for (auto& image : images)
		{
			imgs.push_back(&image);
		}
		return imgs;
	}
}