#include "vkn/vulkan_imgui_editor.h"
#include "vkn/vulkan_window.h"
#include "vkn/vulkan_api_help.h"
#include "vkn/backend.h"
#include "vkn/wrapper/device.h"
#include "vkn/wrapper/instance.h"
#include "vkn/wrapper/queue.h"
#include "imgui/imgui_impl_vulkan.h"
#include "imgui/imgui_impl_sdl2.h"
#include "data/global.h"
#include "event/event_system.h"
#include "tinyimageformat/tinyimageformat_apis.h"
namespace vkn {
	using namespace api;
	static Name ImguiPassName{"ImguiPass"};
	static VkDescriptorPool g_DescriptorPool = VK_NULL_HANDLE;
	// Vulkan的ImGui接入比较麻烦，参考教程: https://frguthmann.github.io/posts/vulkan_imgui/
	VkDescriptorPool CreateDescriptorPool(VkDevice device) {
		VkDescriptorPoolSize pool_sizes[] =
		{
			{ VK_DESCRIPTOR_TYPE_SAMPLER, 1000 },
			{ VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, 1000 },
			{ VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE, 1000 },
			{ VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, 1000 },
			{ VK_DESCRIPTOR_TYPE_UNIFORM_TEXEL_BUFFER, 1000 },
			{ VK_DESCRIPTOR_TYPE_STORAGE_TEXEL_BUFFER, 1000 },
			{ VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 1000 },
			{ VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 1000 },
			{ VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER_DYNAMIC, 1000 },
			{ VK_DESCRIPTOR_TYPE_STORAGE_BUFFER_DYNAMIC, 1000 },
			{ VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT, 1000 }
		};
		VkDescriptorPoolCreateInfo pool_info = {};
		pool_info.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;
		pool_info.flags = VK_DESCRIPTOR_POOL_CREATE_FREE_DESCRIPTOR_SET_BIT;
		pool_info.maxSets = 1000 * IM_ARRAYSIZE(pool_sizes);
		pool_info.poolSizeCount = (uint32_t)IM_ARRAYSIZE(pool_sizes);
		pool_info.pPoolSizes = pool_sizes;
		VkDescriptorPool descriptorPool;
		vkCreateDescriptorPool(device, &pool_info, VK_NULL_HANDLE, &descriptorPool);
		return descriptorPool;
	}
	VkRenderPass CreateRenderPass(TinyImageFormat format, VkDevice device) {
		VkAttachmentDescription colorAttachment = {};
		colorAttachment.format = (VkFormat)TinyImageFormat_ToVkFormat(format);
		colorAttachment.samples = VK_SAMPLE_COUNT_1_BIT;
		colorAttachment.loadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE;
		colorAttachment.storeOp = VK_ATTACHMENT_STORE_OP_STORE;
		colorAttachment.stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE;
		colorAttachment.stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE;
		colorAttachment.initialLayout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;
		colorAttachment.finalLayout = VK_IMAGE_LAYOUT_PRESENT_SRC_KHR;

		VkAttachmentDescription stencilAttachment = {};
		stencilAttachment.format = VK_FORMAT_S8_UINT;
		stencilAttachment.samples = VK_SAMPLE_COUNT_1_BIT;
		stencilAttachment.loadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE;
		stencilAttachment.storeOp = VK_ATTACHMENT_STORE_OP_DONT_CARE;
		stencilAttachment.stencilLoadOp = VK_ATTACHMENT_LOAD_OP_CLEAR;
		stencilAttachment.stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE;
		stencilAttachment.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;
		stencilAttachment.finalLayout = VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL;

		VkAttachmentReference colorAttachmentRef = {};
		colorAttachmentRef.attachment = 0;
		colorAttachmentRef.layout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;

		VkSubpassDescription subpass = {};
		subpass.pipelineBindPoint = VK_PIPELINE_BIND_POINT_GRAPHICS;
		subpass.colorAttachmentCount = 1;
		subpass.pColorAttachments = &colorAttachmentRef;

		VkSubpassDependency dependency = {};
		dependency.srcSubpass = VK_SUBPASS_EXTERNAL;
		dependency.dstSubpass = 0;
		dependency.srcAccessMask = VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT;
		dependency.srcStageMask = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT;
		dependency.dstStageMask = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT |
			VK_PIPELINE_STAGE_EARLY_FRAGMENT_TESTS_BIT;
		dependency.dstAccessMask = VK_ACCESS_COLOR_ATTACHMENT_READ_BIT | VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT |
			VK_ACCESS_DEPTH_STENCIL_ATTACHMENT_WRITE_BIT | VK_ACCESS_DEPTH_STENCIL_ATTACHMENT_READ_BIT;
		VkAttachmentDescription attachmentList[2] = { colorAttachment , stencilAttachment };
		VkRenderPassCreateInfo info = {};
		info.sType = VK_STRUCTURE_TYPE_RENDER_PASS_CREATE_INFO;
		info.attachmentCount = 2;
		info.pAttachments = attachmentList;
		info.subpassCount = 1;
		info.pSubpasses = &subpass;
		info.dependencyCount = 1;
		info.pDependencies = &dependency;
		VkRenderPass renderPass;
		if (vkCreateRenderPass(device, &info, VK_NULL_HANDLE, &renderPass) != VK_SUCCESS)
			throw std::runtime_error("Could not create Dear ImGui's render pass");
		return renderPass;
	}
	void VulkanImguiEditor::Initialize()
	{
		VulkanAPI* API = VulkanAPI::Ptr();
		VulkanWindow* window = VulkanWindow::Ptr();
		Backend& backend = API->GetBackend();
		Queue* pQueue = backend.GetDevice().GetQueue(Queue::RenderQueue);
		VkDescriptorPool descriptorPool = CreateDescriptorPool(backend.GetDevice().Ptr());
		TextureDesc surface = API->context.surface;
		VkRenderPass renderPass = CreateRenderPass(surface.format, backend.GetDevice().Ptr());

		ImGui_ImplVulkan_InitInfo init_info = {};
		init_info.Instance = backend.GetInstance().Ptr();
		init_info.PhysicalDevice = backend.GetDevice().GetPhysical();
		init_info.Device = backend.GetDevice().Ptr();
		init_info.QueueFamily = pQueue->QueueFamilyIndex();
		init_info.Queue = pQueue->Ptr();
		init_info.DescriptorPool = descriptorPool;
		init_info.MinImageCount = 2;
		init_info.ImageCount = API->context.frameCount;
		init_info.RenderPass = renderPass;
		init_info.PipelineCache = VK_NULL_HANDLE;
		init_info.Subpass = 0;
		init_info.MSAASamples = VK_SAMPLE_COUNT_1_BIT;
		init_info.Allocator = VK_NULL_HANDLE;
		ImGui_ImplVulkan_Init(&init_info);
		API->SetRenderPassInfo(ImguiPassName, renderPass);
		EventSystem::Ptr()->BeginRenderFrame.Subscribe(&VulkanImguiEditor::OnBeginRenderFrame, this);
	}
	void VulkanImguiEditor::Finalize()
	{

	}
	ImTextureID VulkanImguiEditor::AddTexture(ImageViewPtr imageview, SamplerPtr sampler, ResourceState state)
	{
		VkDescriptorSet descriptorSet = ImGui_ImplVulkan_AddTexture((VkSampler)sampler, (VkImageView)imageview, vkApiGetImageLayout(state));
		return reinterpret_cast<ImTextureID>(descriptorSet);
	}
	void VulkanImguiEditor::Render(FrameGraph& graph, RenderEditorContext& ctx)
	{
		for (auto win : mWindows)
		{
			win->Draw(graph, ctx);
		}
	}
	void VulkanImguiEditor::OnBeginRenderFrame(FrameGraph& graph, uint32_t frame)
	{
		graph.mIsRenderEditorSurface = gEngineConfig.IsRenderEditorSurface;
		if (gEngineConfig.IsRenderEditorSurface) {
			graph.mEditorSurfaceID = graph.mSurfaceID;
			graph.mSurfaceID = graph.GetTextureID(FrameGraph::NameEditorSurface, frame);
		}
		graph.AddRenderPass<VulkanImguiEditor>();
	}
	void VulkanImguiEditor::Setup(FrameGraph& graph, RenderPassBuilder& builder)
	{
		TextureDesc stencil = graph.GetRenderSurface();
		stencil.id = 0;
		stencil.format = TinyImageFormat_S8_UINT;
		stencil.image = nullptr;
		stencil.usage = TextureUsage::STENCIL_ATTACHMENT | TextureUsage::DEPTH_ATTACHMENT;
		graph.AcquireTexture(stencil);
		builder.Name(ImguiPassName)
			.Type(RenderPassNodeType::Imgui, RenderPassNodeFlag::Output)
			.Write(graph.GetRenderSurface(), ResourceState::COLOR_ATTACHMENT)
			.Write(stencil, ResourceState::DEPTH_ATTACHMENT);
		if (gEngineConfig.IsRenderEditorSurface) {
			builder.Read(graph.GetSurface(), ResourceState::READ_ONLY);
		}
	}
	void VulkanImguiEditor::Execute(FrameGraph& graph, RenderPassContext& context)
	{
		auto& surface = graph.GetRenderSurface();
		context->SetViewport(0.0f, 0.0f, (float)surface.width, (float)surface.height, 0.f, 1.f);
		context->SetScissor(0, 0, surface.width, surface.height);

		graph.GetRenderSurface().state = ResourceState::PRESENT;
		ImGui_ImplVulkan_NewFrame();
		ImGui_ImplSDL2_NewFrame();
		ImGui::NewFrame();

		VulkanImguiEditor* editor = VulkanImguiEditor::Ptr();
		RenderEditorContext editorContext{.editor = editor, .frame = context->frame, .frameCount = context->frameCount };
		editor->Render(graph, editorContext);

		ImGui::Render();
		VulkanContext& ctx = *(VulkanContext*)context.parent;
		ImGui_ImplVulkan_RenderDrawData(ImGui::GetDrawData(), ctx.command);
		ImGuiIO& io = ImGui::GetIO();
		// 更新并渲染平台窗口
		if (io.ConfigFlags & ImGuiConfigFlags_ViewportsEnable)
		{
			ImGui::UpdatePlatformWindows();
			ImGui::RenderPlatformWindowsDefault(nullptr, ctx.command);
		}
	}
}