FunnyGirl/Source/FunnyGirlEditor/Private/FunnyGirlEditorModule.cpp

// Copyright Epic Games, Inc. All Rights Reserved.

#include "FunnyGirlEditorModule.h"

#include "AssetRegistry/AssetRegistryModule.h"
#include "AssetRegistry/IAssetRegistry.h"
#include "AssetTypeActions/AssetTypeActions_DirectorPrompt.h"
#include "AssetTypeActions/AssetTypeActions_Identity.h"
#include "AssetTypeActions/AssetTypeActions_LLMProvider.h"
#include "AssetTypeActions/AssetTypeActions_PersonalityTrait.h"
#include "AssetTypeActions/AssetTypeActions_SystemEvent.h"
#include "AssetTypeActions/AssetTypeActions_WorldContext.h"
#include "AssetToolsModule.h"
#include "IAssetTools.h"
#include "IAssetTypeActions.h"
#include "Misc/FileHelper.h"
#include "Misc/MessageDialog.h"
#include "Misc/Paths.h"
#include "Modules/ModuleManager.h"
#include "Morph/Editor/LipSyncMappingAssetTypeActions.h"
#include "Morph/Editor/MorphPresetAssetTypeActions.h"
#include "Morph/MorphPresetDataAsset.h"
#include "Morph/Editor/MorphPresetDetailsCustomization.h"
#include "PropertyEditorModule.h"
#include "ToolMenus.h"

#define LOCTEXT_NAMESPACE "FunnyGirlEditor"

EAssetTypeCategories::Type FFunnyGirlEditorModule::LLMAssetCategory = EAssetTypeCategories::Misc;
EAssetTypeCategories::Type FFunnyGirlEditorModule::NarrativeAssetCategory = EAssetTypeCategories::Misc;

void FFunnyGirlEditorModule::StartupModule()
{
	IAssetTools& AssetTools = FModuleManager::LoadModuleChecked<FAssetToolsModule>("AssetTools").Get();
	LLMAssetCategory = AssetTools.RegisterAdvancedAssetCategory(FName(TEXT("LLM")), LOCTEXT("LLMAssetCategory", "LLM"));
	NarrativeAssetCategory = AssetTools.RegisterAdvancedAssetCategory(FName(TEXT("Narrative")), LOCTEXT("NarrativeAssetCategory", "Narrative"));

	RegisterAssetTypeAction(MakeShared<FAssetTypeActions_LLMProvider>());
	RegisterAssetTypeAction(MakeShared<FAssetTypeActions_DirectorPrompt>());
	RegisterAssetTypeAction(MakeShared<FAssetTypeActions_Identity>());
	RegisterAssetTypeAction(MakeShared<FAssetTypeActions_PersonalityTrait>());
	RegisterAssetTypeAction(MakeShared<FAssetTypeActions_SystemEvent>());
	RegisterAssetTypeAction(MakeShared<FAssetTypeActions_WorldContext>());
	RegisterAssetTypeAction(MakeShared<FMorphPresetAssetTypeActions>());
	RegisterAssetTypeAction(MakeShared<FLipSyncMappingAssetTypeActions>());

	FPropertyEditorModule& PropertyEditorModule = FModuleManager::LoadModuleChecked<FPropertyEditorModule>(TEXT("PropertyEditor"));
	PropertyEditorModule.RegisterCustomClassLayout(
		UMorphPresetDataAsset::StaticClass()->GetFName(),
		FOnGetDetailCustomizationInstance::CreateStatic(&FMorphPresetDetailsCustomization::MakeInstance));
	PropertyEditorModule.NotifyCustomizationModuleChanged();
	bMorphPresetCustomizationRegistered = true;

	UToolMenus::RegisterStartupCallback(FSimpleMulticastDelegate::FDelegate::CreateRaw(this, &FFunnyGirlEditorModule::RegisterMenus));
}

void FFunnyGirlEditorModule::ShutdownModule()
{
	if (UObjectInitialized() && FModuleManager::Get().IsModuleLoaded(TEXT("ToolMenus")))
	{
		UToolMenus::UnRegisterStartupCallback(this);
		UToolMenus::UnregisterOwner(this);
	}

	if (FModuleManager::Get().IsModuleLoaded(TEXT("AssetTools")))
	{
		IAssetTools& AssetTools = FModuleManager::GetModuleChecked<FAssetToolsModule>("AssetTools").Get();
		for (const TSharedPtr<IAssetTypeActions>& AssetTypeAction : RegisteredAssetTypeActions)
		{
			if (AssetTypeAction.IsValid())
			{
				AssetTools.UnregisterAssetTypeActions(AssetTypeAction.ToSharedRef());
			}
		}
	}

	RegisteredAssetTypeActions.Empty();

	if (bMorphPresetCustomizationRegistered && FModuleManager::Get().IsModuleLoaded(TEXT("PropertyEditor")))
	{
		FPropertyEditorModule& PropertyEditorModule = FModuleManager::GetModuleChecked<FPropertyEditorModule>(TEXT("PropertyEditor"));
		PropertyEditorModule.UnregisterCustomClassLayout(UMorphPresetDataAsset::StaticClass()->GetFName());
		PropertyEditorModule.NotifyCustomizationModuleChanged();
		bMorphPresetCustomizationRegistered = false;
	}
}

