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

#include "FunnyGirlEditorModule.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 "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"

#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;
}

void FFunnyGirlEditorModule::ShutdownModule()
{
	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);
}

#undef LOCTEXT_NAMESPACE

IMPLEMENT_MODULE(FFunnyGirlEditorModule, FunnyGirlEditor)