// Fill out your copyright notice in the Description page of Project Settings.


#include "PortaFrameGrabber.h"
#include "IImageWrapper.h"
#include "IImageWrapperModule.h"
#include "ImageUtils.h"
#include "PortaFrameGrabberSettings.h"
#include "PortaSequenceSubsystem.h"
#include "PortaConnectionSubsystem.h"
#include "PortaLevelSubsystem.h"
#include "Engine/GameViewportClient.h"
#include "Misc/FileHelper.h"
#include "LevelSequencePlayer.h"
#include "PortaUtils.h"
#include "Engine/TextureRenderTarget2D.h"
#include "UnrealClient.h"
#include "Modules/ModuleManager.h"
#include "PortaInterface.h"
#include "Misc/Base64.h"

// Called when the game starts or when spawned
void UPortaFrameGrabber::Initialize(FSubsystemCollectionBase& Collection)
{
	screenshotQueueSize = 0;
	shouldProccessQueue = false;

	GEngine->GameViewport->OnScreenshotCaptured().AddUObject(this, &UPortaFrameGrabber::AcceptScreenshot);
}

TStatId UPortaFrameGrabber::GetStatId() const
{
	RETURN_QUICK_DECLARE_CYCLE_STAT(UPortaFrameGrabber, STATGROUP_Tickables);
}

bool UPortaFrameGrabber::IsTickable() const
{
	return !IsDefaultSubobject() && IsValid(this) && !WaitingOnLevel && (!screenshotQueue.IsEmpty() || currentReq);
}

void UPortaFrameGrabber::Tick(float DeltaTime)
{
	if (ScreenshotDelay > 0)
		ScreenshotDelay -= DeltaTime;
	else
		ProcessQueue();
}

void UPortaFrameGrabber::RequestScreenshotStatic(const FString& BatchRequest, const TArray<FString>& Levels, const TArray<FString>& Sequences, FScreenshotDelegate Callback)
{
	PortaUtils::CallWorldSubsystemMember<UPortaFrameGrabber, const FString&, const TArray<FString>&, const TArray<FString>&, FScreenshotDelegate>(&UPortaFrameGrabber::RequestScreenshot, BatchRequest, Levels, Sequences, Callback);

}

void UPortaFrameGrabber::RequestScreenshot(const FString& BatchRequest, const TArray<FString>& Levels, const TArray<FString>& Sequences, FScreenshotDelegate Callback)
{
	auto req = new FPortaScreenshotRequest();
	req->BatchRequest = BatchRequest;
	req->Levels = Levels;
	req->Sequences = Sequences;
	req->Callback = Callback;


	screenshotQueue.Enqueue(req);
	screenshotQueueSize++;
}

void UPortaFrameGrabber::ProcessQueue()
{
	if (currentReq)
	{
		FScreenshotRequest::RequestScreenshot(false);
		return;
	}
	if (screenshotQueueSize >= QUEUE_THRESHOLD)
		shouldProccessQueue = true;

	if (screenshotQueue.IsEmpty())
		shouldProccessQueue = false;

	if (!screenshotQueue.IsEmpty() && shouldProccessQueue)
	{
		screenshotQueue.Dequeue(currentReq);
		screenshotQueueSize--;

		if (!currentReq->Levels.IsEmpty())
		{
			WaitingOnLevel = true;
			if (currentReq->Levels.Num() > 1)
			{
				UE_LOG(LogPorta, Error, TEXT("Screenshot grabbing does not support loading more than one streaming level. Only the first provided will be loaded."));
			}
			FLevelLoaded Delegate;
			Delegate.BindUFunction(this, TEXT("OnLevelLoaded"));
			UPortaLevelSubsystem::LoadStreamingLevelStatic(currentReq->Levels[0], true, Delegate);
		}
		else
		{
			ScreenshotDelay = 0.3f;
			RunBatchAndSequences();
		}

		shouldProccessQueue = false;
	}
}

void UPortaFrameGrabber::OnLevelLoaded()
{
	WaitingOnLevel = false;
	ScreenshotDelay = 0.3f;
	RunBatchAndSequences();
}

void UPortaFrameGrabber::RunBatchAndSequences()
{
	if (!currentReq)
		return;

	if (!currentReq->BatchRequest.IsEmpty())
	{
		auto Connect = UPortaConnectionSubsystem::Get();
		Connect->RunBatch(currentReq->BatchRequest);
	}

	for (auto& SequencePath : currentReq->Sequences)
	{
		UPortaSequenceSubsystem::ShowEndOfSequenceStatic(SequencePath);
	}
}


void UPortaFrameGrabber::AcceptScreenshot(int32 InSizeX, int32 InSizeY, const TArray<FColor>& InImageData)
{
	auto settings = GetDefault<UPortaFrameGrabberSettings>();

	if (!currentReq)
	{
		// something isn't right.
		return;
	}

	IImageWrapperModule& ImageWrapperModule = FModuleManager::LoadModuleChecked<IImageWrapperModule>(FName("ImageWrapper"));
	TSharedPtr<IImageWrapper> ImageWrapper = ImageWrapperModule.CreateImageWrapper(EImageFormat::JPEG);

	if (settings->ShouldUseCustomResolution)
	{
		TArray<FColor> resizedImageData;
		resizedImageData.Reserve(settings->CustomResolutionWidth * settings->CustomResolutionHeight);
		FImageUtils::ImageResize(InSizeX, InSizeY, InImageData, settings->CustomResolutionWidth, settings->CustomResolutionHeight, resizedImageData, false);

		ImageWrapper->SetRaw(&resizedImageData[0], resizedImageData.Num() * sizeof(FColor), settings->CustomResolutionWidth, settings->CustomResolutionHeight, ERGBFormat::BGRA, 8);
	}
	else
	{
		ImageWrapper->SetRaw(&InImageData[0], InImageData.Num() * sizeof(FColor), InSizeX, InSizeY, ERGBFormat::BGRA, 8);
	}

	auto CompressedImage = ImageWrapper->GetCompressed(settings->JPEGQuality);
	FString imgData64 = FBase64::Encode(CompressedImage.GetData(), CompressedImage.Num());
	currentReq->Callback.Execute(imgData64);
	shouldProccessQueue = true;
	delete currentReq;
	currentReq = nullptr;
}