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

#include "LODRecipeManager.h"
#include "LODRecipe.h"
#include "SimplygonPluginMeshReductionModule.h"
#include "ISimplygonPluginUtilitiesModule.h"
#include "SimplygonPluginUtilities/Private/SimplygonProcessor.h"
#include "SimplygonDataConverter.h"
#include "LODRecipe/Private/LODRecipeRegistry.h"
#include "SimplygonModule.h"

#include "Misc/ScopedSlowTask.h"
#include "Editor.h"
#include "Subsystems/AssetEditorSubsystem.h"
#include "Engine/SkeletalMesh.h"
#include "Rendering/SkeletalMeshLODModel.h"
#include "MaterialUtilities.h"
#include "AssetRegistryModule.h"

#include "Telemetry.h"
#include "TelemetryManager.h"

DEFINE_LOG_CATEGORY_STATIC(LogLODRecipeManager, Verbose, All);
#define LOCTEXT_NAMESPACE "LODRecipeManager"

FLODRecipeManager::FLODRecipeManager()
	: SimplygonWorkCompleted(0)
	, SimplygonWorkProgress(0)
{
	BindDelegates();

	// Cache game thread ticker
	ISimplygonPluginUtilitiesModule& PluginUtilitiesModule = FModuleManager::LoadModuleChecked<ISimplygonPluginUtilitiesModule>("SimplygonPluginUtilities");
	ISimplygonPluginUtilities& PluginUtils = PluginUtilitiesModule.GetUtilities();
	SimplygonProcessor = PluginUtils.GetSimplygonProcessor();
}

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

void FLODRecipeManager::BuildAll()
{
	FAssetRegistryModule& AssetRegistryModule = FModuleManager::LoadModuleChecked<FAssetRegistryModule>("AssetRegistry");
	TArray<FAssetData> AssetDataList;
	FARFilter Filter;
	Filter.ClassNames.Add(ULODRecipe::StaticClass()->GetFName());
	AssetRegistryModule.Get().GetAssets(Filter, AssetDataList);

	TArray<ULODRecipe*> RecipesToBuild;
	for (const FAssetData& AssetData : AssetDataList)
	{
		ULODRecipe* LODRecipe = Cast<ULODRecipe>(AssetData.GetAsset());
		if (!LODRecipe)
			continue;

		RecipesToBuild.AddUnique(LODRecipe);
	}

	Build(RecipesToBuild);
}

void FLODRecipeManager::Build(TSoftObjectPtr<UObject> Mesh)
{
	TSoftObjectPtr<ULODRecipe> LODRecipe = FLODRecipeRegistry::Get().FindLODRecipe(Mesh);
	if (!LODRecipe.IsValid())
	{
		UE_LOG(LogLODRecipeManager, Warning, TEXT("Mesh does not have a LOD recipe attached"));
		return;
	}

	FLODRecipeContext Context;
	Context.SpecificMeshes.Add(Mesh.ToSoftObjectPath());
	Build(TArray<ULODRecipe*>{ LODRecipe.Get() }, Context);
}

