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

#include "AssetCompilingManager.h"
#include "Engine/Engine.h"
#include "Engine/GameViewportClient.h"
#include "Engine/World.h"
#include "GameFramework/PlayerController.h"
#include "Misc/CommandLine.h"
#include "Misc/Parse.h"
#include "ShaderCompiler.h"

// ---------------------------------------------------------------------------------------------------------
// Context
// ---------------------------------------------------------------------------------------------------------

void FDescentScenarioContext::SetNamed(FName Key, AActor* Actor)
{
	Named.Add(Key, Actor);
}

AActor* FDescentScenarioContext::GetNamed(FName Key) const
{
	const TWeakObjectPtr<AActor>* Found = Named.Find(Key);
	return Found ? Found->Get() : nullptr;
}

void FDescentScenarioContext::Fail(const FString& Reason)
{
	if (!bFailed)
	{
		bFailed = true;
		FailReason = Reason;
	}
}

UDescentSimSubsystem* FDescentScenarioContext::Sim() const
{
	return UDescentSimSubsystem::Get(World);
}

TSharedRef<FJsonObject> FDescentScenarioContext::Snapshot() const
{
	return FDescentStateRegistry::Get().BuildSnapshot(World, ScenarioName);
}

// ---------------------------------------------------------------------------------------------------------
// Registry
// ---------------------------------------------------------------------------------------------------------

FDescentScenarioRegistry& FDescentScenarioRegistry::Get()
{
	static FDescentScenarioRegistry Instance;
	return Instance;
}

void FDescentScenarioRegistry::Register(const TCHAR* Name, const TCHAR* Description, FBuildFn Build)
{
	FEntry& E = Entries.FindOrAdd(Name);
	E.Name = Name;
	E.Description = Description;
	E.Build = Build;
}

const FDescentScenarioRegistry::FEntry* FDescentScenarioRegistry::Find(const FString& Name) const
{
	return Entries.Find(Name);
}

TArray<FString> FDescentScenarioRegistry::GetNames() const
{
	TArray<FString> Names;
	Entries.GetKeys(Names);
	Names.Sort();
	return Names;
}

// ---------------------------------------------------------------------------------------------------------
// Builder
// ---------------------------------------------------------------------------------------------------------

namespace
{
	/** True while shaders or async assets (meshes, textures, Nanite builds) are still being prepared. */
	bool IsContentStillPreparing()
	{
		if (GShaderCompilingManager && GShaderCompilingManager->IsCompiling())
		{
			return true;
		}
		return FAssetCompilingManager::Get().GetNumRemainingAssets() > 0;
	}
}

FDescentScenarioBuilder& FDescentScenarioBuilder::TimeLimit(float Seconds)
{
	TimeLimitSeconds = Seconds;
	return *this;
}

FDescentScenarioBuilder& FDescentScenarioBuilder::Do(const FString& Label, TFunction<void(FDescentScenarioContext&)> Fn)
{
	Steps.Add({ Label, [Fn = MoveTemp(Fn)](FDescentScenarioContext& C)
	{
		Fn(C);
		return true;
	} });
	return *this;
}

FDescentScenarioBuilder& FDescentScenarioBuilder::Until(const FString& Label, FDescentStepFn Fn)
{
	Steps.Add({ Label, MoveTemp(Fn) });
	return *this;
}

FDescentScenarioBuilder& FDescentScenarioBuilder::WaitFrames(int32 Frames)
{
	Steps.Add({ FString::Printf(TEXT("WaitFrames(%d)"), Frames), [Frames](FDescentScenarioContext& C)
	{
		return C.FramesInStep() >= Frames;
	} });
	return *this;
}

FDescentScenarioBuilder& FDescentScenarioBuilder::WaitSeconds(float Seconds)
{
	Steps.Add({ FString::Printf(TEXT("WaitSeconds(%.2f)"), Seconds), [Seconds](FDescentScenarioContext& C)
	{
		return C.SecondsInStep() >= Seconds;
	} });
	return *this;
}

