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

#include "IRecordable.hpp"
#include "Vulkan/Image.hpp"
#include "Vulkan/Context.hpp"
#include "Vulkan/Resources/ResourceManager.hpp"
#include "Vulkan/Scene/VulkanShader.hpp"
#include "Scene/Texture.hpp"
#include "Scene/Shader/Shader.hpp"

namespace OpenVulkano::Vulkan
{
	class VulkanTexture : public IRenderResource<Scene::Texture>, public IRecordable, public Image
	{
	public:
		vk::Sampler m_sampler;
		vk::DescriptorSet m_descriptorSet;

		VulkanTexture() : IRenderResource<Scene::Texture>(nullptr) {}

		virtual void Init(ResourceManager* resManager, Scene::Texture* texture, vk::DescriptorSetLayout* descriptorSetLayout, const DescriptorSetLayoutBinding& binding)
		{
			Image::Init(resManager->GetContext()->device.get(), texture->format, { texture->resolution.x, texture->resolution.y, texture->resolution.z });
			resManager->CopyDataToImage(texture->size, texture->textureBuffer, this);
			texture->updated = false;

			m_sampler = resManager->CreateSampler(reinterpret_cast<const vk::SamplerCreateInfo&>(*texture->m_samplerConfig));
			SetDescriptorSet(resManager, descriptorSetLayout, binding);

			UpdateAddress(texture);
			UpdateRenderResource(GetOwnerResource());
		}

		void InitSharedMem(ResourceManager* resManager, Scene::Texture* texture, vk::DescriptorSetLayout* descriptorSetLayout, const DescriptorSetLayoutBinding& binding)
		{
			Image::Init(resManager->GetContext()->device.get(), texture->format, { texture->resolution.x, texture->resolution.y, texture->resolution.z }, vk::MemoryPropertyFlagBits::eDeviceLocal | vk::MemoryPropertyFlagBits::eHostVisible);
			texture->updated = false;
			texture->textureBuffer = Map();

			m_sampler = resManager->CreateSampler(reinterpret_cast<const vk::SamplerCreateInfo&>(*texture->m_samplerConfig));
			SetDescriptorSet(resManager, descriptorSetLayout, binding);

			UpdateAddress(texture);
			UpdateRenderResource(GetOwnerResource());
		}

		virtual ~VulkanTexture() override
		{
			if (m_sampler) Close();
		}

		void Close() override
		{
			ResetRenderResource(GetOwnerResource());
			Image::Close();
		}

		void SetDescriptorSet(ResourceManager* resManager, vk::DescriptorSetLayout* descriptorSetLayout, const DescriptorSetLayoutBinding& binding)
		{
			// Setup Descriptor set
			const vk::DescriptorSetAllocateInfo descSetAllocInfo = { ResourceManager::INSTANCE->descriptorPool, 1, descriptorSetLayout };
			m_descriptorSet = resManager->GetContext()->device->device.allocateDescriptorSets(descSetAllocInfo)[0];
			vk::DescriptorImageInfo imageInfo = { m_sampler, view, vk::ImageLayout::eShaderReadOnlyOptimal };
			vk::WriteDescriptorSet writeDescriptorSet = { m_descriptorSet, binding.bindingId, 0, 1 };
			writeDescriptorSet.descriptorType = static_cast<vk::DescriptorType>(binding.descriptorType);
			writeDescriptorSet.pImageInfo = &imageInfo;
			resManager->GetContext()->device->device.updateDescriptorSets(1, &writeDescriptorSet, 0, nullptr);
		}

		void Record(VulkanDrawContext* context) override
		{
			int setId = -1, i = 0;
			for (const auto& descriptorSet : context->GetShader()->GetOwner()->descriptorSets)
			{
				for (const auto& descriptor : descriptorSet)
				{
					if (descriptor.descriptorType < 6)
					{
						if (setId != -1) [[unlikely]]
						{
							Logger::RENDER->error("Found multiple possible descriptor set for texture! Use 'Record(VulkanDrawContext* context, int setId)' explicitly when using multiple sets containing textures!");
						}
						setId = i;
						break;
					}
				}
				i++;
			}
			if (setId == -1) [[unlikely]]
			{
				Logger::RENDER->error("Failed to discover setId for texture binding.");
				return;
			}
			else [[likely]]
			{
				Record(context, setId + 2); // Set 0 and 1 are automatically bound special sets for the camera and the node
			}
		}

		virtual void Record(VulkanDrawContext* context, int setId)
		{
			context->commandBuffer.bindDescriptorSets(vk::PipelineBindPoint::eGraphics, context->GetShader()->pipelineLayout, setId, 1, &m_descriptorSet, 0, nullptr);
		}

		void Release() override
		{
			//TODO
		}
	};

	class VulkanTextureDynamic : public VulkanTexture
	{
	public:
		void Init(ResourceManager* resManager, Scene::Texture* texture, vk::DescriptorSetLayout* descriptorSetLayout, const DescriptorSetLayoutBinding& binding) override
		{
			VulkanTexture::Init(resManager, texture, descriptorSetLayout, binding);
		}

		void Record(VulkanDrawContext* context) override
		{
			if(GetOwner()->updated)
			{
				context->renderer->GetResourceManager().CopyDataToImage(GetOwner()->size, GetOwner()->textureBuffer, this);
				GetOwner()->updated = false;
			}
			VulkanTexture::Record(context);
		}

		void Record(VulkanDrawContext* context, int setId) override
		{
			if(GetOwner()->updated)
			{
				context->renderer->GetResourceManager().CopyDataToImage(GetOwner()->size, GetOwner()->textureBuffer, this);
				GetOwner()->updated = false;
			}
			VulkanTexture::Record(context, setId);
		}
	};
}