void FLODRecipeManager::Build(const TArray<ULODRecipe*>& InLODRecipes, const FLODRecipeContext& InContext)
{
	ISimplygonPluginMeshReductionModule& ReductionModule = FModuleManager::LoadModuleChecked<ISimplygonPluginMeshReductionModule>("SimplygonPluginMeshReduction");
	ISimplygonReduction* SimplygonProcessingInterface = ReductionModule.GetSimplygonReduction();
	if (SimplygonProcessingInterface == nullptr)
	{
		UE_LOG(LogLODRecipeManager, Error, TEXT("Simplygon processing interface not found."));
		return;
	}

	ISimplygonPluginUtilitiesModule& PluginUtilitiesModule = FModuleManager::LoadModuleChecked<ISimplygonPluginUtilitiesModule>("SimplygonPluginUtilities");
	ISimplygonPluginUtilities& SimplygonUtilities = PluginUtilitiesModule.GetUtilities();

	Simplygon::ISimplygon* SimplygonAPI = ISimplygonModule::Get().GetSDK(true);
	if (SimplygonAPI == nullptr)
		return;

	// Filter
	FLODRecipeMeshMap FilteredLODRecipes;
	FilterLODRecipes(InLODRecipes, InContext, FilteredLODRecipes);

	int32 NumMeshes = InitProgress(FilteredLODRecipes);
	int32 SimplygonAmountOfWork = NumMeshes;

	float DataConvertionAmountOfWork = (float)SimplygonAmountOfWork * 0.2f;
	FText DataConvertionProgressText = LOCTEXT("FLODRecipeManagerBuild_ConvertData", "Converting mesh data ({0}/{1})");
	int32 DataConvertionWorkComplete = 0;

	bool bUsePipelineBatch = SimplygonProcessingInterface->IsUsingDistribution();

	FScopedSlowTask SlowTask(SimplygonAmountOfWork + DataConvertionAmountOfWork, (LOCTEXT("FLODRecipeManagerBuild_InitialProgress", "Preparing to generate LODs")));
	SlowTask.MakeDialog();

	Simplygon::spPipelineBatch PipelineBatch;
	TArray<FSimplygonLODRecipeJobData>* JobDataArray = nullptr;
	if( bUsePipelineBatch )
	{
		PipelineBatch = SimplygonAPI->CreatePipelineBatch();
		JobDataArray = new TArray<FSimplygonLODRecipeJobData>();
		SimplygonAmountOfWork = 1;
	}

	FTelemetryBuilder Builder;
	Builder.SetProperties({ FTelemetry::Prop(TEXT("action"), TEXT("build_recipes")),
		FTelemetry::Prop(TEXT("filter_criteria"), (uint8)InContext.FilterCriteria) ,
		FTelemetry::Prop(TEXT("filter_specificmeshes"), InContext.SpecificMeshes.Num() > 0),
		FTelemetry::Prop(TEXT("num_recipes"), InLODRecipes.Num()),
		FTelemetry::Prop(TEXT("num_filtered_recipes"), FilteredLODRecipes.Num()),
		FTelemetry::Prop(TEXT("num_meshes"), NumMeshes)
		});

	FTelemetry::Record(TEXT("Build"), SGTC::LODRecipe, MoveTemp(Builder));

	for (const auto& LODRecipeToMeshes : FilteredLODRecipes)
	{
		ULODRecipe* LODRecipe = LODRecipeToMeshes.Key;
		const TArray<TSoftObjectPtr<UObject>>& Meshes = LODRecipeToMeshes.Value;
		for (auto MeshSoftPtr : Meshes)
		{
			SlowTask.EnterProgressFrame(
				DataConvertionAmountOfWork / (float)NumMeshes,
				FText::Format(DataConvertionProgressText, FText::AsNumber(DataConvertionWorkComplete + 1), FText::AsNumber(NumMeshes)));

			UObject* Mesh = MeshSoftPtr.LoadSynchronous();

			FSimplygonFlattenMaterialSettings FlattenMaterialSettings;
			Simplygon::spPipeline Pipeline = LODRecipe->CreatePipeline(SimplygonAPI, FlattenMaterialSettings);

			// RecipeJobData is released in Complete lambda
			FSimplygonLODRecipeJobData RecipeJobData{};
			RecipeJobData.LODRecipe = LODRecipe;
			RecipeJobData.OriginalMesh = Mesh;

			Simplygon::spScene SceneToProcess = SimplygonAPI->CreateScene();
			FSimplygonMeshMaterialInfo MeshMaterialInfo(false);
			if (Mesh->IsA<UStaticMesh>())
			{
				UStaticMesh* StaticMesh = Cast<UStaticMesh>(Mesh);
				FSimplygonDataConverter::AddToScene(SimplygonAPI, SceneToProcess, FlattenMaterialSettings, StaticMesh, MeshMaterialInfo);
			}
			else if (Mesh->IsA<USkeletalMesh>())
			{
				USkeletalMesh* SkeletalMesh = Cast<USkeletalMesh>(Mesh);

				// Create Geometry Data from Skeletal Mesh
				Simplygon::spGeometryData GeometryData = FSimplygonDataConverter::CreateGeometryData(SimplygonAPI, SceneToProcess, FlattenMaterialSettings, SkeletalMesh, MeshMaterialInfo);

				// Add morph targets to Geometry Data
				if (LODRecipe->IsMorphTargetsAllowed())
					FSimplygonDataConverter::CreateMorphFieldsFromSkeletalMesh(GeometryData, SkeletalMesh, 0);

				// Lastly, Add Geometry Data to Simplygon Scene 
				FSimplygonDataConverter::AddToScene(SimplygonAPI, SceneToProcess, GeometryData, SkeletalMesh->GetName());
			}

			if (MeshMaterialInfo.bBakeMaterials)
			{
				FSimplygonDataConverter::AddToScene(SimplygonAPI, SceneToProcess, MeshMaterialInfo);
			}

			// Need to be able to map back to original sections for LODs that doesn't bake materials. 
			for (const FMeshMaterialSectionId& Id : MeshMaterialInfo.MaterialOutputOrder)
			{
				RecipeJobData.MaterialIndexToSlotName.Append(MeshMaterialInfo.MaterialIndexToSlotName[Id]);
			}

			if( bUsePipelineBatch )
			{
				// 6. Add to pipeline batch
				PipelineBatch->Queue(Pipeline, SceneToProcess);
				JobDataArray->Add(RecipeJobData);
			}
			else
			{
				FSimplygonJobData JobData{};
				JobData.JobGuid = FGuid::NewGuid();
				JobData.UserData = new FSimplygonLODRecipeJobData(RecipeJobData);
				SimplygonProcessingInterface->RunSceneAsync(SceneToProcess, Pipeline, JobData, ProgressDelegate, CompleteDelegate, FailedDelegate);
			}

			++DataConvertionWorkComplete;
		}
	}

	if( bUsePipelineBatch )
	{
		// Launch the batch processing
		UE_LOG(LogLODRecipeManager, Log, TEXT("Processing LOD recipe in batch"));
		FSimplygonJobData JobData;
		JobData.JobGuid = FGuid::NewGuid();
		JobData.UserData = JobDataArray;
		SimplygonProcessingInterface->RunBatchAsync(PipelineBatch, JobData, ProgressDelegate, BatchCompleteDelegate, BatchFailedDelegate);
	}

	//Block until all jobs has returned
	WaitForComplete(SimplygonAmountOfWork, SlowTask);
}