FDescentScenarioBuilder& FDescentScenarioBuilder::WaitSim(double Seconds)
{
	// The sim clock at the first call of this step is the zero of the wait.
	TSharedRef<double> StartTime = MakeShared<double>(-1.0);
	Steps.Add({ FString::Printf(TEXT("WaitSim(%.2f)"), Seconds), [Seconds, StartTime](FDescentScenarioContext& C)
	{
		const UDescentSimSubsystem* Sim = C.Sim();
		if (!Sim)
		{
			C.Fail(TEXT("WaitSim: there is no sim subsystem in this world"));
			return false;
		}
		if (*StartTime < 0.0)
		{
			*StartTime = Sim->SimTime();
		}
		return Sim->SimTime() - *StartTime >= Seconds - 1e-9;
	} });
	return *this;
}

FDescentScenarioBuilder& FDescentScenarioBuilder::FastForward(double Seconds, double Dt)
{
	return Do(FString::Printf(TEXT("FastForward(%.2f s, dt %.4f)"), Seconds, Dt), [Seconds, Dt](FDescentScenarioContext& C)
	{
		UDescentSimSubsystem* Sim = C.Sim();
		if (!Sim)
		{
			C.Fail(TEXT("FastForward: there is no sim subsystem in this world"));
			return;
		}
		const double T0 = Sim->SimTime();
		const int64 Steps = Sim->FastForward(Seconds, Dt);
		UE_LOG(LogDescent, Display, TEXT("FastForward: %lld steps, sim time %.3f -> %.3f s"), Steps, T0, Sim->SimTime());
		if (Steps == 0 && Seconds > 0.0)
		{
			C.Fail(FString::Printf(TEXT("FastForward(%.2f s, dt %.4f) ran no steps"), Seconds, Dt));
		}
	});
}

FDescentScenarioBuilder& FDescentScenarioBuilder::Expect(const FString& Label, TFunction<bool(FDescentScenarioContext&)> Fn)
{
	return Do(FString::Printf(TEXT("Expect(%s)"), *Label), [Label, Fn = MoveTemp(Fn)](FDescentScenarioContext& C)
	{
		if (Fn(C))
		{
			UE_LOG(LogDescent, Display, TEXT("Expect '%s': ok"), *Label);
		}
		else
		{
			C.Fail(FString::Printf(TEXT("Expect '%s' failed"), *Label));
		}
	});
}

FDescentScenarioBuilder& FDescentScenarioBuilder::ExpectNear(const FString& Label, TFunction<double(FDescentScenarioContext&)> ValueFn, double Expected, double Tolerance)
{
	return Do(FString::Printf(TEXT("ExpectNear(%s)"), *Label), [Label, ValueFn = MoveTemp(ValueFn), Expected, Tolerance](FDescentScenarioContext& C)
	{
		const double Value = ValueFn(C);
		if (FMath::IsFinite(Value) && FMath::Abs(Value - Expected) <= Tolerance)
		{
			UE_LOG(LogDescent, Display, TEXT("ExpectNear '%s': ok (value %.6g, expected %.6g +- %.6g)"), *Label, Value, Expected, Tolerance);
		}
		else
		{
			C.Fail(FString::Printf(TEXT("ExpectNear '%s' failed: value %.6g, expected %.6g +- %.6g (off by %.6g)"),
				*Label, Value, Expected, Tolerance, Value - Expected));
		}
	});
}

FDescentScenarioBuilder& FDescentScenarioBuilder::ExpectNear(const FString& Label, double Value, double Expected, double Tolerance)
{
	return ExpectNear(Label, [Value](FDescentScenarioContext&) { return Value; }, Expected, Tolerance);
}

FDescentScenarioBuilder& FDescentScenarioBuilder::Exec(const FString& Command)
{
	return Do(FString::Printf(TEXT("Exec(%s)"), *Command), [Command](FDescentScenarioContext& C)
	{
		UE_LOG(LogDescent, Display, TEXT("Scenario exec: %s"), *Command);
		if (GEngine && !GEngine->Exec(C.World, *Command))
		{
			// Many valid commands return false (they only set a value), so this is information, not failure.
			UE_LOG(LogDescent, Verbose, TEXT("Exec '%s' returned false"), *Command);
		}
	});
}

