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

#include "Camera/PlayerCameraManager.h"
#include "Engine/Engine.h"
#include "Engine/GameViewportClient.h"
#include "Engine/World.h"
#include "EngineUtils.h"
#include "GameFramework/GameModeBase.h"
#include "GameFramework/Pawn.h"
#include "GameFramework/PlayerController.h"
#include "HAL/FileManager.h"
#include "Misc/App.h"
#include "Misc/CommandLine.h"
#include "Misc/EngineVersion.h"
#include "Misc/FileHelper.h"
#include "Misc/Parse.h"
#include "Misc/Paths.h"
#include "Policies/PrettyJsonPrintPolicy.h"
#include "Serialization/JsonSerializer.h"
#include "Serialization/JsonWriter.h"
#include "UnrealClient.h"

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

void FDescentStateRegistry::Register(FName Section, FProvider Provider)
{
	Providers.Add(Section, MoveTemp(Provider));
}

bool FDescentStateRegistry::Unregister(FName Section)
{
	return Providers.Remove(Section) > 0;
}

void FDescentStateRegistry::RegisterPart(FName Section, FName Part, FProvider Provider)
{
	Parts.FindOrAdd(Section).Add(Part, MoveTemp(Provider));
}

bool FDescentStateRegistry::UnregisterPart(FName Section, FName Part)
{
	TMap<FName, FProvider>* Found = Parts.Find(Section);
	if (!Found)
	{
		return false;
	}
	const bool bRemoved = Found->Remove(Part) > 0;
	if (Found->IsEmpty())
	{
		Parts.Remove(Section);
	}
	return bRemoved;
}

TArray<FName> FDescentStateRegistry::GetSectionNames() const
{
	TArray<FName> Names;
	Providers.GetKeys(Names);
	for (const TPair<FName, TMap<FName, FProvider>>& P : Parts)
	{
		Names.AddUnique(P.Key);
	}
	Names.Sort(FNameLexicalLess());
	return Names;
}

TSharedRef<FJsonObject> FDescentStateRegistry::BuildSnapshot(UWorld* World, const FString& SnapshotName) const
{
	TSharedRef<FJsonObject> Root = MakeShared<FJsonObject>();
	Root->SetNumberField(TEXT("schema"), 1);
	Root->SetStringField(TEXT("name"), SnapshotName);

	for (const FName& Name : GetSectionNames())
	{
		TSharedRef<FJsonObject> Section = MakeShared<FJsonObject>();
		if (const FProvider* Whole = Providers.Find(Name))
		{
			(*Whole)(World, *Section);
		}
		if (const TMap<FName, FProvider>* PartMap = Parts.Find(Name))
		{
			TArray<FName> PartNames;
			PartMap->GetKeys(PartNames);
			PartNames.Sort(FNameLexicalLess());
			for (const FName& PartName : PartNames)
			{
				TSharedRef<FJsonObject> PartObject = MakeShared<FJsonObject>();
				(*PartMap)[PartName](World, *PartObject);
				Section->SetObjectField(PartName.ToString(), PartObject);
			}
		}
		Root->SetObjectField(Name.ToString(), Section);
	}
	return Root;
}

FString FDescentStateRegistry::GetStateDir()
{
	return FPaths::ConvertRelativePathToFull(FPaths::ProjectSavedDir() / TEXT("State"));
}

FString FDescentStateRegistry::GetShotDir()
{
	return FPaths::ConvertRelativePathToFull(FPaths::ProjectSavedDir() / TEXT("Shots"));
}

bool FDescentStateRegistry::IsValidFileStem(const FString& Stem)
{
	if (Stem.IsEmpty() || Stem.Contains(TEXT("..")))
	{
		return false;
	}
	for (const TCHAR C : Stem)
	{
		const bool bOk = FChar::IsAlnum(C) || C == TEXT('_') || C == TEXT('-') || C == TEXT('.');
		if (!bOk)
		{
			return false;
		}
	}
	return true;
}

namespace
{
	using FPrettyWriter = TJsonWriter<TCHAR, TPrettyJsonPrintPolicy<TCHAR>>;