void FLODRecipeManager::BindDelegates()
{
	ProgressDelegate.BindLambda([&WorkProgress = SimplygonWorkProgress](const FSimplygonJobData& InJobData, int32 ProgressPercentage, int32 ProgressPercentageDelta)
		{
			FPlatformAtomics::InterlockedAdd(&WorkProgress, ProgressPercentageDelta);
		});

	CompleteDelegate.BindLambda([=, &WorkCompleted = SimplygonWorkCompleted](const FSimplygonJobData& InJobData, const Simplygon::spPipeline& InPipeline)
		{
			FSimplygonProcessGameThreadDelegate CompleteJobOnGameThread;
			CompleteJobOnGameThread.BindLambda([=]()
				{
					FSimplygonLODRecipeJobData* RecipeJobData = (FSimplygonLODRecipeJobData*)InJobData.UserData;

					ULODRecipe* LODRecipe = RecipeJobData->LODRecipe;
					UStaticMesh* StaticMesh = Cast<UStaticMesh>(RecipeJobData->OriginalMesh);
					USkeletalMesh* SkeletalMesh = Cast<USkeletalMesh>(RecipeJobData->OriginalMesh);
					if (StaticMesh)
						LODRecipe->ProcessCompleted(StaticMesh, *RecipeJobData, InPipeline, InPipeline->GetProcessedScene());
					else if (SkeletalMesh)
						LODRecipe->ProcessSkeletalMeshCompleted(SkeletalMesh, *RecipeJobData, InPipeline, InPipeline->GetProcessedScene());

					FTelemetryBuilder Builder;
					Builder.SetProperties({
						FTelemetry::Prop(TEXT("asset_type"), StaticMesh != nullptr ? TEXT("StaticMesh") : TEXT("SkeletalMesh"))
						});

					FTelemetry::Record(TEXT("JobCompleted"), SGTC::LODRecipe, MoveTemp(Builder));

					delete RecipeJobData;
				});
			

			SimplygonProcessor->Add(CompleteJobOnGameThread);

			FPlatformAtomics::InterlockedAdd(&WorkCompleted, 1);
		});

	BatchCompleteDelegate.BindLambda([=, &WorkCompleted = SimplygonWorkCompleted](const FSimplygonJobData& InJobData, const Simplygon::spPipelineBatch& InPipelineBatch)
		{
			Simplygon::ISimplygon* SimplygonAPI = ISimplygonModule::Get().GetSDK(true);
			if (SimplygonAPI == nullptr)
				return;

			FSimplygonProcessGameThreadDelegate CompleteJobOnGameThread;
			CompleteJobOnGameThread.BindLambda([=]()
				{
					TArray<FSimplygonLODRecipeJobData>* JobDataArray = (TArray<FSimplygonLODRecipeJobData>*)InJobData.UserData;

					for(unsigned int ipipeline = 0; ipipeline < InPipelineBatch->GetBatchCount(); ++ipipeline)
					{
						FSimplygonLODRecipeJobData* RecipeJobData = &((*JobDataArray)[ipipeline]);

						Simplygon::spSceneImporter SceneImporter = SimplygonAPI->CreateSceneImporter();
						SceneImporter->SetImportFilePath(InPipelineBatch->GetOutputPath(ipipeline));
						SceneImporter->RunImport();

						Simplygon::spScene Scene = SceneImporter->GetScene();
						if(Scene.IsNull())
						{
							UE_LOG(LogLODRecipeManager, Error, TEXT("LOD recipe processing failed for %s"), *RecipeJobData->OriginalMesh->GetFName().ToString());
							continue;
						}

						Simplygon::spPipelineSerializer PipelineSerializer = SimplygonAPI->CreatePipelineSerializer();
						Simplygon::spPipeline Pipeline = PipelineSerializer->LoadPipelineFromFile(InPipelineBatch->GetPipelinePath(ipipeline));

						ULODRecipe* LODRecipe = RecipeJobData->LODRecipe;
						UStaticMesh* StaticMesh = Cast<UStaticMesh>(RecipeJobData->OriginalMesh);
						USkeletalMesh* SkeletalMesh = Cast<USkeletalMesh>(RecipeJobData->OriginalMesh);
						if (StaticMesh)
							LODRecipe->ProcessCompleted(StaticMesh, *RecipeJobData, Pipeline, Scene);
						else if (SkeletalMesh)
							LODRecipe->ProcessSkeletalMeshCompleted(SkeletalMesh, *RecipeJobData, Pipeline, Scene);

						FTelemetryBuilder Builder;
						Builder.SetProperties({
							FTelemetry::Prop(TEXT("asset_type"), StaticMesh != nullptr ? TEXT("StaticMesh") : TEXT("SkeletalMesh"))
							});

						FTelemetry::Record(TEXT("JobCompleted"), SGTC::LODRecipe, MoveTemp(Builder));
					}

					delete JobDataArray;
				});			

			SimplygonProcessor->Add(CompleteJobOnGameThread);

			FPlatformAtomics::InterlockedAdd(&WorkCompleted, 1);
		});

	FailedDelegate.BindLambda([&WorkCompleted = SimplygonWorkCompleted](const FSimplygonJobData& InJobData, const FString& InMessage)
		{
			FTelemetryBuilder Builder;
			Builder.SetProperties({ 
				FTelemetry::Prop(TEXT("error_message"), InMessage)
				});

			FTelemetry::Record(TEXT("JobFailed"), SGTC::LODRecipe, MoveTemp(Builder));

			FSimplygonLODRecipeJobData* RecipeJobData = (FSimplygonLODRecipeJobData*)InJobData.UserData;
			delete RecipeJobData;

			UE_LOG(LogLODRecipeManager, Error, TEXT("%s"), *InMessage);
			FPlatformAtomics::InterlockedAdd(&WorkCompleted, 1);
		});

	BatchFailedDelegate.BindLambda([&WorkCompleted = SimplygonWorkCompleted](const FSimplygonJobData& InJobData, const FString& InMessage)
		{
			FTelemetryBuilder Builder;
			Builder.SetProperties({ 
				FTelemetry::Prop(TEXT("error_message"), InMessage)
				});

			FTelemetry::Record(TEXT("JobFailed"), SGTC::LODRecipe, MoveTemp(Builder));

			TArray<FSimplygonLODRecipeJobData>* JobDataArray = (TArray<FSimplygonLODRecipeJobData>*)InJobData.UserData;
			delete JobDataArray;

			UE_LOG(LogLODRecipeManager, Error, TEXT("%s"), *InMessage);
			FPlatformAtomics::InterlockedAdd(&WorkCompleted, 1);
		});
}

