#include "CameraVideoEncoder.h"
#include "VulkanRHIPrivate.h"
#include "Runtime/AVEncoder/Public/VideoEncoderFactory.h"
#include "Runtime/Engine/Public/ScreenRendering.h"
#include "Runtime/RenderCore/Public/CommonRenderResources.h"
#include "Runtime/RHI/Public/RHIGPUReadback.h"
#include "Utility/DriverSimUtility.h"

void OnEncodedPacket(FCameraVideoEncoder* Encoder, uint32 InCameraIndex, uint32 InLayerIndex, const AVEncoder::FVideoEncoderInputFrame* InFrame, const AVEncoder::FCodecPacket& InPacket)
{
	if (!Encoder->bEncoding)
	{
		return;
	}

	Encoder->EncodedFrameQueue->Enqueue(MakeShared<FEncodedFrame>(InCameraIndex, InPacket));
}

void FCameraVideoEncoder::Init(uint32 InCameraIndex, TQueue<TSharedPtr<FEncodedFrame>>* InEncodedFrameQueue,
	uint32 InWidth, uint32 InHeight, uint32 InEncodeSpan,
	uint32 FrameRate, EEncodeQuality InEncodeQuality)
{
	bEncoding = false;
	EncodeQuality = InEncodeQuality;
	CameraIndex = InCameraIndex;
	Width = InWidth;
	Height = InHeight;
	EncodeSpanCounter = 0;
	EncodeSpan = InEncodeSpan;
	EncodedFrameQueue = InEncodedFrameQueue;

	bCudaSupported = FModuleManager::Get().IsModuleLoaded("CUDA") &&
		FModuleManager::GetModuleChecked<FCUDAModule>("CUDA").IsAvailable();
	if (!bCudaSupported)
	{
		DSLog(FString::Printf(TEXT("Failed to find cuda module!")), ELogLevel::Warning, ELogTag::Sensor);
		return;
	}

	// 创建h264视频编码器
	auto& Available = AVEncoder::FVideoEncoderFactory::Get().GetAvailable();
	if (Available.Num() == 0)
	{
		DSLog(FString::Printf(TEXT("Failed to find available encoder!")), ELogLevel::Fatal, ELogTag::Sensor);
		return;
	}

	bEncoding = true;
	bFrontFrame = true;

	// 如果已经完成初始化，清理帧缓存
	if (VideoEncoderInput)
	{
		return;
	}

	// 初始化视频编码输入
	FindVideoEncoderInput();

	// 设置视频编码属性
	AVEncoder::FVideoEncoder::FLayerConfig VideoEncoderConfig;
	VideoEncoderConfig.Width = Width;
	VideoEncoderConfig.Height = Height;
	VideoEncoderConfig.MaxFramerate = FrameRate;
	uint32 BitRate = EncodeQuality == EEncodeQuality::EQ_Low ? 1000000 : (EncodeQuality == EQ_Medium ? 5000000 : 30000000);
	VideoEncoderConfig.MaxBitrate = BitRate;
	VideoEncoderConfig.TargetBitrate = BitRate;
	VideoEncoder = AVEncoder::FVideoEncoderFactory::Get().Create(Available[0].ID, VideoEncoderInput, VideoEncoderConfig);

	if (!VideoEncoder)
	{
		DSLog(FString::Printf(TEXT("Failed to create %dth video encoder!"), InCameraIndex), ELogLevel::Fatal, ELogTag::Debug);
		return;
	}
	DSLog(FString::Printf(TEXT("Create %dth video encoder"), InCameraIndex), ELogLevel::Info, ELogTag::Sensor);

	// 注册视频编码事件
	VideoEncoder->SetOnEncodedPacket([this](uint32 InLayerIndex, const AVEncoder::FVideoEncoderInputFrame* InFrame, const AVEncoder::FCodecPacket& InPacket)
	{ OnEncodedPacket(this, CameraIndex, InLayerIndex, InFrame, InPacket); });

	// 初始化TextureReadback回读数据
	TextureReadback = MakeShared<class FRHIGPUTextureReadback>(TEXT("GPUTextureReadback"));
	TextureData = new uint8[Width * Height *4];
	bEnqueueCopied = false;
	bReadbacked = false;
}