FDescentScenarioBuilder& FDescentScenarioBuilder::Camera(const FVector& LocationM, const FRotator& Rotation, float FovDeg)
{
	return Do(FString::Printf(TEXT("Camera(%.1f,%.1f,%.1f)"), LocationM.X, LocationM.Y, LocationM.Z),
		[LocationM, Rotation, FovDeg](FDescentScenarioContext& C)
		{
			if (!C.Debug || !C.Debug->PlaceCamera(LocationM, Rotation, FovDeg))
			{
				C.Fail(TEXT("could not place the scenario camera (no player controller?)"));
			}
		});
}

FDescentScenarioBuilder& FDescentScenarioBuilder::LookAt(const FVector& LocationM, const FVector& TargetM, float FovDeg)
{
	return Camera(LocationM, (TargetM - LocationM).Rotation(), FovDeg);
}

FDescentScenarioBuilder& FDescentScenarioBuilder::Shot(const FString& Label, int32 WarmupFrames, float WarmupSeconds)
{
	struct FShotState
	{
		int32 Phase = 0;
		int64 QuietSinceFrame = -1;
		double QuietSinceTime = 0.0;
		bool bDone = false;
		FDescentShotResult Result;
	};
	TSharedRef<FShotState> St = MakeShared<FShotState>();

	Steps.Add({ FString::Printf(TEXT("Shot(%s)"), *Label), [St, Label, WarmupFrames, WarmupSeconds](FDescentScenarioContext& C)
	{
		if (St->Phase == 0)
		{
			// Phase 0: let shaders / Nanite builds / async meshes finish, then keep rendering a few frames so
			// temporal AA, Lumen and exposure have settled. Only then ask for the frame.
			if (IsContentStillPreparing())
			{
				St->QuietSinceFrame = -1;
				return false;
			}
			if (St->QuietSinceFrame < 0)
			{
				St->QuietSinceFrame = static_cast<int64>(GFrameCounter);
				St->QuietSinceTime = FPlatformTime::Seconds();
			}
			const bool bFramesOk = static_cast<int64>(GFrameCounter) - St->QuietSinceFrame >= WarmupFrames;
			const bool bTimeOk = FPlatformTime::Seconds() - St->QuietSinceTime >= WarmupSeconds;
			if (!bFramesOk || !bTimeOk)
			{
				return false;
			}
			if (!C.Debug)
			{
				C.Fail(TEXT("no debug subsystem"));
				return false;
			}
			const FString Stem = C.ScenarioName + TEXT("_") + Label;
			TSharedRef<FShotState> Captured = St;
			const bool bAccepted = C.Debug->RequestShot(Stem, [Captured](const FDescentShotResult& R)
			{
				Captured->Result = R;
				Captured->bDone = true;
			});
			if (!bAccepted && !St->bDone)
			{
				C.Fail(TEXT("screenshot request was refused"));
			}
			St->Phase = 1;
			return false;
		}

		// Phase 1: wait for the capture callback, then judge the result.
		if (!St->bDone)
		{
			return false;
		}
		const FDescentShotResult& R = St->Result;
		if (!R.bOk)
		{
			C.Fail(FString::Printf(TEXT("Shot(%s) failed: %s"), *Label, *R.Error));
			return false;
		}
		UE_LOG(LogDescent, Display, TEXT("Shot(%s): %s  %dx%d  mean luma %.3f  centre %.3f  p99 %.3f  non-black %.1f%%"),
			*Label, *R.Path, R.Size.X, R.Size.Y, R.Stats.MeanLuma, R.Stats.CentreLuma, R.Stats.P99Luma, R.Stats.NonBlack * 100.0);
		if (R.Size != C.RequestedResolution)
		{
			C.Fail(FString::Printf(TEXT("Shot(%s) is %dx%d but %dx%d was requested (-ResX/-ResY)"),
				*Label, R.Size.X, R.Size.Y, C.RequestedResolution.X, C.RequestedResolution.Y));
			return false;
		}
		if (R.NonBlackFraction < 0.001)
		{
			C.Fail(FString::Printf(TEXT("Shot(%s) is black: the frame was not rendered"), *Label));
			return false;
		}
		C.LastShotStats = R.Stats;
		return true;
	} });
	return *this;
}

