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

#include "LODRecipeCommandlet.h"
#include "LODRecipeManager.h"
#include "LODRecipeContext.h"

#include "HAL/FileManager.h"
#include "Misc/CommandLine.h"
#include "Misc/FileHelper.h"
#include "Misc/Paths.h"
#include "AssetRegistryModule.h"
#include "IAssetRegistry.h"
#include "LODRecipe.h"
#include "Editor.h"
#include "FileHelpers.h"
#include "Misc/FeedbackContext.h"
#include "ISourceControlModule.h"
#include "ISourceControlProvider.h"
#include "SourceControlHelpers.h"
#include "Engine/StaticMesh.h"
#include "Engine/SkeletalMesh.h"
#include "HAL/PlatformFilemanager.h"
#include "UObject/SavePackage.h"

DEFINE_LOG_CATEGORY_STATIC(LogLODRecipeCommandlet, VeryVerbose, All);

struct FPackageSaveContext
{
	bool bShouldAttemptToSave;
	bool bShouldAttemptToAdd;
	UPackage* PackageToSave;
	FString PackageFilename;
	FPackageSaveContext()
		: bShouldAttemptToSave(false)
		, bShouldAttemptToAdd(false)
		, PackageToSave(nullptr)
		, PackageFilename(TEXT(""))
	{}
};

static FString GetSwitchString(const ELODRecipeCommandletSwitches Switch)
{
	const UEnum* EnumPtr = FindObject<UEnum>(ANY_PACKAGE, TEXT("ELODRecipeCommandletSwitches"), true);
	if (!EnumPtr)
		return FString();

	return EnumPtr->GetNameStringByValue((int64)Switch);
}

static bool HasSwitch(const ELODRecipeCommandletSwitches Switch, const TArray<FString>& Switches)
{
	FString SwitchName = GetSwitchString(Switch);
	if (SwitchName.IsEmpty())
		return false;

	return Switches.Contains(SwitchName);
}

static void GetPackagesFromRecipe(const ULODRecipe& Recipe, TArray<UPackage*>& Packages, const FLODRecipeContext& Context)
{
	TArray<TSoftObjectPtr<UStaticMesh>> StaticMeshes = Recipe.GetStaticMeshes();
	TArray<TSoftObjectPtr<USkeletalMesh>> SkeletalMeshes = Recipe.GetSkeletalMeshes();

	if (Context.IncludeStaticMeshes())
	{
		for (TSoftObjectPtr<UStaticMesh>& Mesh : StaticMeshes)
		{
			Mesh.LoadSynchronous();
			Packages.Add(Mesh->GetOutermost());
		}
	}

	if (Context.IncludeSkeletalMeshes())
	{
		for (TSoftObjectPtr<USkeletalMesh>& Mesh : SkeletalMeshes)
		{
			Mesh.LoadSynchronous();
			Packages.Add(Mesh->GetOutermost());
		}
	}

}

ULODRecipeCommandlet::ULODRecipeCommandlet()
	: bRunOnlyStaticMeshes(false)
	, bRunOnlySkeletalMeshes(false)
	, bShowHelp(false)
	, bAllowSourceControl(false)
	, bHasSourceControl(false)
	, bSaveAssets(false)
	, bBuildSelectedRecipes(false)
{
}

void ULODRecipeCommandlet::PrintUsage()
{
	FString HelpText = TEXT("\n\n LOD Recipe Commandlet\n\tExample usage: UE4Editor-Cmd.exe <ProjectName> -run=LODRecipeCommandlet -Save");

	const UEnum* EnumPtr = FindObject<UEnum>(ANY_PACKAGE, TEXT("ELODRecipeCommandletSwitches"), true);
	if (EnumPtr)
	{
		for (int32 SwitchIndex = 0; SwitchIndex < (int32)ELODRecipeCommandletSwitches::NUM; ++SwitchIndex)
		{
			FString EnumName = EnumPtr->GetNameStringByValue(SwitchIndex);
			FText Description = EnumPtr->GetToolTipTextByIndex(SwitchIndex);

			HelpText.Append(FString::Printf(TEXT("\n\n -%s"), *EnumName));
			HelpText.Append(FString::Printf(TEXT("\n\t%s"), *Description.ToString()));
		}
	}

	UE_LOG(LogLODRecipeCommandlet, Display, TEXT("%s"), *HelpText);
}