void FCameraVideoEncoder::Stop()
{
	bEncoding = false;
}

void FCameraVideoEncoder::Shutdown()
{
	if (!bCudaSupported)
	{
		return;
	}

	ClearFrames();
	VideoEncoderInput.Reset();
	VideoEncoder->ClearOnEncodedPacket();
	VideoEncoder->Shutdown();
	VideoEncoder.Reset();
	TextureReadback.Reset();

	SAFE_DELETE(TextureData);
}

void FCameraVideoEncoder::Tick(float DeltaTime)
{
	if (EncodeQuality != EQ_Lossless || !bEncoding)
	{
		return;
	}

	ENQUEUE_RENDER_COMMAND(Readback)(
		[this](FRHICommandList& RHICmdList)
	{
		if (bEnqueueCopied && TextureReadback && TextureReadback->IsReady())
		{
			int32 ActualWidth, ActualHeight;
			void* DataPtr = TextureReadback->Lock(ActualWidth, ActualHeight);
			UDriverSimUtility::CopyMapStagingSurfaceData(DataPtr, TextureData,
				ActualWidth, ActualHeight, Width, Height);
			TextureReadback->Unlock();

			bEnqueueCopied = false;
			bReadbacked = true;
		}
	});

	if (bReadbacked)
	{
		EncodedFrameQueue->Enqueue(MakeShared<FEncodedFrame>(CameraIndex, TextureData, Width * Height * 4));
		bReadbacked = false;
	}
}

void FCameraVideoEncoder::Encode(const FTexture2DRHIRef& Texture)
{
	if (!bCudaSupported)
	{
		return;
	}

	FRHICommandListImmediate& RHICmdList = GRHICommandList.GetImmediateCommandList();

	if (EncodeQuality == EQ_Lossless)
	{
		TextureReadback->EnqueueCopy(RHICmdList, Texture);
		RHICmdList.Transition(FRHITransitionInfo(Texture, ERHIAccess::Unknown, ERHIAccess::ReadOnlyMask));
		bEnqueueCopied = true;
	}
	else
	{
		AVEncoder::FVideoEncoderInputFrame* InputFrame = ObtainInputFrame();
		if (!InputFrame)
		{
			return;
		}

		const int64 TimestampUs = UDriverSimUtility::GetTimeStampUS();
		InputFrame->SetTimestampUs(TimestampUs);
		InputFrame->SetTimestampRTP(TimestampUs);
		RHICmdList.Transition(FRHITransitionInfo(Texture, ERHIAccess::Unknown, ERHIAccess::ReadOnlyMask));

		// Actual texture copy (i.e the actual "capture")
		CopyTexture(Texture, InputFrameTextureMap[InputFrame]);

		AVEncoder::FVideoEncoder::FEncodeOptions EncodeOptions;
		if (bFrontFrame)
		{
			EncodeOptions.bForceKeyFrame = true;
			bFrontFrame = false;
		}
		else if (EncodeSpanCounter >= EncodeSpan)
		{
			EncodeOptions.bForceKeyFrame = true;
			EncodeSpanCounter = 0;
		}
		else
		{
			EncodeOptions.bForceKeyFrame = false;
		}
		EncodeSpanCounter++;

		VideoEncoder->Encode(InputFrame, EncodeOptions);

		InputFrame->Release();
	}
}

