#include "Debug/DescentDebugSubsystem.h"
#include "Debug/DescentScenario.h"
#include "Debug/DescentStateRegistry.h"
#include "TheDescent.h"
#include "DescentUnits.h"

#include "Camera/CameraActor.h"
#include "Camera/CameraComponent.h"
#include "Engine/Engine.h"
#include "Engine/GameInstance.h"
#include "Engine/GameViewportClient.h"
#include "Engine/World.h"
#include "GameFramework/Pawn.h"
#include "GameFramework/PlayerController.h"
#include "HAL/FileManager.h"
#include "HAL/PlatformMisc.h"
#include "ImageCore.h"
#include "ImageUtils.h"
#include "Misc/CommandLine.h"
#include "Misc/FileHelper.h"
#include "Misc/Parse.h"
#include "Misc/Paths.h"
#include "Serialization/JsonReader.h"
#include "Serialization/JsonSerializer.h"
#include "UnrealClient.h"

UDescentDebugSubsystem* UDescentDebugSubsystem::Get(const UWorld* World)
{
	UGameInstance* GI = World ? World->GetGameInstance() : nullptr;
	return GI ? GI->GetSubsystem<UDescentDebugSubsystem>() : nullptr;
}

void UDescentDebugSubsystem::Initialize(FSubsystemCollectionBase& Collection)
{
	Super::Initialize(Collection);

	// -DescentScenario=<name> on the command line starts the scenario as soon as the world is up.
	// Not tied to a particular map: the runner itself waits for the world to begin play.
	FString ScenarioName;
	if (FParse::Value(FCommandLine::Get(), TEXT("DescentScenario="), ScenarioName) && !ScenarioName.IsEmpty())
	{
		StartScenario(ScenarioName);
	}
}

void UDescentDebugSubsystem::Deinitialize()
{
	Runner.Reset();
	if (ShotCapturedHandle.IsValid())
	{
		UGameViewportClient::OnScreenshotCaptured().Remove(ShotCapturedHandle);
		ShotCapturedHandle.Reset();
	}
	Super::Deinitialize();
}

// ---------------------------------------------------------------------------------------------------------
// State dump
// ---------------------------------------------------------------------------------------------------------

bool UDescentDebugSubsystem::DumpState(UWorld* World, const FString& Stem, FString& OutPathOrError)
{
	if (!FDescentStateRegistry::IsValidFileStem(Stem))
	{
		OutPathOrError = FString::Printf(TEXT("'%s' is not a valid file name (use letters, digits, _ - . only)"), *Stem);
		return false;
	}
	const FString Path = FDescentStateRegistry::GetStateDir() / (Stem + TEXT(".json"));
	const TSharedRef<FJsonObject> Snapshot = FDescentStateRegistry::Get().BuildSnapshot(World, Stem);
	FString Error;
	if (!FDescentStateRegistry::WriteJsonFile(Snapshot, Path, Error))
	{
		OutPathOrError = Error;
		return false;
	}
	OutPathOrError = Path;
	return true;
}

// ---------------------------------------------------------------------------------------------------------
// Screenshots
// ---------------------------------------------------------------------------------------------------------

FString UDescentDebugSubsystem::GetShotSidecarPath(const FString& Stem)
{
	return FDescentStateRegistry::GetShotDir() / (Stem + TEXT(".json"));
}

bool UDescentDebugSubsystem::RequestShot(const FString& Stem, TFunction<void(const FDescentShotResult&)> Done)
{
	FDescentShotResult Fail;
	if (!FDescentStateRegistry::IsValidFileStem(Stem))
	{
		Fail.Error = FString::Printf(TEXT("'%s' is not a valid file name (use letters, digits, _ - . only)"), *Stem);
		Done(Fail);
		return false;
	}
	if (bShotPending || 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;
	}

	PendingShotPath = FDescentStateRegistry::GetShotDir() / (Stem + TEXT(".png"));
	IFileManager::Get().MakeDirectory(*FPaths::GetPath(PendingShotPath), /*Tree*/ true);
	// A stale file from an earlier run must never be mistaken for this run's result.
	IFileManager::Get().Delete(*PendingShotPath, /*RequireExists*/ false, /*EvenReadOnly*/ true);
	IFileManager::Get().Delete(*GetShotSidecarPath(Stem), /*RequireExists*/ false, /*EvenReadOnly*/ true);

	ShotDone = MoveTemp(Done);
	bShotPending = true;
	ShotCapturedHandle = UGameViewportClient::OnScreenshotCaptured().AddUObject(this, &UDescentDebugSubsystem::OnScreenshotCaptured);

	// While UGameViewportClient::OnScreenshotCaptured (a static delegate on the viewport client, not the one on
	// FScreenshotRequest) has a listener, the engine hands us the pixels instead of writing its own file
	// (UGameViewportClient::ProcessScreenShots), so we choose the path and the PNG encoding. The pixels are the
	// game viewport's back buffer: the same frame a player would see, no Slate UI.
	FScreenshotRequest::RequestScreenshot(PendingShotPath, /*bInShowUI*/ false, /*bAddFilenameSuffix*/ false);
	return true;
}