	/**
	 * Writes a JSON value tree ourselves rather than through FJsonSerializer: that one prints doubles with 17
	 * significant digits (2.8 becomes 2.7999999999999998), which makes dumps hard to read and to diff. Numbers
	 * here use the shortest form that round-trips at the 4-decimal precision the providers round to.
	 */
	void WriteValueTree(FPrettyWriter& W, const FString* Name, const TSharedPtr<FJsonValue>& V);

	void WriteObjectBody(FPrettyWriter& W, const FJsonObject& Obj)
	{
		for (const TPair<FString, TSharedPtr<FJsonValue>>& Field : Obj.Values)
		{
			WriteValueTree(W, &Field.Key, Field.Value);
		}
	}

	void WriteValueTree(FPrettyWriter& W, const FString* Name, const TSharedPtr<FJsonValue>& V)
	{
		if (!V.IsValid() || V->Type == EJson::Null)
		{
			Name ? W.WriteNull(*Name) : W.WriteNull();
			return;
		}
		switch (V->Type)
		{
		case EJson::Number:
		{
			const FString Text = FString::SanitizeFloat(V->AsNumber(), 0);
			Name ? W.WriteRawJSONValue(*Name, Text) : W.WriteRawJSONValue(Text);
			break;
		}
		case EJson::String:
			Name ? W.WriteValue(*Name, V->AsString()) : W.WriteValue(V->AsString());
			break;
		case EJson::Boolean:
			Name ? W.WriteValue(*Name, V->AsBool()) : W.WriteValue(V->AsBool());
			break;
		case EJson::Array:
			Name ? W.WriteArrayStart(*Name) : W.WriteArrayStart();
			for (const TSharedPtr<FJsonValue>& Item : V->AsArray())
			{
				WriteValueTree(W, nullptr, Item);
			}
			W.WriteArrayEnd();
			break;
		case EJson::Object:
			Name ? W.WriteObjectStart(*Name) : W.WriteObjectStart();
			WriteObjectBody(W, *V->AsObject());
			W.WriteObjectEnd();
			break;
		default:
			Name ? W.WriteNull(*Name) : W.WriteNull();
			break;
		}
	}
}

bool FDescentStateRegistry::WriteJsonFile(const TSharedRef<FJsonObject>& Object, const FString& AbsPath, FString& OutError)
{
	FString Text;
	TSharedRef<FPrettyWriter> Writer = TJsonWriterFactory<TCHAR, TPrettyJsonPrintPolicy<TCHAR>>::Create(&Text);
	Writer->WriteObjectStart();
	WriteObjectBody(*Writer, *Object);
	Writer->WriteObjectEnd();
	Writer->Close();

	IFileManager::Get().MakeDirectory(*FPaths::GetPath(AbsPath), /*Tree*/ true);
	if (!FFileHelper::SaveStringToFile(Text, *AbsPath, FFileHelper::EEncodingOptions::ForceUTF8WithoutBOM))
	{
		OutError = FString::Printf(TEXT("could not write %s"), *AbsPath);
		return false;
	}
	return true;
}

namespace Descent::Json
{
	TSharedRef<FJsonValue> Vec3(const FVector& V)
	{
		TArray<TSharedPtr<FJsonValue>> A;
		A.Add(MakeShared<FJsonValueNumber>(Round(V.X)));
		A.Add(MakeShared<FJsonValueNumber>(Round(V.Y)));
		A.Add(MakeShared<FJsonValueNumber>(Round(V.Z)));
		return MakeShared<FJsonValueArray>(A);
	}

	TSharedRef<FJsonValue> Rot3(const FRotator& R)
	{
		return Vec3(FVector(R.Pitch, R.Yaw, R.Roll));
	}
}

// ---------------------------------------------------------------------------------------------------------
// Built-in sections
// ---------------------------------------------------------------------------------------------------------

namespace
{
	const FName SecMeta(TEXT("meta"));
	const FName SecCamera(TEXT("camera"));
	const FName SecPawn(TEXT("pawn"));
	const FName SecWorld(TEXT("world"));

	APlayerController* FirstPC(UWorld* World)
	{
		return World ? World->GetFirstPlayerController() : nullptr;
	}