FDescentScenarioBuilder& FDescentScenarioBuilder::DumpState(const FString& Label)
{
	return Do(FString::Printf(TEXT("DumpState(%s)"), *Label), [Label](FDescentScenarioContext& C)
	{
		FString PathOrError;
		if (UDescentDebugSubsystem::DumpState(C.World, C.ScenarioName + TEXT("_") + Label, PathOrError))
		{
			UE_LOG(LogDescent, Display, TEXT("DumpState(%s): %s"), *Label, *PathOrError);
		}
		else
		{
			C.Fail(FString::Printf(TEXT("DumpState(%s) failed: %s"), *Label, *PathOrError));
		}
	});
}

FDescentScenarioBuilder& FDescentScenarioBuilder::Quit(int32 ExitCode)
{
	return Do(TEXT("Quit"), [ExitCode](FDescentScenarioContext& C)
	{
		C.RequestQuit(ExitCode);
	});
}

// ---------------------------------------------------------------------------------------------------------
// Runner
// ---------------------------------------------------------------------------------------------------------

FDescentScenarioRunner::FDescentScenarioRunner(UDescentDebugSubsystem* InOwner, const FString& InName, FDescentScenarioBuilder&& InBuilder)
	: Owner(InOwner)
	, Name(InName)
	, Builder(MoveTemp(InBuilder))
{
	Context.Debug = InOwner;
	Context.ScenarioName = InName;

	int32 ResX = 0, ResY = 0;
	FParse::Value(FCommandLine::Get(), TEXT("ResX="), ResX);
	FParse::Value(FCommandLine::Get(), TEXT("ResY="), ResY);
	if (ResX > 0 && ResY > 0)
	{
		Context.RequestedResolution = FIntPoint(ResX, ResY);
	}

	// First step of every scenario: make the viewport the size that was asked for. The window normally comes
	// up at -ResX/-ResY already; if the OS clamped it (small desktop, DPI scaling), r.SetRes forces it.
	const FIntPoint Wanted = Context.RequestedResolution;
	Builder.Steps.Insert({ TEXT("EnsureResolution"), [Wanted](FDescentScenarioContext& C)
	{
		if (!GEngine || !GEngine->GameViewport || !GEngine->GameViewport->Viewport)
		{
			return true; // nothing to size; Shot() will fail by itself if it matters
		}
		const FIntPoint Now = GEngine->GameViewport->Viewport->GetSizeXY();
		if (Now == Wanted)
		{
			return true;
		}
		if (C.FramesInStep() == 0)
		{
			UE_LOG(LogDescent, Warning, TEXT("Viewport is %dx%d, wanted %dx%d; sending r.SetRes"), Now.X, Now.Y, Wanted.X, Wanted.Y);
			GEngine->Exec(C.World, *FString::Printf(TEXT("r.SetRes %dx%dw"), Wanted.X, Wanted.Y));
		}
		return C.FramesInStep() >= 10;
	} }, 0);
}

FDescentScenarioRunner::~FDescentScenarioRunner()
{
	if (TickHandle.IsValid())
	{
		FTSTicker::GetCoreTicker().RemoveTicker(TickHandle);
	}
}

void FDescentScenarioRunner::FailRun(const FString& Reason)
{
	Context.Fail(Reason);
}

void FDescentScenarioRunner::Start()
{
	WaitStartTime = FPlatformTime::Seconds();
	UE_LOG(LogDescent, Display, TEXT("Scenario '%s' queued: %d steps, limit %.0f s, %dx%d"),
		*Name, Builder.Steps.Num(), Builder.TimeLimitSeconds, Context.RequestedResolution.X, Context.RequestedResolution.Y);
	TickHandle = FTSTicker::GetCoreTicker().AddTicker(FTickerDelegate::CreateRaw(this, &FDescentScenarioRunner::Tick));
}