void UDescentDebugSubsystem::OnScreenshotCaptured(int32 Width, int32 Height, const TArray<FColor>& Colors)
{
	if (!bShotPending)
	{
		return; // someone else's screenshot
	}

	FDescentShotResult Result;
	Result.Path = PendingShotPath;
	Result.Size = FIntPoint(Width, Height);

	const int64 NumPixels = static_cast<int64>(Width) * Height;
	if (Colors.Num() != NumPixels || NumPixels == 0)
	{
		Result.Error = FString::Printf(TEXT("pixel count %d does not match %dx%d"), Colors.Num(), Width, Height);
		FinishShot(Result);
		return;
	}

	// Health statistics so a scenario can tell a black or unrendered frame from a dark scene, and so the look
	// tests can assert luminance bands without opening the PNG.
	Result.Stats = Descent::Image::ComputeStats(TConstArrayView64<FColor>(Colors.GetData(), Colors.Num()), Width, Height);
	Result.MeanLuma = Result.Stats.MeanLuma;
	Result.NonBlackFraction = Result.Stats.NonBlack;

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

	// The sidecar: the same numbers as JSON next to the PNG (DESIGN_UE.md section 9.2).
	const FDescentLumaStats& S = Result.Stats;
	TSharedRef<FJsonObject> Side = MakeShared<FJsonObject>();
	Side->SetNumberField(TEXT("w"), Width);
	Side->SetNumberField(TEXT("h"), Height);
	Side->SetNumberField(TEXT("mean_luma"), Descent::Json::Round(S.MeanLuma));
	Side->SetNumberField(TEXT("centre_luma"), Descent::Json::Round(S.CentreLuma));
	Side->SetNumberField(TEXT("non_black"), Descent::Json::Round(S.NonBlack));
	Side->SetNumberField(TEXT("p99_luma"), Descent::Json::Round(S.P99Luma));
	Side->SetNumberField(TEXT("centre_detail"), Descent::Json::Round(S.CentreDetail));
	Side->SetNumberField(TEXT("centre_rms"), Descent::Json::Round(S.CentreRms));
	TArray<TSharedPtr<FJsonValue>> Rgb;
	Rgb.Add(MakeShared<FJsonValueNumber>(Descent::Json::Round(S.MeanR)));
	Rgb.Add(MakeShared<FJsonValueNumber>(Descent::Json::Round(S.MeanG)));
	Rgb.Add(MakeShared<FJsonValueNumber>(Descent::Json::Round(S.MeanB)));
	Side->SetArrayField(TEXT("mean_rgb"), Rgb);
	FString SideError;
	const FString Stem = FPaths::GetBaseFilename(PendingShotPath);
	if (!FDescentStateRegistry::WriteJsonFile(Side, GetShotSidecarPath(Stem), SideError))
	{
		Result.Error = SideError;
		FinishShot(Result);
		return;
	}

	Result.bOk = true;
	FinishShot(Result);
}

void UDescentDebugSubsystem::FinishShot(const FDescentShotResult& Result)
{
	if (ShotCapturedHandle.IsValid())
	{
		UGameViewportClient::OnScreenshotCaptured().Remove(ShotCapturedHandle);
		ShotCapturedHandle.Reset();
	}
	bShotPending = false;
	TFunction<void(const FDescentShotResult&)> Cb = MoveTemp(ShotDone);
	ShotDone = nullptr;
	if (Cb)
	{
		Cb(Result);
	}
}

// ---------------------------------------------------------------------------------------------------------
// Measuring a part of a shot that was taken earlier
// ---------------------------------------------------------------------------------------------------------