int32 FLODRecipeManager::InitProgress(const FLODRecipeMeshMap& LODRecipes)
{
	//Reset old progress data
	SimplygonWorkCompleted = 0;
	SimplygonWorkProgress = 0;

	int32 TotalAmountOfWork = 0;
	for (const auto& LODRecipeToMeshes : LODRecipes)
	{
			const auto& Meshes = LODRecipeToMeshes.Value;
			TotalAmountOfWork += Meshes.Num();
	}

	return TotalAmountOfWork;
}

void FLODRecipeManager::WaitForComplete(const int32 TotalAmountOfWork, FScopedSlowTask& SlowTask)
{
	FText SimplygonProgressText = LOCTEXT("FLODRecipeManager_SimplygonProcess", "Generating LOD for mesh ({0}/{1})");
	int32 LastWorkProgress = 0;
	SlowTask.EnterProgressFrame(0, FText::Format(SimplygonProgressText, FText::AsNumber(SimplygonWorkCompleted + 1), FText::AsNumber(TotalAmountOfWork)));
	while (SimplygonWorkCompleted < TotalAmountOfWork)
	{
		if (SimplygonWorkProgress > LastWorkProgress)
		{
			int32 CurrentProgress = SimplygonWorkProgress;
			int32 Delta = CurrentProgress - LastWorkProgress;
			SlowTask.EnterProgressFrame((float)Delta / 100.0f, FText::Format(SimplygonProgressText, FText::AsNumber(SimplygonWorkCompleted + 1), FText::AsNumber(TotalAmountOfWork)));
			LastWorkProgress = CurrentProgress;
		}
		SimplygonProcessor->Tick(0.0f);
		FPlatformProcess::Sleep(0.1f);
	}
}