void FCameraVideoEncoder::FindVideoEncoderInput()
{
	if (GDynamicRHI)
	{
		FString RHIName = GDynamicRHI->GetName();

		bool bIsRHISupported = false;

		if (RHIName == TEXT("Vulkan"))
		{
			if (IsRHIDeviceAMD())
			{
				FVulkanDynamicRHI* DynamicRHI = static_cast<FVulkanDynamicRHI*>(GDynamicRHI);
				AVEncoder::FVulkanDataStruct VulkanData = { DynamicRHI->GetInstance(), DynamicRHI->GetDevice()->GetPhysicalHandle(), DynamicRHI->GetDevice()->GetInstanceHandle() };

				VideoEncoderInput = AVEncoder::FVideoEncoderInput::CreateForVulkan(&VulkanData, Width, Height, true);
				bIsRHISupported = true;
			}
			else if (IsRHIDeviceNVIDIA())
			{
				VideoEncoderInput = AVEncoder::FVideoEncoderInput::CreateForCUDA(FModuleManager::GetModuleChecked<FCUDAModule>("CUDA").GetCudaContext(), Width, Height, true);
				bIsRHISupported = true;
			}
		}
#if PLATFORM_WINDOWS
		else if (RHIName == TEXT("D3D11"))
		{
			VideoEncoderInput = AVEncoder::FVideoEncoderInput::CreateForD3D11(GDynamicRHI->RHIGetNativeDevice(), Width, Height, true, IsRHIDeviceAMD());
			bIsRHISupported = true;
		}
		else if (RHIName == TEXT("D3D12"))
		{
			VideoEncoderInput = AVEncoder::FVideoEncoderInput::CreateForD3D12(GDynamicRHI->RHIGetNativeDevice(), Width, Height, true, IsRHIDeviceNVIDIA());
			bIsRHISupported = true;
		}
#endif

		if (!bIsRHISupported)
		{
			DSLog(FString::Printf(TEXT("Failed to find a proper video encoder input!")), ELogLevel::Fatal, ELogTag::Sensor);
		}
		if (VideoEncoderInput)
		{
			VideoEncoderInput->SetResolution(Width, Height);
			VideoEncoderInput->SetMaxNumBuffers(3);
		}
	}
}

