#include "UI/DescentUiShot.h"
#include "Debug/DescentStateRegistry.h"
#include "TheDescent.h"

#include "Engine/Engine.h"
#include "Engine/GameViewportClient.h"
#include "HAL/FileManager.h"
#include "ImageUtils.h"
#include "Misc/FileHelper.h"
#include "Misc/Paths.h"
#include "UnrealClient.h"

namespace
{
	struct FUiShotState
	{
		bool bPending = false;
		FString Path;
		FDelegateHandle Handle;
		TFunction<void(const Descent::Ui::FUiShotResult&)> Done;
	};
	FUiShotState GUiShot;

	void UiShotFinish(const Descent::Ui::FUiShotResult& R)
	{
		if (GUiShot.Handle.IsValid())
		{
			UGameViewportClient::OnScreenshotCaptured().Remove(GUiShot.Handle);
			GUiShot.Handle.Reset();
		}
		GUiShot.bPending = false;
		TFunction<void(const Descent::Ui::FUiShotResult&)> Cb = MoveTemp(GUiShot.Done);
		GUiShot.Done = nullptr;
		if (Cb)
		{
			Cb(R);
		}
	}

	void UiShotCaptured(int32 Width, int32 Height, const TArray<FColor>& Colors)
	{
		if (!GUiShot.bPending)
		{
			return;
		}
		Descent::Ui::FUiShotResult R;
		R.Path = GUiShot.Path;
		R.Size = FIntPoint(Width, Height);
		if (Colors.Num() != Width * Height || Colors.Num() == 0)
		{
			R.Error = FString::Printf(TEXT("pixel count %d does not match %dx%d"), Colors.Num(), Width, Height);
			UiShotFinish(R);
			return;
		}
		// The window back buffer can carry alpha 0; a PNG with that looks empty in a viewer.
		TArray<FColor> Opaque = Colors;
		for (FColor& C : Opaque)
		{
			C.A = 255;
		}
		R.Stats = Descent::Image::ComputeStats(TConstArrayView64<FColor>(Opaque.GetData(), Opaque.Num()), Width, Height);

		TArray64<uint8> Png;
		FImageUtils::PNGCompressImageArray(Width, Height, TArrayView64<const FColor>(Opaque.GetData(), Opaque.Num()), Png);
		if (Png.Num() == 0 || !FFileHelper::SaveArrayToFile(Png, *GUiShot.Path))
		{
			R.Error = FString::Printf(TEXT("could not write %s"), *GUiShot.Path);
			UiShotFinish(R);
			return;
		}

		TSharedRef<FJsonObject> Side = MakeShared<FJsonObject>();
		Side->SetNumberField(TEXT("w"), Width);
		Side->SetNumberField(TEXT("h"), Height);
		Side->SetBoolField(TEXT("with_ui"), true);
		Side->SetNumberField(TEXT("mean_luma"), Descent::Json::Round(R.Stats.MeanLuma));
		Side->SetNumberField(TEXT("centre_luma"), Descent::Json::Round(R.Stats.CentreLuma));
		Side->SetNumberField(TEXT("non_black"), Descent::Json::Round(R.Stats.NonBlack));
		Side->SetNumberField(TEXT("p99_luma"), Descent::Json::Round(R.Stats.P99Luma));
		FString SideError;
		const FString Stem = FPaths::GetBaseFilename(GUiShot.Path);
		if (!FDescentStateRegistry::WriteJsonFile(Side, FDescentStateRegistry::GetShotDir() / (Stem + TEXT(".json")), SideError))
		{
			R.Error = SideError;
			UiShotFinish(R);
			return;
		}
		R.bOk = true;
		UiShotFinish(R);
	}
}

namespace Descent::Ui
{
	bool IsUiShotPending()
	{
		return GUiShot.bPending;
	}

	bool RequestUiShot(const FString& Stem, TFunction<void(const FUiShotResult&)> Done)
	{
		FUiShotResult Fail;
		if (!FDescentStateRegistry::IsValidFileStem(Stem))
		{
			Fail.Error = FString::Printf(TEXT("'%s' is not a valid file name"), *Stem);
			Done(Fail);
			return false;
		}
		if (GUiShot.bPending || FScreenshotRequest::IsScreenshotRequested())
		{
			Fail.Error = TEXT("another screenshot is still pending");
			Done(Fail);
			return false;
		}
		if (!GEngine || !GEngine->GameViewport || !GEngine->GameViewport->Viewport)
		{
			Fail.Error = TEXT("there is no game viewport to capture");
			Done(Fail);
			return false;
		}
		GUiShot.Path = FDescentStateRegistry::GetShotDir() / (Stem + TEXT(".png"));
		IFileManager::Get().MakeDirectory(*FPaths::GetPath(GUiShot.Path), true);
		IFileManager::Get().Delete(*GUiShot.Path, false, true);
		GUiShot.Done = MoveTemp(Done);
		GUiShot.bPending = true;
		GUiShot.Handle = UGameViewportClient::OnScreenshotCaptured().AddStatic(&UiShotCaptured);
		FScreenshotRequest::RequestScreenshot(GUiShot.Path, /*bInShowUI*/ true, /*bAddFilenameSuffix*/ false);
		return true;
	}
}