EAssetTypeCategories::Type FFunnyGirlEditorModule::GetLLMAssetCategory()
{
	return LLMAssetCategory;
}

EAssetTypeCategories::Type FFunnyGirlEditorModule::GetNarrativeAssetCategory()
{
	return NarrativeAssetCategory;
}

void FFunnyGirlEditorModule::RegisterAssetTypeAction(const TSharedRef<IAssetTypeActions>& AssetTypeAction)
{
	IAssetTools& AssetTools = FModuleManager::LoadModuleChecked<FAssetToolsModule>("AssetTools").Get();
	AssetTools.RegisterAssetTypeActions(AssetTypeAction);
	RegisteredAssetTypeActions.Add(AssetTypeAction);
}

void FFunnyGirlEditorModule::RegisterMenus()
{
	FToolMenuOwnerScoped OwnerScoped(this);

	UToolMenu* ToolsMenu = UToolMenus::Get()->ExtendMenu(TEXT("LevelEditor.MainMenu.Tools"));
	if (!ToolsMenu)
	{
		return;
	}

	FToolMenuSection& Section = ToolsMenu->FindOrAddSection(TEXT("FunnyGirl"));
	Section.Label = LOCTEXT("FunnyGirlToolsSection", "Funny Girl");
	Section.AddMenuEntry(
		TEXT("ExportMorphPresetDescriptions"),
		LOCTEXT("ExportMorphPresetDescriptionsLabel", "Export Morph Preset Descriptions"),
		LOCTEXT("ExportMorphPresetDescriptionsTooltip", "Export all Morph Preset assets to Saved/Exports/MorphPresetDescriptions.tsv."),
		FSlateIcon(),
		FUIAction(FExecuteAction::CreateRaw(this, &FFunnyGirlEditorModule::ExportMorphPresetDescriptions)));
}

void FFunnyGirlEditorModule::ExportMorphPresetDescriptions()
{
	FAssetRegistryModule& AssetRegistryModule = FModuleManager::LoadModuleChecked<FAssetRegistryModule>(TEXT("AssetRegistry"));
	IAssetRegistry& AssetRegistry = AssetRegistryModule.Get();
	AssetRegistry.SearchAllAssets(true);

	FARFilter Filter;
	Filter.ClassPaths.Add(UMorphPresetDataAsset::StaticClass()->GetClassPathName());
	Filter.bRecursiveClasses = true;

	TArray<FAssetData> PresetAssets;
	AssetRegistry.GetAssets(Filter, PresetAssets);
	PresetAssets.Sort([](const FAssetData& A, const FAssetData& B)
	{
		return A.GetSoftObjectPath().ToString() < B.GetSoftObjectPath().ToString();
	});

	FString Output;
	Output += TEXT("AssetName\tPresetName\tDescription\n");

	for (const FAssetData& AssetData : PresetAssets)
	{
		const UMorphPresetDataAsset* Preset = Cast<UMorphPresetDataAsset>(AssetData.GetAsset());
		if (!Preset)
		{
			continue;
		}

		FString AssetName = AssetData.AssetName.ToString();
		FString PresetName = Preset->DisplayName.ToString().TrimStartAndEnd();
		FString Description = Preset->Description.ToString().TrimStartAndEnd();

		if (PresetName.IsEmpty())
		{
			PresetName = AssetName;
		}

		AssetName.ReplaceInline(TEXT("\t"), TEXT(" "));
		PresetName.ReplaceInline(TEXT("\t"), TEXT(" "));
		Description.ReplaceInline(TEXT("\t"), TEXT(" "));
		Description.ReplaceInline(TEXT("\r\n"), TEXT(" "));
		Description.ReplaceInline(TEXT("\n"), TEXT(" "));
		Description.ReplaceInline(TEXT("\r"), TEXT(" "));

		Output += FString::Printf(TEXT("%s\t%s\t%s\n"), *AssetName, *PresetName, *Description);
	}

	const FString ExportDir = FPaths::Combine(FPaths::ProjectSavedDir(), TEXT("Exports"));
	IFileManager::Get().MakeDirectory(*ExportDir, true);
	const FString ExportPath = FPaths::Combine(ExportDir, TEXT("MorphPresetDescriptions.tsv"));

	if (!FFileHelper::SaveStringToFile(Output, *ExportPath, FFileHelper::EEncodingOptions::ForceUTF8WithoutBOM))
	{
		FMessageDialog::Open(EAppMsgType::Ok, FText::Format(
			LOCTEXT("ExportMorphPresetDescriptionsFailed", "Failed to write Morph Preset descriptions to:\n{0}"),
			FText::FromString(ExportPath)));
		return;
	}

	FMessageDialog::Open(EAppMsgType::Ok, FText::Format(
		LOCTEXT("ExportMorphPresetDescriptionsSuccess", "Exported {0} Morph Preset assets to:\n{1}"),
		FText::AsNumber(PresetAssets.Num()),
		FText::FromString(ExportPath)));
}

#undef LOCTEXT_NAMESPACE

IMPLEMENT_MODULE(FFunnyGirlEditorModule, FunnyGirlEditor)