bool ULODRecipeCommandlet::ParseParams(const FString& InParams)
{
	TArray<FString> Tokens;
	TArray<FString> Switches;
	TMap<FString, FString> ParamVals;

	ParseCommandLine(*InParams, Tokens, Switches, ParamVals);

	if (HasSwitch(ELODRecipeCommandletSwitches::Help, Switches) || Switches.Contains(TEXT("?")))
	{
		bShowHelp = true;
	}

	bRunOnlyStaticMeshes = HasSwitch(ELODRecipeCommandletSwitches::OnlyStaticMeshes, Switches);
	bRunOnlySkeletalMeshes = HasSwitch(ELODRecipeCommandletSwitches::OnlySkeletalMeshes, Switches);
	bAllowSourceControl = HasSwitch(ELODRecipeCommandletSwitches::SourceControl, Switches);
	bSaveAssets = HasSwitch(ELODRecipeCommandletSwitches::Save, Switches);

	FString LODRecipesStr = GetSwitchString(ELODRecipeCommandletSwitches::LODRecipes);
	if (!LODRecipesStr.IsEmpty() && ParamVals.Contains(LODRecipesStr))
	{
		FString LODRecipeValues = ParamVals[LODRecipesStr];
		LODRecipeValues.ParseIntoArray(SelectedLODRecipes, TEXT(","), true);
		bBuildSelectedRecipes = SelectedLODRecipes.Num() > 0;
	}

	return true;
}