namespace
{
	/** Reads Saved/Shots/<Stem>.png back as sRGB BGRA8. */
	bool LoadShotPixels(const FString& Stem, TArray<FColor>& OutPixels, int32& OutW, int32& OutH, FString& OutError)
	{
		if (!FDescentStateRegistry::IsValidFileStem(Stem))
		{
			OutError = FString::Printf(TEXT("'%s' is not a valid file name (use letters, digits, _ - . only)"), *Stem);
			return false;
		}
		const FString Path = FDescentStateRegistry::GetShotDir() / (Stem + TEXT(".png"));
		if (!IFileManager::Get().FileExists(*Path))
		{
			OutError = FString::Printf(TEXT("there is no shot %s (take it first)"), *Path);
			return false;
		}
		FImage Loaded;
		if (!FImageUtils::LoadImage(*Path, Loaded))
		{
			OutError = FString::Printf(TEXT("could not read %s"), *Path);
			return false;
		}
		FImage Bgra;
		Loaded.CopyTo(Bgra, ERawImageFormat::BGRA8, EGammaSpace::sRGB);
		OutW = Bgra.SizeX;
		OutH = Bgra.SizeY;
		const TArrayView64<FColor> View = Bgra.AsBGRA8();
		OutPixels.SetNumUninitialized(static_cast<int32>(View.Num()));
		FMemory::Memcpy(OutPixels.GetData(), View.GetData(), View.Num() * sizeof(FColor));
		return true;
	}

	/** Adds one entry to the "regions" array of the shot's sidecar (creating the sidecar if the PNG came from elsewhere). */
	bool AppendRegion(const FString& Stem, int32 W, int32 H, const TSharedRef<FJsonObject>& Entry, FString& OutError)
	{
		const FString SidePath = UDescentDebugSubsystem::GetShotSidecarPath(Stem);
		TSharedPtr<FJsonObject> Side;
		FString Text;
		if (FFileHelper::LoadFileToString(Text, *SidePath))
		{
			FJsonSerializer::Deserialize(TJsonReaderFactory<>::Create(Text), Side);
		}
		if (!Side.IsValid())
		{
			Side = MakeShared<FJsonObject>();
			Side->SetNumberField(TEXT("w"), W);
			Side->SetNumberField(TEXT("h"), H);
		}
		TArray<TSharedPtr<FJsonValue>> Regions;
		if (const TArray<TSharedPtr<FJsonValue>>* Existing = nullptr; Side->TryGetArrayField(TEXT("regions"), Existing))
		{
			Regions = *Existing;
		}
		Regions.Add(MakeShared<FJsonValueObject>(Entry));
		Side->SetArrayField(TEXT("regions"), Regions);
		return FDescentStateRegistry::WriteJsonFile(Side.ToSharedRef(), SidePath, OutError);
	}
}

bool UDescentDebugSubsystem::MeasureShotRegion(const FString& Stem, double X0, double Y0, double X1, double Y1, double& OutMeanLuma, FString& OutError)
{
	TArray<FColor> Pixels;
	int32 W = 0, H = 0;
	if (!LoadShotPixels(Stem, Pixels, W, H, OutError))
	{
		return false;
	}
	OutMeanLuma = Descent::Image::MeanLumaRect(TConstArrayView64<FColor>(Pixels.GetData(), Pixels.Num()), W, H, X0, Y0, X1, Y1);

	TSharedRef<FJsonObject> E = MakeShared<FJsonObject>();
	E->SetStringField(TEXT("kind"), TEXT("rect"));
	E->SetNumberField(TEXT("x0"), X0);
	E->SetNumberField(TEXT("y0"), Y0);
	E->SetNumberField(TEXT("x1"), X1);
	E->SetNumberField(TEXT("y1"), Y1);
	E->SetNumberField(TEXT("mean_luma"), Descent::Json::Round(OutMeanLuma));
	return AppendRegion(Stem, W, H, E, OutError);
}

bool UDescentDebugSubsystem::MeasureShotRing(const FString& Stem, double Cx, double Cy, double R0, double R1, double& OutMeanLuma, FString& OutError)
{
	TArray<FColor> Pixels;
	int32 W = 0, H = 0;
	if (!LoadShotPixels(Stem, Pixels, W, H, OutError))
	{
		return false;
	}
	OutMeanLuma = Descent::Image::MeanLumaRing(TConstArrayView64<FColor>(Pixels.GetData(), Pixels.Num()), W, H, Cx, Cy, R0, R1);

	TSharedRef<FJsonObject> E = MakeShared<FJsonObject>();
	E->SetStringField(TEXT("kind"), TEXT("ring"));
	E->SetNumberField(TEXT("cx"), Cx);
	E->SetNumberField(TEXT("cy"), Cy);
	E->SetNumberField(TEXT("r0"), R0);
	E->SetNumberField(TEXT("r1"), R1);
	E->SetNumberField(TEXT("mean_luma"), Descent::Json::Round(OutMeanLuma));
	return AppendRegion(Stem, W, H, E, OutError);
}

