/*
 * 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 "Renderer.hpp"
#include "VulkanDrawContext.hpp"
#include "Scene/Shader/Shader.hpp"
#include "Scene/Geometry.hpp"
#include "Scene/VulkanGeometry.hpp"
#include "Scene/VulkanNode.hpp"
#include "Scene/VulkanShader.hpp"
#include "Base/UI/IVulkanWindow.hpp"
#include "Host/PlatformProducer.hpp"
#include "Vulkan/Scene/VulkanCamera.hpp"
#include <stdexcept>

namespace OpenVulkano::Vulkan
{
	void Renderer::Init(IGraphicsAppManager* graphicsAppManager, IWindow* window)
	{
		logger = Logger::RENDER;
		logger->info("Initializing Vulkan renderer ...");
		IVulkanWindow* vulkanWindow = window->GetVulkanWindow();
		if (!vulkanWindow)
		{
			logger->error("The provided window is not compatible with Vulkan.");
			throw std::runtime_error("The provided window is not compatible with Vulkan.");
		}
		context.Init(graphicsAppManager, vulkanWindow);
		for (int i = 0; i < context.swapChain.GetImageCount(); i++)
		{
			waitSemaphores.emplace_back();
			waitSemaphores[i].renderComplete.push_back(context.device->device.createSemaphore({}));
			waitSemaphores[i].renderReady.resize(2);
		}
		resourceManager.Init(&context, context.swapChain.GetImageCount());
		uiRenderer.Init(&context);
		threadPool.resize(EngineConfiguration::GetEngineConfiguration()->GetNumThreads() - 1);

		//Setup cmd pools and buffers
		commands.resize(threadPool.size() + 2); // One extra cmd object for the primary buffer and one for the main thread
		for(uint32_t i = 0; i < commands.size(); i++)
		{
			commands[i] = std::vector<CommandHelper>(context.swapChain.GetImageCount());
			for(size_t j = 0; j < commands[i].size(); j++)
			{
				commands[i][j].Init(context.device->device, context.device->queueIndices.GetGraphics(),
				(i == commands.size() - 1) ? vk::CommandBufferLevel::ePrimary : vk::CommandBufferLevel::eSecondary);
			}
		}
		submitBuffers.resize(context.swapChain.GetImageCount());
		for(uint32_t i = 0; i < submitBuffers.size(); i++)
		{
			submitBuffers[i].resize(commands.size() - 1);
			for (size_t j = 0; j < submitBuffers[i].size(); j++)
			{
				submitBuffers[i][j] = commands[j][i].cmdBuffer;
			}
		}

		logger->info("Vulkan renderer initialized");
	}

	void Renderer::Tick()
	{
		currentImageId = context.swapChain.AcquireNextImage();
		scene->GetCamera()->SetSize(context.window->GetWidth(), context.window->GetHeight());
		Render();
	}

	void Renderer::Close()
	{
		context.device->WaitIdle();
		for (int i = 0; i < context.swapChain.GetImageCount(); i++)
		{
			//context.device->device.destroySemaphore(waitSemaphores[i]);
		}
		uiRenderer.Close();
		resourceManager.Close();
		context.Close();
		scene = nullptr;
	}

	std::string Renderer::GetMainRenderDeviceName()
	{
		return (context.device) ? context.device->GetDeviceName() : "Unknown";
	}

	void Renderer::Resize(const uint32_t newWidth, const uint32_t newHeight)
	{
		if (!context.device) return;
		context.Resize(newWidth, newHeight);
		resourceManager.Resize();
	}

	CommandHelper* Renderer::GetCommandData(uint32_t poolId)
	{
		return &commands[poolId][currentImageId];
	}

	void Renderer::RunThread(Renderer* renderer, Data::ReadOnlyAtomicArrayQueue<Scene::Drawable*>* jobQueue, uint32_t id)
	{
		renderer->RecordSecondaryBuffer(jobQueue, id);
	}

	void Renderer::StartThreads(Data::ReadOnlyAtomicArrayQueue<Scene::Drawable*>* jobQueue)
	{
		for(uint32_t i = 0; i < threadPool.size(); i++)
		{
			threadPool[i] = std::thread(RunThread, this, jobQueue, i);
		}
	}

	void Renderer::RecordPrimaryBuffer()
	{
		CommandHelper* cmdHelper = GetCommandData(commands.size() - 1);
		cmdHelper->Reset();
		cmdHelper->cmdBuffer.begin(vk::CommandBufferBeginInfo(vk::CommandBufferUsageFlagBits::eOneTimeSubmit));
		context.swapChainRenderPass.Begin(cmdHelper->cmdBuffer);
	}

	void Renderer::Submit()
	{
		for (auto& thread : threadPool) { thread.join(); } // Wait till everything is recorded
		CommandHelper* cmdHelper = GetCommandData(commands.size() - 1);
		cmdHelper->cmdBuffer.executeCommands(submitBuffers[currentImageId].size(), submitBuffers[currentImageId].data());
		context.swapChainRenderPass.End(cmdHelper->cmdBuffer);
		uiRenderer.DrawUiFrame(cmdHelper->cmdBuffer);
		cmdHelper->cmdBuffer.end();
		std::array<vk::PipelineStageFlags, 2> stateFlags = { vk::PipelineStageFlags(vk::PipelineStageFlagBits::eColorAttachmentOutput), vk::PipelineStageFlags(vk::PipelineStageFlagBits::eColorAttachmentOutput) };
		waitSemaphores[currentImageId].renderReady[0] = resourceManager.EndFrame();
		waitSemaphores[currentImageId].renderReady[1] = context.swapChain.GetCurrentSemaphore();
		vk::SubmitInfo si = vk::SubmitInfo(
				waitSemaphores[currentImageId].renderReady.size(), waitSemaphores[currentImageId].renderReady.data(), stateFlags.data(),
				1, &cmdHelper->cmdBuffer,
				waitSemaphores[currentImageId].renderComplete.size(), waitSemaphores[currentImageId].renderComplete.data());
		auto submitResult = context.device->graphicsQueue.submit(1, &si, context.swapChain.GetCurrentSubmitFence());
		if (submitResult == vk::Result::eSuccess) [[likely]]
			context.swapChain.Present(context.device->graphicsQueue, waitSemaphores[currentImageId].renderComplete);
		else [[unlikely]]
			Logger::RENDER->error("Failed to submit draw command buffer: {}", to_string(submitResult));
	}

	void Renderer::Render()
	{
		resourceManager.StartFrame(currentImageId);
		Data::ReadOnlyAtomicArrayQueue<Scene::Drawable*> jobQueue(scene->shapeList);
		StartThreads(&jobQueue);
		RecordPrimaryBuffer();
		RecordSecondaryBuffer(&jobQueue, threadPool.size());
		Submit();
	}

	void Renderer::RecordSecondaryBuffer(Data::ReadOnlyAtomicArrayQueue<Scene::Drawable*>* jobQueue, uint32_t poolId)
	{
		CommandHelper* cmdHelper = GetCommandData(poolId);
		cmdHelper->Reset();
		vk::CommandBufferInheritanceInfo inheritance = { context.swapChainRenderPass.renderPass, 0, context.swapChainRenderPass.GetFrameBuffer()->GetCurrentFrameBuffer() };
		cmdHelper->cmdBuffer.begin(vk::CommandBufferBeginInfo{ vk::CommandBufferUsageFlagBits::eOneTimeSubmit | vk::CommandBufferUsageFlagBits::eRenderPassContinue, &inheritance });

		Scene::Drawable** drawablePointer;
		VulkanDrawContext drawContext { poolId, currentImageId, cmdHelper->cmdBuffer, this };

		drawContext.SetCamera(scene->GetCamera());

		while((drawablePointer = jobQueue->Pop()) != nullptr)
		{
			Scene::Drawable* drawable = *drawablePointer;
			drawContext.EncodeShader(drawable->GetShader());
			drawable->GetEncoder().vulkan(drawable, &drawContext);
		}
		cmdHelper->cmdBuffer.end();
	}

	RendererProducerRegistration<Renderer> rendererRegistration(RenderAPI::Vulkan);
}