// Copyright (c) Microsoft Corporation. 
// Licensed under the MIT license.

#include "LODRecipeRegistry.h"

#include "AssetRegistryModule.h"
#include "AssetData.h"
#include "Engine/AssetManager.h"

#include "Logging/MessageLog.h"
#include "Misc/UObjectToken.h"

#include "Engine/StaticMesh.h"
#include "Engine/SkeletalMesh.h"

DEFINE_LOG_CATEGORY_STATIC(LogSimplygonLODRecipeRegistry, Verbose, All);
#define LOCTEXT_NAMESPACE "LODRecipeRegistry"

FLODRecipeRegistry::FLODRecipeRegistry()
{
	RefreshRegistry();
}

FLODRecipeRegistry& FLODRecipeRegistry::Get()
{
	static FLODRecipeRegistry Singleton;
	return Singleton;
}

TSoftObjectPtr<ULODRecipe> FLODRecipeRegistry::FindLODRecipe(TSoftObjectPtr<UObject> Mesh)
{
	if (MeshToRecipeMap.Contains(Mesh.ToSoftObjectPath()))
		return TSoftObjectPtr<ULODRecipe>(MeshToRecipeMap[Mesh.ToSoftObjectPath()]);

	return nullptr;
}

void FLODRecipeRegistry::Add(TSoftObjectPtr<UObject> Mesh, ULODRecipe* LODRecipe)
{
	MeshToRecipeMap.Add(Mesh.ToSoftObjectPath(), LODRecipe);
}

void FLODRecipeRegistry::Remove(TSoftObjectPtr<UObject> Mesh)
{
	MeshToRecipeMap.Remove(Mesh.ToSoftObjectPath());
}

void FLODRecipeRegistry::RefreshRegistry()
{
	FAssetRegistryModule& AssetRegistryModule = FModuleManager::LoadModuleChecked<FAssetRegistryModule>("AssetRegistry");
	if (AssetRegistryModule.Get().IsLoadingAssets())
	{
		// We are still discovering assets, listen for the completion delegate before building the graph
		if (!AssetRegistryModule.Get().OnFilesLoaded().IsBoundToObject(this))
		{
			OnAssetRegFilesLoaded = AssetRegistryModule.Get().OnFilesLoaded().AddRaw(this, &FLODRecipeRegistry::OnInitialAssetRegistrySearchComplete);
		}
		return;
	}

	//Reset map
	MeshToRecipeMap.Empty();

	TArray<FAssetData> AllLODRecipeAssets;
	{
		FARFilter Filter;
		Filter.bIncludeOnlyOnDiskAssets = !GIsEditor;
		Filter.ClassNames.Add(ULODRecipe::StaticClass()->GetFName());
		Filter.bRecursiveClasses = true;
		AssetRegistryModule.Get().GetAssets(Filter, AllLODRecipeAssets);
	}
	
	// Store how many recipe references each mesh has
	TMap<FName, TArray<FSoftObjectPath>> PackageToRecipeReferences;

	TArray<FName> AllRecipeDependencies;
	for (const FAssetData& RecipeAssetData : AllLODRecipeAssets)
	{
		TArray<FName> Dependencies;
		FName RecipePackageName = *RecipeAssetData.ToSoftObjectPath().GetLongPackageName();
		AssetRegistryModule.Get().GetDependencies(RecipePackageName, Dependencies, UE::AssetRegistry::EDependencyCategory::Package);
		AllRecipeDependencies.Append(Dependencies);

		for (const FName& MeshPackage : Dependencies)
		{
			PackageToRecipeReferences.FindOrAdd(MeshPackage).Add(RecipeAssetData.ToSoftObjectPath());
		}
	}

	//Get static and skeletal meash asset data from recipe dependencies
	TArray<FAssetData> MeshAssetData;
	{
		FARFilter Filter;
		Filter.bIncludeOnlyOnDiskAssets = !GIsEditor;
		Filter.PackageNames.Append(AllRecipeDependencies);
		Filter.ClassNames.Add(USkeletalMesh::StaticClass()->GetFName());
		Filter.ClassNames.Add(UStaticMesh::StaticClass()->GetFName());
		AssetRegistryModule.Get().GetAssets(Filter, MeshAssetData);
	}

	TArray<TSharedRef<FTokenizedMessage>> Errors;
	for (const FAssetData& MeshAsset : MeshAssetData)
	{
		const TArray<FSoftObjectPath>* LODRecipes = PackageToRecipeReferences.Find(MeshAsset.PackageName);

		if (!LODRecipes)
			continue;
	
		MeshToRecipeMap.Add(MeshAsset.ToSoftObjectPath(), *LODRecipes->begin());

		if (LODRecipes->Num() <= 1)
			continue;

		TSharedRef<FTokenizedMessage> ErrorMessage = FTokenizedMessage::Create(EMessageSeverity::Error);
		ErrorMessage->AddToken(FAssetNameToken::Create(MeshAsset.PackageName.ToString()))
					->AddToken(FTextToken::Create(FText::FromString(TEXT("has multiple LODRecipe references:"))));

		for (const FSoftObjectPath& LODRecipePath : *LODRecipes)
		{
			ErrorMessage->AddToken( FAssetNameToken::Create(LODRecipePath.GetLongPackageName()) );
		}

		Errors.Add(ErrorMessage);
	}

	if (Errors.Num() > 0)
	{
		static const FName LogName = "Simplygon";
		FMessageLog MessageLog(LogName);
		MessageLog.NewPage(LOCTEXT("SimplygonLODRecipeRegistryNewPage", "LOD Recipe Registry"));
		MessageLog.SuppressLoggingToOutputLog(false);

		MessageLog.AddMessages(Errors);
		MessageLog.Open(EMessageSeverity::Error);
	}

	if (AssetRegistryModule.Get().OnFilesLoaded().IsBoundToObject(this))
	{
		AssetRegistryModule.Get().OnFilesLoaded().Remove(OnAssetRegFilesLoaded);
		OnAssetRegFilesLoaded.Reset();
	}
}

void FLODRecipeRegistry::OnInitialAssetRegistrySearchComplete()
{
	RefreshRegistry();
}

#undef LOCTEXT_NAMESPACE