// ---------------------------------------------------------------------------------------------------------
// Teleport, camera, quit
// ---------------------------------------------------------------------------------------------------------

bool UDescentDebugSubsystem::Teleport(UWorld* World, const FVector& LocationM)
{
	APlayerController* PC = World ? World->GetFirstPlayerController() : nullptr;
	AActor* Target = PC ? (PC->GetPawn() ? static_cast<AActor*>(PC->GetPawn()) : PC->GetViewTarget()) : nullptr;
	if (!Target)
	{
		return false;
	}
	return Target->TeleportTo(Descent::MetresToUU(LocationM), Target->GetActorRotation(), /*bIsATest*/ false, /*bNoCheck*/ true);
}

bool UDescentDebugSubsystem::PlaceCamera(const FVector& LocationM, const FRotator& Rotation, float FovDeg)
{
	UWorld* World = GetWorld();
	APlayerController* PC = World ? World->GetFirstPlayerController() : nullptr;
	if (!PC)
	{
		return false;
	}
	if (!ScenarioCamera || !IsValid(ScenarioCamera))
	{
		FActorSpawnParameters Params;
		Params.Name = TEXT("DescentScenarioCamera");
		Params.ObjectFlags |= RF_Transient;
		Params.SpawnCollisionHandlingOverride = ESpawnActorCollisionHandlingMethod::AlwaysSpawn;
		ScenarioCamera = World->SpawnActor<ACameraActor>(Params);
	}
	if (!ScenarioCamera)
	{
		return false;
	}
	ScenarioCamera->SetActorLocationAndRotation(Descent::MetresToUU(LocationM), Rotation, /*bSweep*/ false, nullptr, ETeleportType::TeleportPhysics);
	UCameraComponent* Cam = ScenarioCamera->GetCameraComponent();
	Cam->SetFieldOfView(FovDeg);
	Cam->bConstrainAspectRatio = false; // no letterbox bars in screenshots
	if (PC->GetViewTarget() != ScenarioCamera)
	{
		PC->SetViewTarget(ScenarioCamera);
	}
	return true;
}

void UDescentDebugSubsystem::QuitProcess(int32 ExitCode, const TCHAR* Why)
{
	UE_LOG(LogDescent, Display, TEXT("Quit requested (exit code %d): %s"), ExitCode, Why);
	if (GIsEditor)
	{
		// Inside the editor (Play In Editor) any exit request would end the user's whole editor session, and the
		// runtime module cannot reach the editor's "stop PIE" call. Ignore it; Esc stops Play. (-game, the
		// packaged game and the scenario runner have GIsEditor == false and exit normally.)
		UE_LOG(LogDescent, Warning, TEXT("Quit ignored inside the editor; press Esc to stop Play."));
		return;
	}
	if (ExitCode != 0)
	{
		// A forced exit is the only way to get a non-zero code out of a normal Windows shutdown, which
		// reports 0 whatever happened. Flush first so the reason is in the log.
		if (GLog)
		{
			GLog->Flush();
		}
		FPlatformMisc::RequestExitWithStatus(/*Force*/ true, static_cast<uint8>(ExitCode));
	}
	else
	{
		FPlatformMisc::RequestExitWithStatus(/*Force*/ false, 0);
	}
}

bool UDescentDebugSubsystem::StartScenario(const FString& Name)
{
	const FDescentScenarioRegistry::FEntry* Entry = FDescentScenarioRegistry::Get().Find(Name);
	if (!Entry)
	{
		UE_LOG(LogDescent, Error, TEXT("DESCENT_SCENARIO_RESULT FAIL scenario=%s reason=\"no such scenario; registered: %s\""),
			*Name, *FString::Join(FDescentScenarioRegistry::Get().GetNames(), TEXT(", ")));
		QuitProcess(3, TEXT("unknown scenario"));
		return false;
	}
	if (Runner && !Runner->IsFinished())
	{
		UE_LOG(LogDescent, Warning, TEXT("A scenario is already running"));
		return false;
	}
	FDescentScenarioBuilder Builder;
	Entry->Build(Builder);
	Runner = MakeUnique<FDescentScenarioRunner>(this, Name, MoveTemp(Builder));
	Runner->Start();
	return true;
}