#include "vkn/vulkan_api.h"
#include "vkn/vulkan_window.h"
#include "vkn/vulkan_api_help.h"
#include "vkn/wrapper/buffer.h"
#include "vkn/wrapper/device.h"
#include "vkn/thread/buffer_worker.h"
#include "vkn/thread/command_worker.h"
#include "render/asset/mesh.h"
#include "meta/enum.h"
#include "tinyimageformat/tinyimageformat_apis.h"
#include "zlog.h"
namespace vkn {
	inline bool operator==(const RenderPassKey& k1, const RenderPassKey& k2) {
		if (k1.initialColorLayoutMask != k2.initialColorLayoutMask) return false;
		for (int i = 0; i < MAX_SUPPORTED_RENDER_TARGET_COUNT; i++) {
			if (k1.colorFormat[i] != k2.colorFormat[i]) return false;
		}
		if (k1.depthFormat != k2.depthFormat) return false;
		if (k1.clear != k2.clear) return false;
		if (k1.discardStart != k2.discardStart) return false;
		if (k1.discardEnd != k2.discardEnd) return false;
		if (k1.samples != k2.samples) return false;
		if (k1.needsResolveMask != k2.needsResolveMask) return false;
		if (k1.subpassMask != k2.subpassMask) return false;
		return true;
	}
	inline bool operator==(const FramebufferKey& k1, const FramebufferKey& k2) {
		if (k1.pass != k2.pass) return false;
		if (k1.attachmentCount != k2.attachmentCount) return false;
		for (int i = 0; i < k1.attachmentCount; i++) {
			if (k1.imageViews[i] != k2.imageViews[i]) return false;
		}
		if (k1.height != k2.height) return false;
		if (k1.width != k2.width) return false;
		if (k1.layers != k2.layers) return false;
		return true;
	}
	VulkanAPI::VulkanAPI() : RenderAPI(new VulkanContext())
		, window(*VulkanWindow::Ptr())
		, backend(VulkanEngineName) 
	{

	}
	void VulkanAPI::Init()
	{

	}
	void VulkanAPI::Shutdown()
	{

	}
	void VulkanAPI::SetStaticMesh(Mesh& mesh)
	{
		auto& Indices = mesh.GetIndices();
		auto& Vertices = mesh.GetVertices();
		MeshVAO& VAO = MeshTable[mesh.GetGuid()];
		VAO.indexCount = Indices.size();
		VAO.vertexCount = Vertices.size();

		Buffer indexBuffer{};
		indexBuffer.ppBuffer = &VAO.indexBuffer;
		indexBuffer.pCpuData = Indices.data();
		indexBuffer.size = sizeof(decltype(Indices[0])) * Indices.size();
		indexBuffer.usage = VK_BUFFER_USAGE_INDEX_BUFFER_BIT;
		Backend::TransferWorker->Invoke(indexBuffer);

		Buffer vertexBuffer{};
		vertexBuffer.ppBuffer = &VAO.vertexBuffer;
		vertexBuffer.pCpuData = Vertices.data();
		vertexBuffer.size = Vertices.data_size();
		vertexBuffer.usage = VK_BUFFER_USAGE_VERTEX_BUFFER_BIT;
		Backend::TransferWorker->Invoke(vertexBuffer);
	}
	void VulkanAPI::DrawStaticMesh(Mesh& mesh)
	{

	}
	void VulkanAPI::LoadShader(Shader& shader)
	{
	}
	ImagePtr VulkanAPI::CreateTexture(TextureDesc desc)
	{
		if (desc.image) {
			return desc.image;
		}
		uint32_t depth = any(desc.dimension & TextureDimension::TEX_3D) ? 3 : 1;
		VkImageCreateInfo imageCreateInfo = {
			.sType = VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO,
			.pNext = NULL,
			.flags = vkApiGetImageCreateFlag(desc.dimension, desc.arraySize),
			.imageType = vkApiGetImageType(desc.dimension),
			.format = (VkFormat)TinyImageFormat_ToVkFormat(desc.format),
			.extent = {
				.width = (uint32_t)desc.width,
				.height = (uint32_t)desc.height,
				.depth = (uint32_t)desc.depth,
				},
			.mipLevels = desc.mipLevel,
			.arrayLayers = desc.arraySize,
			.samples = vkApiGetSmpleCountFlag(desc.sampleCount),
			.tiling = VK_IMAGE_TILING_OPTIMAL,
			.usage = vkApiGetImageUsageFlags(desc.state),
			.sharingMode = VK_SHARING_MODE_EXCLUSIVE,
			.queueFamilyIndexCount = 0,
			.pQueueFamilyIndices = NULL,
			.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED
		};
		return ImagePtr();
	}
	ImageViewPtr VulkanAPI::CreateTextureView(TextureViewDesc desc)
	{
		VkImageViewCreateInfo createInfo = {};
		createInfo.sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO;
		createInfo.image = (VkImage)desc.image;
		createInfo.viewType = vkApiGetImageViewType(desc.dimension, desc.layerCount);
		createInfo.format = (VkFormat)TinyImageFormat_ToVkFormat(desc.format);

		// components字段允许调整颜色通道的最终的映射逻辑
		// 比如，我们可以将所有颜色通道映射为红色通道，以实现单色纹理，我们也可以将通道映射具体的常量数值0和1
		// 这里用默认的
		createInfo.components.r = VK_COMPONENT_SWIZZLE_R;
		createInfo.components.g = VK_COMPONENT_SWIZZLE_G;
		createInfo.components.b = VK_COMPONENT_SWIZZLE_B;
		createInfo.components.a = VK_COMPONENT_SWIZZLE_A;

		// subresourceRangle字段用于描述图像的使用目标是什么，以及可以被访问的有效区域
		// 这个图像用作填充color还是depth stencil等
		createInfo.subresourceRange.aspectMask = vkApiGetImageAspectMask(createInfo.format, false);
		// 默认处理所有Mipmap
		createInfo.subresourceRange.baseMipLevel = desc.baseMipLevel;
		createInfo.subresourceRange.levelCount = desc.levelCount;
		// 默认处理所有Layers
		createInfo.subresourceRange.baseArrayLayer = desc.baseArrayLayer;
		createInfo.subresourceRange.layerCount = desc.layerCount;

		VkImageView imageView;
		vkCreateImageView(backend.GetDevice().Ptr(), &createInfo, nullptr, &imageView);
		return (ImageViewPtr)imageView;
	}
	void VulkanAPI::BeginFrame()
	{
		VulkanContext& ctx = *(VulkanContext*)&context;
		window.Aquire(ctx);
		ctx.command.BeginRecord(VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT);
	}
	void VulkanAPI::EndFrame()
	{
		VulkanContext& ctx = *(VulkanContext*)&context;
		ctx.command.EndRecord();
		VkPipelineStageFlags waitStages[] = { VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT };  // 等待渲染阶段
		VkSemaphore waitSemaphores[] = { ctx.surfaceSemaphore };
		VkSemaphore signalSemaphores[] = { ctx.presentSemaphore };// 渲染完成信号量
		VkSubmitInfo submitInfo{};
		submitInfo.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;
		submitInfo.commandBufferCount = 1;
		submitInfo.pCommandBuffers = &ctx.command.Ptr();

		submitInfo.waitSemaphoreCount = 1;
		submitInfo.pWaitSemaphores = waitSemaphores;
		submitInfo.pWaitDstStageMask = waitStages;

		submitInfo.signalSemaphoreCount = 1;
		submitInfo.pSignalSemaphores = signalSemaphores;
		vkQueueSubmit(Backend::RenderWorker->GetQueue().Ptr(), 1, &submitInfo, ctx.surfaceFence);
		window.Present(ctx);
	}
	void VulkanAPI::ExecuteResourceBarriers(const ResourceBarrierDesc& desc) {
		VulkanContext& ctx = *(VulkanContext*)&context;
		pmr::vector<VkBufferMemoryBarrier> bufferBarriers{ FramePool() };
		bufferBarriers.reserve(desc.bufferBarriersCount);
		pmr::vector<VkImageMemoryBarrier> imageBarriers{ FramePool() };
		imageBarriers.reserve(desc.textureBarriersCount);
		VkPipelineStageFlags srcStageMask = 0, dstStageMask = 0;
		for (uint32_t i = 0; i < desc.textureBarriersCount; i++)
		{
			auto& barrier = desc.pTextureBarriers[i];
			auto desc = vkApiGetTextureTransition(srcStageMask, dstStageMask, barrier);
			imageBarriers.push_back(desc);
		}
		if (dstStageMask == 0) {
			dstStageMask = VK_PIPELINE_STAGE_BOTTOM_OF_PIPE_BIT;
		}
		vkCmdPipelineBarrier(ctx.command.Ptr(), srcStageMask, dstStageMask, 0, 0, NULL, 0, NULL, imageBarriers.size(), imageBarriers.data());
	}
	void VulkanAPI::BeginRenderPass(RenderPassNode* node)
	{
		RenderPassKey  config{};
		FramebufferKey frameKey{.layers = 1};
		VkClearValue   clearValues[2 * MAX_SUPPORTED_RENDER_TARGET_COUNT + 1] = { 0 };
		int i = 0;
		for (auto& it : node->outEdges) {
			if (it->IsAttachment()) {
				auto& desc = it->CastTo<api::AttachmentDesc>();

				config.colorFormat[i] = (VkFormat)TinyImageFormat_ToVkFormat(desc.colorFormat);
				config.samples = vkApiGetSmpleCountFlag(desc.sampleCount);

				TargetBufferFlags flag = TargetBufferFlags(int(TargetBufferFlags::COLOR0) << i);
				config.clear |= flag;
				it.ResolveView(&graph);
				frameKey.imageViews[i] = (VkImageView)desc.imageView;
				frameKey.height = desc.height;
				frameKey.width = desc.width;

				//clearValues[i] = 
				i++;
			}
		}
		frameKey.attachmentCount = i;
		VkRenderPass pass = GetRenderPass(config);
		frameKey.pass = pass;
		auto it = FramebufferCache.find(frameKey);
		VkFramebuffer framebuffer = it->second;
		if (it == FramebufferCache.end()) {
			VkFramebufferCreateInfo framebufferInfo = {
				.sType = VK_STRUCTURE_TYPE_FRAMEBUFFER_CREATE_INFO,
				.renderPass = pass,
				.attachmentCount = frameKey.attachmentCount,
				.pAttachments = frameKey.imageViews,
				.width = frameKey.width,
				.height = frameKey.height,
				.layers = frameKey.layers
			};
			vkCreateFramebuffer(backend.GetDevice().Ptr(), &framebufferInfo, nullptr, &framebuffer);
			FramebufferCache.emplace(frameKey, framebuffer);
		}
		VkRect2D renderAarea = {
			.offset = {0,0},
			.extent = {frameKey.width,frameKey.height}
		};
		VkRenderPassBeginInfo beginInfo = {
			.sType = VK_STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO,
			.pNext = VK_NULL_HANDLE,
			.renderPass = pass,
			.framebuffer = framebuffer,
			.renderArea = renderAarea,
			.clearValueCount = frameKey.attachmentCount,
			.pClearValues = clearValues
		};
		VulkanContext& ctx = *(VulkanContext*)&context;
		vkCmdBeginRenderPass(ctx.command.Ptr(), &beginInfo, VK_SUBPASS_CONTENTS_INLINE);
	}
	void VulkanAPI::EndRenderPass(RenderPassNode* node)
	{
		VulkanContext& ctx = *(VulkanContext*)&context;
		vkCmdEndRenderPass(ctx.command.Ptr());
	}
	VkRenderPass VulkanAPI::GetRenderPass(RenderPassKey& config) {
		auto it = RenderPassCache.find(config);
		if (it != RenderPassCache.end()) {
			return it->second;
		}
		// Set up some const aliases for terseness.
		const VkAttachmentLoadOp kClear = VK_ATTACHMENT_LOAD_OP_CLEAR;
		const VkAttachmentLoadOp kDontCare = VK_ATTACHMENT_LOAD_OP_DONT_CARE;
		const VkAttachmentLoadOp kKeep = VK_ATTACHMENT_LOAD_OP_LOAD;
		const VkAttachmentStoreOp kDisableStore = VK_ATTACHMENT_STORE_OP_DONT_CARE;
		const VkAttachmentStoreOp kEnableStore = VK_ATTACHMENT_STORE_OP_STORE;

		VkAttachmentReference inputAttachmentRef[MAX_SUPPORTED_RENDER_TARGET_COUNT] = {};
		VkAttachmentReference colorAttachmentRefs[2][MAX_SUPPORTED_RENDER_TARGET_COUNT] = {};
		VkAttachmentReference resolveAttachmentRef[MAX_SUPPORTED_RENDER_TARGET_COUNT] = {};
		VkAttachmentReference depthAttachmentRef = {};

		const bool hasSubpasses = config.subpassMask != 0;
		const bool hasDepth = config.depthFormat != VK_FORMAT_UNDEFINED;

		VkSubpassDescription subpasses[2] = { {
			.pipelineBindPoint = VK_PIPELINE_BIND_POINT_GRAPHICS,
			.pInputAttachments = nullptr,
			.pColorAttachments = colorAttachmentRefs[0],
			.pResolveAttachments = resolveAttachmentRef,
			.pDepthStencilAttachment = hasDepth ? &depthAttachmentRef : nullptr
		},
		{
			.pipelineBindPoint = VK_PIPELINE_BIND_POINT_GRAPHICS,
			.pInputAttachments = inputAttachmentRef,
			.pColorAttachments = colorAttachmentRefs[1],
			.pResolveAttachments = resolveAttachmentRef,
			.pDepthStencilAttachment = hasDepth ? &depthAttachmentRef : nullptr
		} };

		// The attachment list contains: Color Attachments, Resolve Attachments, and Depth Attachment.
		// For simplicity, create an array that can hold the maximum possible number of attachments.
		// Note that this needs to have the same ordering as the corollary array in getFramebuffer.
		VkAttachmentDescription attachments[MAX_SUPPORTED_RENDER_TARGET_COUNT + MAX_SUPPORTED_RENDER_TARGET_COUNT + 1] = {};

		// We support 2 subpasses, which means we need to supply 1 dependency struct.
		VkSubpassDependency dependencies[1] = { {
			.srcSubpass = 0,
			.dstSubpass = 1,
			.srcStageMask = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT,
			.dstStageMask = VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT,
			.srcAccessMask = VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT,
			.dstAccessMask = VK_ACCESS_SHADER_READ_BIT,
			.dependencyFlags = VK_DEPENDENCY_BY_REGION_BIT,
		} };
		// Finally, create the VkRenderPass.
		VkRenderPassCreateInfo renderPassInfo{
			.sType = VK_STRUCTURE_TYPE_RENDER_PASS_CREATE_INFO,
			.attachmentCount = 0u,
			.pAttachments = attachments,
			.subpassCount = hasSubpasses ? 2u : 1u,
			.pSubpasses = subpasses,
			.dependencyCount = hasSubpasses ? 1u : 0u,
			.pDependencies = dependencies
		};
		int attachmentIndex = 0;
		// Populate the Color Attachments.
		for (int i = 0; i < MAX_SUPPORTED_RENDER_TARGET_COUNT; i++) {
			if (config.colorFormat[i] == VK_FORMAT_UNDEFINED) {
				continue;
			}
			const VkImageLayout subpassLayout = VK_IMAGE_LAYOUT_GENERAL;
			uint32_t index;

			if (!hasSubpasses) {
				index = subpasses[0].colorAttachmentCount++;
				colorAttachmentRefs[0][index].layout = subpassLayout;
				colorAttachmentRefs[0][index].attachment = attachmentIndex;
			}
			else {

				// The Driver API consolidates all color attachments from the first and second subpasses
				// into a single list, and uses a bitmask to mark attachments that belong only to the
				// second subpass and should be available as inputs. All color attachments in the first
				// subpass are automatically made available to the second subpass.

				// If there are subpasses, we require the input attachment to be the first attachment.
				// Breaking this assumption would likely require enhancements to the Driver API in order
				// to supply Vulkan with all the information needed.
				if (config.subpassMask & (1 << i)) {
					index = subpasses[0].colorAttachmentCount++;
					colorAttachmentRefs[0][index].layout = subpassLayout;
					colorAttachmentRefs[0][index].attachment = attachmentIndex;

					index = subpasses[1].inputAttachmentCount++;
					inputAttachmentRef[index].layout = subpassLayout;
					inputAttachmentRef[index].attachment = attachmentIndex;
				}

				index = subpasses[1].colorAttachmentCount++;
				colorAttachmentRefs[1][index].layout = subpassLayout;
				colorAttachmentRefs[1][index].attachment = attachmentIndex;
			}
			const TargetBufferFlags flag = TargetBufferFlags(int(TargetBufferFlags::COLOR0) << i);
			const bool clear = any(config.clear & flag);
			const bool discard = any(config.discardStart & flag);
			VkImageLayout layout = vkApiGetAttachmentLayout(config.colorFormat[i], true);
			attachments[attachmentIndex++] = {
				.format = config.colorFormat[i],
				.samples = config.samples,
				.loadOp = clear ? kClear : (discard ? kDontCare : kKeep),
				.storeOp = kEnableStore,
				.stencilLoadOp = kDontCare,
				.stencilStoreOp = kDisableStore,
				.initialLayout = layout,
				.finalLayout = layout,
			};
		}
		// Nulling out the zero-sized lists is necessary to avoid VK_ERROR_OUT_OF_HOST_MEMORY on Adreno.
		if (subpasses[0].colorAttachmentCount == 0) {
			subpasses[0].pColorAttachments = nullptr;
			subpasses[0].pResolveAttachments = nullptr;
			subpasses[1].pColorAttachments = nullptr;
			subpasses[1].pResolveAttachments = nullptr;
		}
		// Populate the Resolve Attachments.
		VkAttachmentReference* pResolveAttachment = resolveAttachmentRef;
		for (int i = 0; i < MAX_SUPPORTED_RENDER_TARGET_COUNT; i++) {
			if (config.colorFormat[i] == VK_FORMAT_UNDEFINED) {
				continue;
			}

			if (!(config.needsResolveMask & (1 << i))) {
				pResolveAttachment->attachment = VK_ATTACHMENT_UNUSED;
				++pResolveAttachment;
				continue;
			}

			pResolveAttachment->attachment = attachmentIndex;
			pResolveAttachment->layout = VK_IMAGE_LAYOUT_GENERAL;
			++pResolveAttachment;

			attachments[attachmentIndex++] = {
				.format = config.colorFormat[i],
				.samples = VK_SAMPLE_COUNT_1_BIT,
				.loadOp = kDontCare,
				.storeOp = kEnableStore,
				.stencilLoadOp = kDontCare,
				.stencilStoreOp = kDisableStore,
				.initialLayout = VK_IMAGE_LAYOUT_GENERAL,
				.finalLayout = VK_IMAGE_LAYOUT_GENERAL,
			};
		}
		// Populate the Depth Attachment.
		if (hasDepth) {
			const bool clear = any(config.clear & TargetBufferFlags::DEPTH);
			const bool discardStart = any(config.discardStart & TargetBufferFlags::DEPTH);
			const bool discardEnd = any(config.discardEnd & TargetBufferFlags::DEPTH);
			depthAttachmentRef.layout = VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL;
			depthAttachmentRef.attachment = attachmentIndex;
			attachments[attachmentIndex++] = {
				.format = config.depthFormat,
				.samples = (VkSampleCountFlagBits)config.samples,
				.loadOp = clear ? kClear : (discardStart ? kDontCare : kKeep),
				.storeOp = discardEnd ? kDisableStore : kEnableStore,
				.stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE,
				.stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE,
				.initialLayout = VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL,
				.finalLayout = VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL,
			};
		}
		renderPassInfo.attachmentCount = attachmentIndex;
		VkRenderPass renderPass;
		VkResult error = vkCreateRenderPass(backend.GetDevice().Ptr(), &renderPassInfo, nullptr, &renderPass);
		RenderPassCache.emplace(config, renderPass);
		return renderPass;
	}
}