// Copyright Epic Games, Inc. All Rights Reserved.

#pragma once

#include "VideoEncoderInput.h"
#include "Containers/Queue.h"
#include "Misc/ScopeLock.h"
#include "Templates/RefCounting.h"

#if WITH_CUDA
#include "CudaModule.h"
#endif

// HACK (M84FIX) need to break these dependencies
#if PLATFORM_WINDOWS
struct ID3D11DeviceContext;
#endif

struct VkDevice_T;

namespace AVEncoder
{
class FVideoEncoderInputFrameImpl;

class FVideoEncoderInputImpl : public FVideoEncoderInput
{
public:
	FVideoEncoderInputImpl() = default;
	virtual ~FVideoEncoderInputImpl();

	// --- properties

	uint32 GetNumActiveFrames() const { FScopeLock Guard(&ProtectFrames); return ActiveFrames.Num(); }
	bool GetHasFreeFrames() const { FScopeLock Guard(&ProtectFrames); return !AvailableFrames.IsEmpty(); }

	// --- construct video encoder input based on expected input frame format

	bool SetupForDummy(uint32 InWidth, uint32 InHeight);
	bool SetupForYUV420P(uint32 InWidth, uint32 InHeight);

	// set up for an encoder that encodes a D3D11 texture
	bool SetupForD3D11(void* InApplicationD3DDevice, uint32 InWidth, uint32 InHeight);

	// set up for an encoder that encodes a D3D12 texture in the context of a D3D11 device (i.e. nvenc)
	bool SetupForD3D12(void* InApplicationD3DDevice, uint32 InWidth, uint32 InHeight);

	// set up an encoder that encodes a CUArray in the CUDA context
	bool SetupForCUDA(void* InApplicationContext, uint32 InWidth, uint32 InHeight);

	// set up an encoder that encodes a VkImage in the context of a VkDevice
	bool SetupForVulkan(void* InApplicationContext, uint32 InWidth, uint32 InHeight);

	// --- available encoders

	// get a list of supported video encoders
	const TArray<FVideoEncoderInfo>& GetAvailableEncoders() override;

	// --- encoder input frames - user managed

	// create a user managed buffer
	FVideoEncoderInputFrame* CreateBuffer(OnFrameReleasedCallback InOnFrameReleased) override;
	// destroy user managed buffer
	void DestroyBuffer(FVideoEncoderInputFrame* Buffer) override;

	// --- encoder input frames - managed by this object

	// obtain a video frame that can be used as a buffer for input to a video encoder
	FVideoEncoderInputFrame* ObtainInputFrame() override;

	// release (free) an input frame and make it available for future use
	void ReleaseInputFrame(FVideoEncoderInputFrame* InFrame) override;

	// destroy/release any frames that are not currently in use
	void Flush() override;

	// --- input properties

	// TODO (M84FIX) make compile no matter the platform
#if PLATFORM_WINDOWS
	TRefCountPtr<ID3D11Device> GetD3D11EncoderDevice() const;
	TRefCountPtr<ID3D11Device> ForceD3D11InputFrames();
#endif

#if WITH_CUDA
	CUcontext GetCUDAEncoderContext() const;
#endif

#if PLATFORM_WINDOWS || PLATFORM_LINUX
	VkDevice_T* GetVulkanDevice() const;
#endif

private:

	// collect any encoder that can handle frame format
	void CollectAvailableEncoders();

	TArray<FVideoEncoderInfo>		AvailableEncoders;

	FVideoEncoderInputFrameImpl* CreateFrame();
	void SetupFrameYUV420P(FVideoEncoderInputFrameImpl* Frame);
	void SetupFrameD3D11(FVideoEncoderInputFrameImpl* Frame);
	void SetupFrameD3D12(FVideoEncoderInputFrameImpl* Frame);
	void SetupFrameCUDA(FVideoEncoderInputFrameImpl* Frame);

	struct FFrameInfoDummy
	{
	}								FrameInfoDummy;

	struct FFrameInfoYUV420P
	{
		uint32						StrideY = 0;
		uint32						StrideU = 0;
		uint32						StrideV = 0;
	}								FrameInfoYUV420P;

#if PLATFORM_WINDOWS
	struct FFrameInfoD3D
	{
		TRefCountPtr<ID3D11Device>	EncoderDeviceD3D11;
		TRefCountPtr<ID3D11DeviceContext>	EncoderDeviceContextD3D11;
		TRefCountPtr<ID3D12Device>	EncoderDeviceD3D12;
	}								FrameInfoD3D;
#endif

#if WITH_CUDA
	struct FFrameInfoCUDA
	{
		// HACK (M84FIX)
		CUcontext					EncoderContextCUDA;
	}								FrameInfoCUDA;
#endif

#if PLATFORM_WINDOWS || PLATFORM_LINUX
	struct FFrameInfoVulkan
	{
		VkDevice_T*					VulkanDevice;
	}								FrameInfoVulkan;
#endif

	mutable FCriticalSection				ProtectFrames;
	TQueue<FVideoEncoderInputFrameImpl*>	AvailableFrames;
	TArray<FVideoEncoderInputFrameImpl*>	ActiveFrames;
	using UserManagedFrame = TPair<FVideoEncoderInputFrameImpl*, OnFrameReleasedCallback>;
	TArray<UserManagedFrame>				UserManagedFrames;
};

class FVideoEncoderInputFrameImpl : public FVideoEncoderInputFrame
{
public:
	explicit FVideoEncoderInputFrameImpl(FVideoEncoderInputImpl* InInput);
	FVideoEncoderInputFrameImpl(const FVideoEncoderInputFrameImpl& InCloneFrom) = delete;
	explicit FVideoEncoderInputFrameImpl(const FVideoEncoderInputFrameImpl& InCloneFrom, FCloneDestroyedCallback InCloneDestroyedCallback);
	~FVideoEncoderInputFrameImpl();

	// Clone frame - this will create a copy that references the original until destroyed
	const FVideoEncoderInputFrame* Clone(FCloneDestroyedCallback InCloneDestroyedCallback) const override;

	// Release (decrease reference count) of this input frame
	void Release() const override;

	void SetFormat(EVideoFrameFormat InFormat) { Format = InFormat; }
	void SetWidth(uint32 InWidth) { Width = InWidth; }
	void SetHeight(uint32 InHeight) { Height = InHeight; }

private:
	FVideoEncoderInputImpl*					Input;
	const FVideoEncoderInputFrame*			ClonedReference = nullptr;
	const FCloneDestroyedCallback			OnCloneDestroyed;
};

} /* namespace AVEncoder */