void FLODRecipeManager::FilterLODRecipes(const TArray<ULODRecipe*>& LODRecipes, const FLODRecipeContext& InContext, FLODRecipeMeshMap& Result)
{
	Result.Empty();
	bool bIncludeStaticMeshOnly = InContext.IncludeStaticMeshes() && !InContext.IncludeSkeletalMeshes();
	bool bIncludeSkeletalMeshOnly = InContext.IncludeSkeletalMeshes() && !InContext.IncludeStaticMeshes();

	for (ULODRecipe* LODRecipe : LODRecipes)
	{
		TArray<TSoftObjectPtr<UObject>> Meshes = LODRecipe->GetAllMeshes();
		if (InContext.SpecificMeshes.Num() > 0)
		{
			Meshes.RemoveAllSwap([&InContext](TSoftObjectPtr<UObject>& Mesh)
				{
					return InContext.SpecificMeshes.Contains(Mesh.ToSoftObjectPath()) == false;
				});
		}
		
		switch (InContext.FilterCriteria)
		{
		case ELODRecipeFilterCriteria::SkeletalMeshes:
			Meshes.RemoveAllSwap([](TSoftObjectPtr<UObject>& Mesh)
				{
					return Mesh->GetClass()->GetName() != USkeletalMesh::StaticClass()->GetName();
				});
			break;

		case ELODRecipeFilterCriteria::StaticMeshes:
			Meshes.RemoveAllSwap([](TSoftObjectPtr<UObject>& Mesh)
				{
					return Mesh->GetClass()->GetName() != UStaticMesh::StaticClass()->GetName();
				});
			break;

		default:			
			break;
		}
		
		Result.Add(LODRecipe, Meshes);
	}
}

#undef LOCTEXT_NAMESPACE