	void ProvideMeta(UWorld* World, FJsonObject& S)
	{
		S.SetStringField(TEXT("engine_version"), FEngineVersion::Current().ToString());
		S.SetStringField(TEXT("project"), FApp::GetProjectName());
		S.SetStringField(TEXT("map"), World ? World->GetMapName() : FString());
		S.SetBoolField(TEXT("has_world"), World != nullptr);
		S.SetNumberField(TEXT("frame"), static_cast<double>(GFrameCounter));
		S.SetNumberField(TEXT("world_time_s"), World ? Descent::Json::Round(World->GetTimeSeconds()) : 0.0);
		S.SetNumberField(TEXT("real_time_s"), Descent::Json::Round(FPlatformTime::Seconds() - GStartTime));
		if (GEngine && GEngine->GameViewport && GEngine->GameViewport->Viewport)
		{
			const FIntPoint Size = GEngine->GameViewport->Viewport->GetSizeXY();
			TArray<TSharedPtr<FJsonValue>> Px;
			Px.Add(MakeShared<FJsonValueNumber>(Size.X));
			Px.Add(MakeShared<FJsonValueNumber>(Size.Y));
			S.SetArrayField(TEXT("viewport_px"), Px);
		}
		FString Scenario;
		if (FParse::Value(FCommandLine::Get(), TEXT("DescentScenario="), Scenario))
		{
			S.SetStringField(TEXT("scenario"), Scenario);
		}
	}

	void ProvideCamera(UWorld* World, FJsonObject& S)
	{
		APlayerController* PC = FirstPC(World);
		S.SetBoolField(TEXT("available"), PC != nullptr);
		if (!PC)
		{
			return;
		}
		FVector Loc;
		FRotator Rot;
		PC->GetPlayerViewPoint(Loc, Rot);
		S.SetField(TEXT("location_m"), Descent::Json::Vec3(Descent::UUToMetres(Loc)));
		S.SetField(TEXT("rotation_deg_pyr"), Descent::Json::Rot3(Rot));
		if (PC->PlayerCameraManager)
		{
			S.SetNumberField(TEXT("fov_deg"), Descent::Json::Round(PC->PlayerCameraManager->GetFOVAngle()));
		}
		if (AActor* VT = PC->GetViewTarget())
		{
			S.SetStringField(TEXT("view_target"), VT->GetName());
		}
	}

	void ProvidePawn(UWorld* World, FJsonObject& S)
	{
		APlayerController* PC = FirstPC(World);
		APawn* Pawn = PC ? PC->GetPawn() : nullptr;
		S.SetBoolField(TEXT("available"), Pawn != nullptr);
		if (!Pawn)
		{
			return;
		}
		S.SetStringField(TEXT("class"), Pawn->GetClass()->GetName());
		S.SetField(TEXT("location_m"), Descent::Json::Vec3(Descent::UUToMetres(Pawn->GetActorLocation())));
		S.SetField(TEXT("velocity_mps"), Descent::Json::Vec3(Descent::UUToMetres(Pawn->GetVelocity())));
	}

	void ProvideWorld(UWorld* World, FJsonObject& S)
	{
		if (!World)
		{
			return;
		}
		AGameModeBase* GM = World->GetAuthGameMode();
		S.SetStringField(TEXT("game_mode"), GM ? GM->GetClass()->GetName() : FString());
		TMap<FString, int32> ByClass;
		int32 Total = 0;
		for (TActorIterator<AActor> It(World); It; ++It)
		{
			++Total;
			ByClass.FindOrAdd(It->GetClass()->GetName())++;
		}
		S.SetNumberField(TEXT("actor_count"), Total);
		TSharedRef<FJsonObject> Classes = MakeShared<FJsonObject>();
		ByClass.KeySort(TLess<FString>());
		for (const TPair<FString, int32>& P : ByClass)
		{
			Classes->SetNumberField(P.Key, P.Value);
		}
		S.SetObjectField(TEXT("actors_by_class"), Classes);
	}
}

void FDescentStateRegistry::RegisterBuiltInProviders()
{
	FDescentStateRegistry& R = Get();
	R.Register(SecMeta, &ProvideMeta);
	R.Register(SecCamera, &ProvideCamera);
	R.Register(SecPawn, &ProvidePawn);
	R.Register(SecWorld, &ProvideWorld);
}

void FDescentStateRegistry::UnregisterBuiltInProviders()
{
	FDescentStateRegistry& R = Get();
	R.Unregister(SecMeta);
	R.Unregister(SecCamera);
	R.Unregister(SecPawn);
	R.Unregister(SecWorld);
}