// Copyright 1998-2019 Epic Games, Inc. All Rights Reserved.

#include "GameplayMediaEncoderSample.h"
#include "GameplayMediaEncoderCommon.h"
#include "Misc/AssertionMacros.h"

GAMEPLAYMEDIAENCODER_START

namespace
{
	TRefCountPtr<IMFMediaBuffer> CloneMediaBuffer(const TRefCountPtr<IMFMediaBuffer>& Src)
	{
		uint8* SrcByteBuffer = nullptr;
		DWORD SrcMaxLen = 0;
		DWORD SrcCurrLen = 0;
		verify(SUCCEEDED(Src->Lock(&SrcByteBuffer, &SrcMaxLen, &SrcCurrLen)));

		TRefCountPtr<IMFMediaBuffer> Dest;
		verify(SUCCEEDED(MFCreateMemoryBuffer(SrcCurrLen, Dest.GetInitReference())));

		uint8* DestByteBuffer = nullptr;
		verify(SUCCEEDED(Dest->Lock(&DestByteBuffer, nullptr, nullptr)));
		FMemory::Memcpy(DestByteBuffer, SrcByteBuffer, SrcCurrLen);
		verify(SUCCEEDED(Dest->Unlock()));
		verify(SUCCEEDED(Src->Unlock()));

		verify(SUCCEEDED(Dest->SetCurrentLength(SrcCurrLen)));
		return Dest;
	}
}

bool FGameplayMediaEncoderSample::CreateSample()
{
	Sample = nullptr;
	CHECK_HR(MFCreateSample(Sample.GetInitReference()));
	return true;
}

FTimespan FGameplayMediaEncoderSample::GetTime() const
{
	int64 Time;
	HRESULT hr = Sample->GetSampleTime(&Time);
	verifyf(SUCCEEDED(hr), TEXT("%s"), *GetComErrorDescription(hr));
	return Time;
}

void FGameplayMediaEncoderSample::SetTime(FTimespan Time)
{
	verify(SUCCEEDED(Sample->SetSampleTime(Time.GetTicks())));
}

FTimespan FGameplayMediaEncoderSample::GetDuration() const
{
	int64 Duration;
	HRESULT hr = Sample->GetSampleDuration(&Duration);
	verifyf(SUCCEEDED(hr), TEXT("%s"), *GetComErrorDescription(hr));
	return Duration;
}

void FGameplayMediaEncoderSample::SetDuration(FTimespan Duration)
{
	verify(SUCCEEDED(Sample->SetSampleDuration(Duration.GetTicks())));
}

bool FGameplayMediaEncoderSample::IsVideoKeyFrame() const
{
	return MediaType == EMediaType::Video && MFGetAttributeUINT32(Sample, MFSampleExtension_CleanPoint, 0) != 0;
}

FGameplayMediaEncoderSample FGameplayMediaEncoderSample::Clone() const
{
	FGameplayMediaEncoderSample Res{ MediaType };
	if (!Sample)
	{
		return Res;
	}

	verify(Res.CreateSample());

	DWORD Flags = 0;
	verify(SUCCEEDED(Sample->GetSampleFlags(&Flags)));
	verify(SUCCEEDED(Res.GetSample()->SetSampleFlags(Flags)));

	Res.SetTime(GetTime());
	Res.SetDuration(GetDuration());

	if (IsVideoKeyFrame())
	{
		verify(SUCCEEDED(Res.GetSample()->SetUINT32(MFSampleExtension_CleanPoint, 1)));
	}

	DWORD BufferCount = 0;
	verify(SUCCEEDED(Sample->GetBufferCount(&BufferCount)));

	for (DWORD i = 0; i != BufferCount; ++i)
	{
		TRefCountPtr<IMFMediaBuffer> SrcBuf;
		verify(SUCCEEDED(Sample->GetBufferByIndex(i, SrcBuf.GetInitReference())));

		TRefCountPtr<IMFMediaBuffer> DestBuf = CloneMediaBuffer(SrcBuf);
		verify(SUCCEEDED(Res.GetSample()->AddBuffer(DestBuf)));
	}

	return Res;
}

GAMEPLAYMEDIAENCODER_END