/*
 * 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 "VulkanShader.hpp"
#include "Vulkan/Context.hpp"
#include "Scene/Shader/Shader.hpp"
#include "Vulkan/Resources/IShaderOwner.hpp"
#include "Scene/Shader/DescriptorInputDescription.hpp"

namespace OpenVulkano::Vulkan
{
	static_assert(sizeof(vk::DescriptorSetLayoutBinding) == sizeof(DescriptorSetLayoutBinding));
	static_assert(sizeof(vk::PushConstantRange) == sizeof(PushConstantRange));

	VulkanShader::~VulkanShader()
	{
		if (owner) owner->RemoveShader(this);
		owner = nullptr;
		device.destroyPipeline(pipeline);
		for(auto& shaderModule : shaderModules)
		{
			device.destroyShaderModule(shaderModule);
		}
		device.destroyPipelineLayout(pipelineLayout);
		for(auto& descriptorSetLayout : descriptorSetLayouts)
			device.destroyDescriptorSetLayout(descriptorSetLayout);
	}

	VulkanShader::VulkanShader(Context* context, Scene::Shader* shader, IShaderOwner* owner)
		: IRenderResource<Scene::Shader>(shader), device(context->device->device), owner(owner), context(context)
	{
		shaderModules.reserve(shader->shaderPrograms.size());
		shaderStageCreateInfo.resize(shader->shaderPrograms.size());
		int i = 0;
		for(const auto& shaderProgram : shader->shaderPrograms)
		{
			const auto shaderModule = context->device->CreateShaderModule(shaderProgram.GetShaderNameVulkan());
			shaderModules.push_back(shaderModule);
			shaderStageCreateInfo[i] = {{}, static_cast<vk::ShaderStageFlagBits>(static_cast<uint32_t>(shaderProgram.type)), shaderModule, "main"};
			i++;
		}
		BuildPipeline();
	}

	void VulkanShader::BuildPipeline()
	{
		Scene::Shader* shader = GetOwner();
		std::vector<vk::VertexInputBindingDescription> vertexBindDesc;
		std::vector<vk::VertexInputAttributeDescription> attributeDescriptions;
		vertexBindDesc.reserve(shader->vertexInputDescriptions.size());

		for(const auto& description : shader->vertexInputDescriptions)
		{
			vertexBindDesc.emplace_back(description.bindingId, description.vertexSize, static_cast<vk::VertexInputRate>(description.stepMode));
			if (shader->vertexInputDescriptions.size() > 1)
			{
				for(const auto& param : description.inputParameters)
				{
					attributeDescriptions.push_back(reinterpret_cast<const vk::VertexInputAttributeDescription&>(param));
				}
			}
		}
		uint32_t attributeDescriptionsSize;
		vk::VertexInputAttributeDescription* attributeDescriptionsData;
		if (shader->vertexInputDescriptions.size() == 1)
		{ // Reuse already existing vertex attribute description
			static_assert(sizeof(VertexInputParameter) == sizeof(vk::VertexInputAttributeDescription));
			attributeDescriptionsSize = static_cast<uint32_t>(shader->vertexInputDescriptions[0].inputParameters.size());
			attributeDescriptionsData = reinterpret_cast<vk::VertexInputAttributeDescription*>(shader->vertexInputDescriptions[0].inputParameters.data());
		}
		else
		{
			attributeDescriptionsSize = static_cast<uint32_t>(attributeDescriptions.size());
			attributeDescriptionsData = attributeDescriptions.data();
		}

		vk::PipelineViewportStateCreateInfo viewportStateCreateInfo = { {}, 1, &context->swapChain.GetFullscreenViewport(), 1, &context->swapChain.GetFullscreenScissor() };
		vk::PipelineVertexInputStateCreateInfo pipelineVertexInputStateCreateInfo = {
				{}, static_cast<uint32_t>(vertexBindDesc.size()), vertexBindDesc.data(),
				attributeDescriptionsSize, attributeDescriptionsData };
		vk::PipelineInputAssemblyStateCreateInfo inputAssembly = { {}, static_cast<vk::PrimitiveTopology>(shader->topology), 0 };
		vk::PipelineRasterizationStateCreateInfo rasterizer = {};
		rasterizer.cullMode = static_cast<vk::CullModeFlagBits>(shader->cullMode);
		if (shader->depthBias)
		{
			rasterizer.depthBiasEnable = VK_TRUE;
			rasterizer.depthBiasClamp = shader->depthBiasClamp;
			rasterizer.depthBiasConstantFactor = shader->depthBiasConstant;
			rasterizer.depthBiasSlopeFactor = shader->depthBiasSlope;
		}
		vk::PipelineMultisampleStateCreateInfo msaa = {};
		vk::PipelineDepthStencilStateCreateInfo depth = { {}, shader->depthTest, shader->depthWrite, static_cast<vk::CompareOp>(shader->depthCompareOp) };
		depth.maxDepthBounds = 1;
		vk::PipelineColorBlendAttachmentState colorBlendAttachment = {};
		colorBlendAttachment.colorWriteMask = vk::ColorComponentFlagBits::eA | vk::ColorComponentFlagBits::eB | vk::ColorComponentFlagBits::eG | vk::ColorComponentFlagBits::eR;
		if (shader->alphaBlend)
		{
			colorBlendAttachment.blendEnable = true;
			colorBlendAttachment.srcColorBlendFactor = vk::BlendFactor::eSrcAlpha;
			colorBlendAttachment.dstColorBlendFactor = vk::BlendFactor::eOneMinusSrcAlpha;
			colorBlendAttachment.colorBlendOp = vk::BlendOp::eAdd;
			colorBlendAttachment.srcAlphaBlendFactor = vk::BlendFactor::eOne;
			colorBlendAttachment.dstAlphaBlendFactor = vk::BlendFactor::eZero;
			colorBlendAttachment.alphaBlendOp = vk::BlendOp::eAdd;
		}
		vk::PipelineColorBlendStateCreateInfo colorInfo = {};
		colorInfo.attachmentCount = 1;
		colorInfo.pAttachments = &colorBlendAttachment;

		CreatePipelineLayout();

		const auto dynSates = Utils::MakeStdArray<vk::DynamicState>( vk::DynamicState::eViewport, vk::DynamicState::eScissor );
		vk::PipelineDynamicStateCreateInfo dynStateInfo { {}, dynSates.size(), dynSates.data() };

		vk::GraphicsPipelineCreateInfo pipelineCreateInfo = { {}, static_cast<uint32_t>(shaderStageCreateInfo.size()),
															  shaderStageCreateInfo.data(), &pipelineVertexInputStateCreateInfo,
															  &inputAssembly, nullptr, &viewportStateCreateInfo, &rasterizer,
															  &msaa, &depth, &colorInfo, shader->dynamicViewport ? &dynStateInfo : nullptr, pipelineLayout,
															  context->swapChainRenderPass.renderPass };
		pipeline = this->device.createGraphicsPipeline({}, pipelineCreateInfo).value;
	}

	void VulkanShader::Resize()
	{
		device.destroyPipeline(pipeline);
		BuildPipeline();
	}

	void VulkanShader::Record(VulkanDrawContext* context)
	{
		context->commandBuffer.bindPipeline(vk::PipelineBindPoint::eGraphics, pipeline);
	}

	void VulkanShader::CreatePipelineLayout()
	{
		if (!descriptorSetLayouts.empty())
		{
			for(auto& descriptorSetLayout : descriptorSetLayouts)
				device.destroyDescriptorSetLayout(descriptorSetLayout);
			descriptorSetLayouts.clear();
		}
		std::array<vk::DescriptorSetLayoutBinding, 1> layoutBindings1 = { reinterpret_cast<const vk::DescriptorSetLayoutBinding&>(Scene::Camera::DESCRIPTOR_SET_LAYOUT_BINDING) };
		std::array<vk::DescriptorSetLayoutBinding, 1> layoutBindings2 = { reinterpret_cast<const vk::DescriptorSetLayoutBinding&>(Scene::Node::DESCRIPTOR_SET_LAYOUT_BINDING) };
		descriptorSetLayouts.push_back(device.createDescriptorSetLayout({ {}, layoutBindings1.size(), layoutBindings1.data() }));
		descriptorSetLayouts.push_back(device.createDescriptorSetLayout({ {}, layoutBindings2.size(), layoutBindings2.data() }));

		Scene::Shader* shader = GetOwner();
		for(const auto& set : shader->descriptorSets)
		{
			vk::DescriptorSetLayoutCreateInfo createInfo { {}, static_cast<uint32_t>(set.size()), reinterpret_cast<const vk::DescriptorSetLayoutBinding*>(set.data()) };
			descriptorSetLayouts.push_back(device.createDescriptorSetLayout(createInfo));
		}
		vk::PushConstantRange* pcRanges = reinterpret_cast<vk::PushConstantRange*>(shader->pushConstantRanges.data());
		vk::PipelineLayoutCreateInfo plci = {{}, static_cast<uint32_t>(descriptorSetLayouts.size()), descriptorSetLayouts.data(), static_cast<uint32_t>(shader->pushConstantRanges.size()), pcRanges };
		pipelineLayout = this->device.createPipelineLayout(plci);
	}

	void VulkanShader::Release()
	{
		if (owner) owner->RemoveShader(this);
	}
}