AVEncoder::FVideoEncoderInputFrame* FCameraVideoEncoder::ObtainInputFrame()
{
	AVEncoder::FVideoEncoderInputFrame* InputFrame = VideoEncoderInput->ObtainInputFrame();
	if (InputFrame == nullptr)
		return nullptr;

	if (!InputFrameTextureMap.Contains(InputFrame))
	{
#if PLATFORM_WINDOWS
		FString RHIName = GDynamicRHI->GetName();
		if (RHIName == TEXT("D3D11"))
		{
			FRHIResourceCreateInfo CreateInfo(TEXT("VideoCapturerBackBuffer"));
			FTexture2DRHIRef Texture = GDynamicRHI->RHICreateTexture2D(Width, Height, EPixelFormat::PF_B8G8R8A8, 1, 1, TexCreate_Shared | TexCreate_RenderTargetable, ERHIAccess::CopyDest, CreateInfo);
			InputFrame->SetTexture((ID3D11Texture2D*)Texture->GetNativeResource(), [this, InputFrame](ID3D11Texture2D* NativeTexture) { InputFrameTextureMap.Remove(InputFrame); });
			InputFrameTextureMap.Add(InputFrame, Texture);
		}
		else if (RHIName == TEXT("D3D12"))
		{
			FRHIResourceCreateInfo CreateInfo(TEXT("VideoCapturerBackBuffer"));
			FTexture2DRHIRef Texture = GDynamicRHI->RHICreateTexture2D(Width, Height, EPixelFormat::PF_B8G8R8A8, 1, 1, TexCreate_Shared | TexCreate_RenderTargetable, ERHIAccess::CopyDest, CreateInfo);
			InputFrame->SetTexture((ID3D12Resource*)Texture->GetNativeResource(), [this, InputFrame](ID3D12Resource* NativeTexture) { InputFrameTextureMap.Remove(InputFrame); });
			InputFrameTextureMap.Add(InputFrame, Texture);
		}
		else if (RHIName == TEXT("Vulkan"))
#endif // PLATFORM_WINDOWS
		{
			if (IsRHIDeviceNVIDIA())
			{
				FRHIResourceCreateInfo CreateInfo(TEXT("VideoCapturerBackBuffer"));

				// Create a texture that can be exposed to external memory
				FTexture2DRHIRef Texture =
					GDynamicRHI->RHICreateTexture2D(Width, Height, EPixelFormat::PF_B8G8R8A8, 1, 1, TexCreate_Shared | TexCreate_RenderTargetable | TexCreate_UAV | TexCreate_External, ERHIAccess::Present, CreateInfo);

				FVulkanTexture2D* VulkanTexture = static_cast<FVulkanTexture2D*>(Texture.GetReference());
				FVulkanDynamicRHI* VulkanDynamicRHI = static_cast<FVulkanDynamicRHI*>(GDynamicRHI);
				VkDevice device = VulkanDynamicRHI->GetDevice()->GetInstanceHandle();

				// Get the CUarray to that textures memory making sure the clear it when done
				// While this operation is safe (and unavoidable) C4191 has been enabled and this will trigger an error with warnings as errors
#pragma warning(push)
#pragma warning(disable : 4191)
#if PLATFORM_WINDOWS
				VkMemoryGetWin32HandleInfoKHR vkMemoryGetWin32HandleInfoKHR = {};
				vkMemoryGetWin32HandleInfoKHR.sType = VK_STRUCTURE_TYPE_MEMORY_GET_WIN32_HANDLE_INFO_KHR;
				vkMemoryGetWin32HandleInfoKHR.pNext = NULL;
				vkMemoryGetWin32HandleInfoKHR.memory = VulkanTexture->Surface.GetAllocationHandle();
				vkMemoryGetWin32HandleInfoKHR.handleType = VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_WIN32_BIT_KHR;

				void* handle;
				PFN_vkGetMemoryWin32HandleKHR fpGetMemoryWin32HandleKHR = (PFN_vkGetMemoryWin32HandleKHR)VulkanRHI::vkGetDeviceProcAddr(device, "vkGetMemoryWin32HandleKHR");
				VERIFYVULKANRESULT(fpGetMemoryWin32HandleKHR(device, &vkMemoryGetWin32HandleInfoKHR, &handle));
#elif PLATFORM_LINUX
				// Generate VkMemoryGetFdInfoKHR
				VkMemoryGetFdInfoKHR vkMemoryGetFdInfoKHR = {};
				vkMemoryGetFdInfoKHR.sType = VK_STRUCTURE_TYPE_MEMORY_GET_FD_INFO_KHR;
				vkMemoryGetFdInfoKHR.pNext = NULL;
				vkMemoryGetFdInfoKHR.memory = VulkanTexture->Surface.GetAllocationHandle();
				vkMemoryGetFdInfoKHR.handleType = VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_FD_BIT_KHR;

				int fd;
				PFN_vkGetMemoryFdKHR fpGetMemoryFdKHR = (PFN_vkGetMemoryFdKHR)VulkanRHI::vkGetDeviceProcAddr(device, "vkGetMemoryFdKHR");
				VERIFYVULKANRESULT(fpGetMemoryFdKHR(device, &vkMemoryGetFdInfoKHR, &fd));
#endif
#pragma warning(pop)

				FCUDAModule::CUDA().cuCtxPushCurrent(FModuleManager::GetModuleChecked<FCUDAModule>("CUDA").GetCudaContext());

				CUexternalMemory mappedExternalMemory = nullptr;

				// generate a cudaExternalMemoryHandleDesc
				CUDA_EXTERNAL_MEMORY_HANDLE_DESC cudaExtMemHandleDesc = {};
#if PLATFORM_WINDOWS
				cudaExtMemHandleDesc.type = CU_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_WIN32;
				cudaExtMemHandleDesc.handle.win32.handle = handle;
				cudaExtMemHandleDesc.handle.win32.name = nullptr;
#elif PLATFORM_LINUX
				cudaExtMemHandleDesc.type = CU_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_FD;
				cudaExtMemHandleDesc.handle.fd = fd;
#endif
				cudaExtMemHandleDesc.size = VulkanTexture->Surface.GetAllocationOffset() + VulkanTexture->Surface.GetMemorySize();

				// import external memory
				auto result = FCUDAModule::CUDA().cuImportExternalMemory(&mappedExternalMemory, &cudaExtMemHandleDesc);
				if (result != CUDA_SUCCESS)
				{
					UE_LOG(LogClass, Error, TEXT("Failed to import external memory from vulkan error: %d"), result);
				}

				CUmipmappedArray mappedMipArray = nullptr;
				CUarray mappedArray = nullptr;

				CUDA_EXTERNAL_MEMORY_MIPMAPPED_ARRAY_DESC mipmapDesc = {};
				mipmapDesc.numLevels = 1;
				mipmapDesc.offset = VulkanTexture->Surface.GetAllocationOffset();
				mipmapDesc.arrayDesc.Width = Texture->GetSizeX();
				mipmapDesc.arrayDesc.Height = Texture->GetSizeY();
				mipmapDesc.arrayDesc.Depth = 0;
				mipmapDesc.arrayDesc.NumChannels = 4;
				mipmapDesc.arrayDesc.Format = CU_AD_FORMAT_UNSIGNED_INT8;
				mipmapDesc.arrayDesc.Flags = CUDA_ARRAY3D_SURFACE_LDST | CUDA_ARRAY3D_COLOR_ATTACHMENT;

				// get the CUarray from the external memory
				result = FCUDAModule::CUDA().cuExternalMemoryGetMappedMipmappedArray(&mappedMipArray, mappedExternalMemory, &mipmapDesc);
				if (result != CUDA_SUCCESS)
				{
					UE_LOG(LogClass, Error, TEXT("Failed to bind mipmappedArray error: %d"), result);
				}

				// get the CUarray from the external memory
				CUresult mipMapArrGetLevelErr = FCUDAModule::CUDA().cuMipmappedArrayGetLevel(&mappedArray, mappedMipArray, 0);
				if (mipMapArrGetLevelErr != CUDA_SUCCESS)
				{
					UE_LOG(LogClass, Error, TEXT("Failed to bind to mip 0."));
				}

				FCUDAModule::CUDA().cuCtxPopCurrent(NULL);

				InputFrame->SetTexture(mappedArray,
					[this, mappedArray, mappedMipArray, mappedExternalMemory, InputFrame](CUarray NativeTexture)
				{
					// free the cuda types
					FCUDAModule::CUDA().cuCtxPushCurrent(FModuleManager::GetModuleChecked<FCUDAModule>("CUDA").GetCudaContext());

					if (mappedArray)
					{
						auto result = FCUDAModule::CUDA().cuArrayDestroy(mappedArray);
						if (result != CUDA_SUCCESS)
						{
							UE_LOG(LogClass, Error, TEXT("Failed to destroy mappedArray: %d"), result);
						}
					}

					if (mappedMipArray)
					{
						auto result = FCUDAModule::CUDA().cuMipmappedArrayDestroy(mappedMipArray);
						if (result != CUDA_SUCCESS)
						{
							UE_LOG(LogClass, Error, TEXT("Failed to destroy mappedMipArray: %d"), result);
						}
					}

					if (mappedExternalMemory)
					{
						auto result = FCUDAModule::CUDA().cuDestroyExternalMemory(mappedExternalMemory);
						if (result != CUDA_SUCCESS)
						{
							UE_LOG(LogClass, Error, TEXT("Failed to destroy mappedExternalMemoryArray: %d"), result);
						}
					}

					FCUDAModule::CUDA().cuCtxPopCurrent(NULL);

					// finally remove the input frame
					InputFrameTextureMap.Remove(InputFrame);
				});
				InputFrameTextureMap.Add(InputFrame, Texture);
			}
			else if (IsRHIDeviceAMD())
			{
				FRHIResourceCreateInfo CreateInfo(TEXT("VideoCapturerBackBuffer"));
				FTexture2DRHIRef Texture =
					GDynamicRHI->RHICreateTexture2D(Width, Height, EPixelFormat::PF_B8G8R8A8, 1, 1, TexCreate_Shared | TexCreate_RenderTargetable | TexCreate_UAV, ERHIAccess::Present, CreateInfo);
				FVulkanTexture2D* VulkanTexture = static_cast<FVulkanTexture2D*>(Texture.GetReference());

				InputFrame->SetTexture(VulkanTexture->Surface.Image, [this, InputFrame](VkImage NativeTexture) { InputFrameTextureMap.Remove(InputFrame); });

				InputFrameTextureMap.Add(InputFrame, Texture);
			}
		}

		UE_LOG(LogClass, Log, TEXT("%d input frame buffer currently allocated."), InputFrameTextureMap.Num());
	}

	return InputFrame;
}