int32 ULODRecipeCommandlet::Main( const FString& Params )
{
	UE_LOG(LogLODRecipeCommandlet, Log, TEXT("Running LODRecipeCommandlet"));
	
	int Result = 0;
	if (!ParseParams(Params) || bShowHelp)
	{
		PrintUsage();
		return Result;
	}

	if (bAllowSourceControl)
	{
		ISourceControlProvider& SourceControlProvider = ISourceControlModule::Get().GetProvider();
		SourceControlProvider.Init();
		bHasSourceControl = SourceControlProvider.IsEnabled();
		if (!bHasSourceControl)
		{
			UE_LOG(LogLODRecipeCommandlet, Error, TEXT("Could not connect to source control!"))
		}
	}
	else
	{
		bHasSourceControl = false;
	}

	FAssetRegistryModule& AssetRegistryModule = FModuleManager::LoadModuleChecked<FAssetRegistryModule>("AssetRegistry");
	
	//Find all LOD recipe assets
	TArray<FAssetData> LODRecipesToBuild;
	AssetRegistryModule.Get().GetAssetsByClass(ULODRecipe::StaticClass()->GetFName(), LODRecipesToBuild);
	
	if (bBuildSelectedRecipes)
	{
		LODRecipesToBuild.RemoveAll([&SelectedLODRecipes = SelectedLODRecipes](const FAssetData& Visited)
			{
				FString AssetName = Visited.AssetName.ToString();
				FString PackageName = Visited.PackageName.ToString();

				// Do not remove if a match is found
				bool bAssetFound = SelectedLODRecipes.Contains(AssetName) || SelectedLODRecipes.Contains(PackageName);
				return !bAssetFound;
			});

		// Validate if all selected recipes got added.
		for (const FString& SelectedRecipe : SelectedLODRecipes)
		{
			bool bFoundSelectedRecipe = LODRecipesToBuild.ContainsByPredicate([SelectedRecipe](const FAssetData& Visited)
				{
					FString Name = Visited.AssetName.ToString();
					FString PackageName = Visited.PackageName.ToString();
					return SelectedRecipe == Name || SelectedRecipe == PackageName;
				});

			if (!bFoundSelectedRecipe)
				UE_LOG(LogLODRecipeCommandlet, Warning, TEXT("LOD recipe with name %s does not exist."), *SelectedRecipe);
		}
	}

	TArray<FPackageSaveContext> SaveContexts;
	FLODRecipeContext Context;
	Context.FilterCriteria = bRunOnlyStaticMeshes && !bRunOnlySkeletalMeshes ? ELODRecipeFilterCriteria::StaticMeshes : Context.FilterCriteria;
	Context.FilterCriteria = bRunOnlySkeletalMeshes && !bRunOnlyStaticMeshes ? ELODRecipeFilterCriteria::SkeletalMeshes : Context.FilterCriteria;
		
	ISourceControlProvider& SourceControlProvider = ISourceControlModule::Get().GetProvider();

	TArray<ULODRecipe*> Recipes;
	for (auto& LODRecipe : LODRecipesToBuild)
	{
		UObject* LODRecipeObj = LODRecipe.FastGetAsset(true);
		ULODRecipe* Recipe = Cast<ULODRecipe>(LODRecipeObj);
				
		if (!Recipe)
			continue;

		TArray<UPackage*> ReferencedPackages;
		TArray<FSourceControlStateRef> ReferencedPackageStates;
					
		GetPackagesFromRecipe(*Recipe, ReferencedPackages, Context);

		bool bUseSourceControl = bHasSourceControl && SourceControlProvider.IsAvailable();
		SourceControlProvider.GetState(ReferencedPackages, ReferencedPackageStates, EStateCacheUsage::ForceUpdate);
		
		bool bLODRecipeHasMeshes = ReferencedPackages.Num() > 0;
		if (bLODRecipeHasMeshes)
			Recipes.AddUnique(Recipe);

		for (int32 PackageIndex = 0; PackageIndex < ReferencedPackages.Num(); ++PackageIndex)
		{
			FPackageSaveContext SaveContext;
			SaveContext.PackageToSave = ReferencedPackages[PackageIndex];
			SaveContext.PackageFilename = SourceControlHelpers::PackageFilename(SaveContext.PackageToSave);

			if (bUseSourceControl)
			{
				FSourceControlStateRef PackageSCState = ReferencedPackageStates[PackageIndex];

				bool bPackageCanBeCheckedOut = false;
				if (PackageSCState->IsCheckedOutOther())
				{
					// Cannot checkout, file is already checked out
					UE_LOG(LogLODRecipeCommandlet, Error, TEXT("%s is already checked out by someone else, can not submit!"), *SaveContext.PackageFilename);
					Result = -10;
				}
				else if (!PackageSCState->IsCurrent())
				{
					// Cannot checkout, file is not at head revision
					UE_LOG(LogLODRecipeCommandlet, Error, TEXT("%s is not at the head revision and cannot be checked out"), *SaveContext.PackageFilename);
					Result = -10;
				}
				else if (PackageSCState->CanCheckout())
				{
					const bool bWasCheckedOut = SourceControlHelpers::CheckOutOrAddFile(SaveContext.PackageFilename);
					SaveContext.bShouldAttemptToSave = bWasCheckedOut;
					if (!bWasCheckedOut)
					{
						UE_LOG(LogLODRecipeCommandlet, Error, TEXT("%s could not be checked out"), *SaveContext.PackageFilename);
						Result = -10;
					}
				}
				else
				{
					// package was not checked out by another user and is at the current head revision and could not be checked out
					// this means it should be added after save because it doesn't exist
					SaveContext.bShouldAttemptToSave = true;
				}
			}
			else
			{
				bool bIsReadOnly = IFileManager::Get().IsReadOnly(*SaveContext.PackageFilename);
				if (bIsReadOnly)
				{
					SaveContext.bShouldAttemptToSave = FPlatformFileManager::Get().GetPlatformFile().SetReadOnly(*SaveContext.PackageFilename, false);
					if (!SaveContext.bShouldAttemptToSave)
					{
						UE_LOG(LogLODRecipeCommandlet, Error, TEXT("%s is read only and could not be made writable.  Will not save"), *SaveContext.PackageFilename);
						Result = -20;
					}
				}
				else
				{
					SaveContext.bShouldAttemptToSave = true;
				}
			}

			if (SaveContext.bShouldAttemptToSave)
			{
				SaveContexts.Add(SaveContext);
			}
		}
	}

	FDelegateHandle NewAssetsEvent = ULODRecipe::OnNewAssetsCreated().AddLambda([&bSaveAssets = bSaveAssets](const TArray<UObject*>* AssetsToSave)
		{
			if (!bSaveAssets)
				return;

			if (!AssetsToSave)
				return;

			for (const UObject* Asset : *AssetsToSave)
			{
				FString PackageFilename = SourceControlHelpers::PackageFilename(Asset->GetOutermost());
				FSavePackageArgs SavePackageArgs;
				SavePackageArgs.TopLevelFlags = RF_Standalone | RF_Public;
				GEditor->SavePackage(Asset->GetOutermost(), nullptr, *PackageFilename, SavePackageArgs);
			}
		});

	FLODRecipeManager::Get().Build(Recipes, Context);
	ULODRecipe::OnNewAssetsCreated().Remove(NewAssetsEvent);

	// Now save packages
	if (bSaveAssets)
	{
		for (FPackageSaveContext SaveContext : SaveContexts)
		{
			if (SaveContext.bShouldAttemptToSave)
			{
				FSavePackageArgs SavePackageArgs;
				SavePackageArgs.TopLevelFlags = RF_Standalone;
				GEditor->SavePackage(SaveContext.PackageToSave,nullptr, *SaveContext.PackageFilename, SavePackageArgs);
			}
		}
	}

	return Result;
}