void FDescentScenarioRunner::Finish(int32 ExitCode, const FString& Reason)
{
	if (bFinished)
	{
		return;
	}
	bFinished = true;
	const double Elapsed = bStarted ? Context.SecondsInScenario() : 0.0;
	if (ExitCode == 0)
	{
		UE_LOG(LogDescent, Display, TEXT("DESCENT_SCENARIO_RESULT PASS scenario=%s seconds=%.1f"), *Name, Elapsed);
	}
	else
	{
		UE_LOG(LogDescent, Error, TEXT("DESCENT_SCENARIO_RESULT FAIL scenario=%s seconds=%.1f exit=%d reason=\"%s\""),
			*Name, Elapsed, ExitCode, *Reason);
	}
	UDescentDebugSubsystem::QuitProcess(ExitCode, *Reason);
}

bool FDescentScenarioRunner::Tick(float /*DeltaTime*/)
{
	if (bFinished)
	{
		return false; // returning false removes the ticker
	}

	UWorld* World = Owner ? Owner->GetWorld() : nullptr;
	if (!bStarted)
	{
		// Startup phase: wait for the map to load and the player to exist. Shader compilation of a cold
		// cache can make this take minutes, so it has its own generous limit.
		const bool bReady = World && World->HasBegunPlay() && World->GetFirstPlayerController();
		if (!bReady)
		{
			if (FPlatformTime::Seconds() - WaitStartTime > 300.0)
			{
				Finish(2, TEXT("the world never began play within 300 s"));
				return false;
			}
			return true;
		}
		// The boot steps ran inside StartPlay, before this point. A failed one means the world is not the one the
		// scenario expects, so the run fails here with the boot step named, not later with a confusing symptom.
		if (const ADescentGameMode* GM = World->GetAuthGameMode<ADescentGameMode>())
		{
			if (!GM->GetBootReport().bOk)
			{
				FString Which;
				for (const FDescentBootReport::FEntry& E : GM->GetBootReport().Ran)
				{
					if (!E.bOk)
					{
						Which += FString::Printf(TEXT("%s%s (%s)"), Which.IsEmpty() ? TEXT("") : TEXT(", "), *E.Name, *E.Error);
					}
				}
				Finish(2, FString::Printf(TEXT("boot step failed: %s"), *Which));
				return false;
			}
		}
		bStarted = true;
		Context.ScenarioStartTime = FPlatformTime::Seconds();
		UE_LOG(LogDescent, Display, TEXT("Scenario '%s' starting on map %s after %.1f s"),
			*Name, *World->GetMapName(), Context.ScenarioStartTime - WaitStartTime);
	}
	Context.World = World;

	if (Context.SecondsInScenario() > Builder.TimeLimitSeconds)
	{
		const FString Label = StepIndex < Builder.Steps.Num() ? Builder.Steps[StepIndex].Label : TEXT("(after last step)");
		Finish(2, FString::Printf(TEXT("time limit of %.0f s exceeded in step %d '%s'"), Builder.TimeLimitSeconds, StepIndex, *Label));
		return false;
	}

	// Run as many instantaneous steps as are ready this frame; stop at the first one that needs more time.
	for (int32 Guard = 0; Guard < 256 && StepIndex < Builder.Steps.Num(); ++Guard)
	{
		FDescentScenarioStep& Step = Builder.Steps[StepIndex];
		if (!bStepEntered)
		{
			bStepEntered = true;
			Context.StepEnterFrame = static_cast<int64>(GFrameCounter);
			Context.StepEnterTime = FPlatformTime::Seconds();
			UE_LOG(LogDescent, Display, TEXT("Scenario '%s' step %d/%d: %s"), *Name, StepIndex + 1, Builder.Steps.Num(), *Step.Label);
		}

		const bool bDone = Step.Run(Context);

		if (Context.bFailed)
		{
			Finish(2, FString::Printf(TEXT("step '%s': %s"), *Step.Label, *Context.FailReason));
			return false;
		}
		if (Context.bQuitRequested)
		{
			Finish(Context.QuitCode, Context.QuitCode == 0 ? TEXT("Quit step reached") : TEXT("Quit step with non-zero code"));
			return false;
		}
		if (!bDone)
		{
			return true;
		}
		++StepIndex;
		bStepEntered = false;
	}

	if (StepIndex >= Builder.Steps.Num())
	{
		// A scenario that forgot Quit() still ends rather than hanging until the timeout.
		Finish(0, TEXT("end of steps"));
		return false;
	}
	return true;
}