void FCameraVideoEncoder::CopyTexture(const FTexture2DRHIRef& SourceTexture, FTexture2DRHIRef& DestinationTexture) const
{
	FRHICommandListImmediate& RHICmdList = FRHICommandListExecutor::GetImmediateCommandList();

	IRendererModule* RendererModule = &FModuleManager::GetModuleChecked<IRendererModule>("Renderer");

	// #todo-renderpasses there's no explicit resolve here? Do we need one?
	FRHIRenderPassInfo RPInfo(DestinationTexture, ERenderTargetActions::Load_Store);

	RHICmdList.BeginRenderPass(RPInfo, TEXT("CopyBackbuffer"));

	{
		RHICmdList.SetViewport(0, 0, 0.0f, DestinationTexture->GetSizeX(), DestinationTexture->GetSizeY(), 1.0f);

		FGraphicsPipelineStateInitializer GraphicsPSOInit;
		RHICmdList.ApplyCachedRenderTargets(GraphicsPSOInit);
		GraphicsPSOInit.BlendState = TStaticBlendState<>::GetRHI();
		GraphicsPSOInit.RasterizerState = TStaticRasterizerState<>::GetRHI();
		GraphicsPSOInit.DepthStencilState = TStaticDepthStencilState<false, CF_Always>::GetRHI();

		// New engine version...
		FGlobalShaderMap* ShaderMap = GetGlobalShaderMap(GMaxRHIFeatureLevel);
		TShaderMapRef<FScreenVS> VertexShader(ShaderMap);
		TShaderMapRef<FScreenPS> PixelShader(ShaderMap);

		GraphicsPSOInit.BoundShaderState.VertexDeclarationRHI = GFilterVertexDeclaration.VertexDeclarationRHI;
		GraphicsPSOInit.BoundShaderState.VertexShaderRHI = VertexShader.GetVertexShader();
		GraphicsPSOInit.BoundShaderState.PixelShaderRHI = PixelShader.GetPixelShader();

		GraphicsPSOInit.PrimitiveType = PT_TriangleList;

		SetGraphicsPipelineState(RHICmdList, GraphicsPSOInit);

		if (DestinationTexture->GetSizeX() != SourceTexture->GetSizeX() || DestinationTexture->GetSizeY() != SourceTexture->GetSizeY())
		{
			PixelShader->SetParameters(RHICmdList, TStaticSamplerState<SF_Bilinear>::GetRHI(), SourceTexture);
		}
		else
		{
			PixelShader->SetParameters(RHICmdList, TStaticSamplerState<SF_Point>::GetRHI(), SourceTexture);
		}

		RendererModule->DrawRectangle(RHICmdList, 0, 0,                // Dest X, Y
			DestinationTexture->GetSizeX(),  // Dest Width
			DestinationTexture->GetSizeY(),  // Dest Height
			0, 0,                            // Source U, V
			1, 1,                            // Source USize, VSize
			DestinationTexture->GetSizeXY(), // Target buffer size
			FIntPoint(1, 1),                 // Source texture size
			VertexShader, EDRF_Default);
	}

	RHICmdList.EndRenderPass();
}

void FCameraVideoEncoder::ClearFrames()
{
	//for (auto& Iter : InputFrameTextureMap)
	//{
	//	Iter.Key->Release();
	//	Iter.Value.SafeRelease();
	//}
	InputFrameTextureMap.Empty();
	VideoEncoderInput->Flush();
}