BallisticsDocs/Source/EasyBallisticsEditor/Private/EasyBallisticsEditor.cpp

// Copyright 2016 Mookie. All Rights Reserved.

#include "EasyBallisticsEditor.h"
#include "AssetToolsModule.h"
#include "PropertyEditorModule.h"
#include "EditorModeRegistry.h"
#include "ComponentVisualizer.h"
#include "UnrealEdGlobals.h"
#include "Editor/UnrealEdEngine.h"

// Asset Factories
#include "EBMaterialResponseMapFactory.h"
#include "EBBulletPropertiesFactory.h"
#include "EBBulletActorFactory.h"
#include "EBMathematicalBallisticsFactory.h"
#include "EBWeaponConfigurationFactory.h"
#include "EBWeaponDataFactory.h"
#include "EBGunActorFactory.h"

// Component Visualizers
#include "EBBarrelComponentFactory.h"
#include "EBMagazineComponentFactory.h"

// Customizations
#include "EBPhysicalMaterialCustomization.h"
#include "EBPlayerWeaponControllerCustomization.h"
#include "EBJsonImportExportTool.h"

// Component Classes
#include "EBBarrel.h"
#include "EBMagazine.h"
#include "EBPlayerWeaponController.h"
#include "PhysicalMaterials/PhysicalMaterial.h"
#include "Editor.h"

#define LOCTEXT_NAMESPACE "FEasyBallisticsEditorModule"

EAssetTypeCategories::Type FEasyBallisticsEditorModule::BallisticsAssetCategory;

void FEasyBallisticsEditorModule::StartupModule()
{
	// Register asset factories
	IAssetTools& AssetTools = FModuleManager::LoadModuleChecked<FAssetToolsModule>("AssetTools").Get();
	
	// Register the Ballistics category
	BallisticsAssetCategory = AssetTools.RegisterAdvancedAssetCategory(FName(TEXT("Ballistics")), LOCTEXT("BallisticsCategory", "Ballistics"));
	
	// Store references to our asset type actions for cleanup
	RegisteredAssetTypeActions.Empty();
	
	// Register Material Response Map factory
	{
		TSharedRef<IAssetTypeActions> MaterialResponseMapActions = MakeShareable(new FEBMaterialResponseMapFactory());
		AssetTools.RegisterAssetTypeActions(MaterialResponseMapActions);
		RegisteredAssetTypeActions.Add(MaterialResponseMapActions);
	}
	
	// Register Bullet Properties Asset factory
	{
		TSharedRef<IAssetTypeActions> BulletPropertiesActions = MakeShareable(new FEBBulletPropertiesAssetFactory());
		AssetTools.RegisterAssetTypeActions(BulletPropertiesActions);
		RegisteredAssetTypeActions.Add(BulletPropertiesActions);
	}
	
	// Register Material Properties Asset factory
	{
		TSharedRef<IAssetTypeActions> MaterialPropertiesActions = MakeShareable(new FEBMaterialPropertiesAssetFactory());
		AssetTools.RegisterAssetTypeActions(MaterialPropertiesActions);
		RegisteredAssetTypeActions.Add(MaterialPropertiesActions);
	}
	
	// Register Mathematical Ballistics Asset factory
	{
		TSharedRef<IAssetTypeActions> MathematicalBallisticsActions = MakeShareable(new FEBMathematicalBallisticsFactory());
		AssetTools.RegisterAssetTypeActions(MathematicalBallisticsActions);
		RegisteredAssetTypeActions.Add(MathematicalBallisticsActions);
	}
	
	// Register Weapon Configuration Asset factory
	{
		TSharedRef<IAssetTypeActions> WeaponConfigurationActions = MakeShareable(new FEBWeaponConfigurationFactory());
		AssetTools.RegisterAssetTypeActions(WeaponConfigurationActions);
		RegisteredAssetTypeActions.Add(WeaponConfigurationActions);
	}
	
	// Register Weapon Data Asset factory
	{
		TSharedRef<IAssetTypeActions> WeaponDataActions = MakeShareable(new FEBWeaponDataFactory());
		AssetTools.RegisterAssetTypeActions(WeaponDataActions);
		RegisteredAssetTypeActions.Add(WeaponDataActions);
	}
	
	// Register UFactory objects (these handle the "Add" menu in Content Browser)
	// These need to be manually registered to appear in the Add menu
	
	// Register Material Response Map UFactory
	if (!GetDefault<UEBMaterialResponseMapFactory>())
	{
		NewObject<UEBMaterialResponseMapFactory>();
	}
	
	// Register Bullet Properties UFactory  
	if (!GetDefault<UEBBulletPropertiesAssetFactory>())
	{
		NewObject<UEBBulletPropertiesAssetFactory>();
	}
	
	// Register Material Properties UFactory
	if (!GetDefault<UEBMaterialPropertiesAssetFactory>())
	{
		NewObject<UEBMaterialPropertiesAssetFactory>();
	}
	
	// Register Mathematical Ballistics UFactory
	if (!GetDefault<UEBMathematicalBallisticsFactory>())
	{
		NewObject<UEBMathematicalBallisticsFactory>();
	}
	
	// Register Weapon Configuration UFactory
	if (!GetDefault<UEBWeaponConfigurationFactory>())
	{
		NewObject<UEBWeaponConfigurationFactory>();
	}
	
	// Register Weapon Data UFactory
	if (!GetDefault<UEBWeaponDataFactory>())
	{
		NewObject<UEBWeaponDataFactory>();
	}

	// Register UObject-based factories for actors
	if (GEditor)
	{
		// Register Weapon Configuration factory
		UEBWeaponConfigurationFactory* WeaponConfigFactory = NewObject<UEBWeaponConfigurationFactory>();
		// UFactory objects are automatically discovered by the editor
		
		// Register Gun Actor factory
		UEBGunActorFactory* GunActorFactory = NewObject<UEBGunActorFactory>();
		GEditor->ActorFactories.Add(GunActorFactory);
		
		// Register Bullet Actor factory
		UEBBulletActorFactory* BulletActorFactory = NewObject<UEBBulletActorFactory>();
		GEditor->ActorFactories.Add(BulletActorFactory);
	}
	
	// Register Component Visualizers
	if (GUnrealEd)
	{
		// Register Barrel Component Visualizer
		TSharedPtr<FEBBarrelComponentVisualizer> BarrelVisualizer = MakeShareable(new FEBBarrelComponentVisualizer());
		GUnrealEd->RegisterComponentVisualizer(UEBBarrel::StaticClass()->GetFName(), BarrelVisualizer);
		RegisteredComponentClassNames.Add(UEBBarrel::StaticClass()->GetFName());
		
		// Register Magazine Component Visualizer
		TSharedPtr<FEBMagazineComponentVisualizer> MagazineVisualizer = MakeShareable(new FEBMagazineComponentVisualizer());
		GUnrealEd->RegisterComponentVisualizer(UEBMagazine::StaticClass()->GetFName(), MagazineVisualizer);
		RegisteredComponentClassNames.Add(UEBMagazine::StaticClass()->GetFName());
	}
	
	// Register Physical Material customization
	FPropertyEditorModule& PropertyModule = FModuleManager::LoadModuleChecked<FPropertyEditorModule>("PropertyEditor");
	PropertyModule.RegisterCustomClassLayout(
		UPhysicalMaterial::StaticClass()->GetFName(),
		FOnGetDetailCustomizationInstance::CreateStatic(&FEBPhysicalMaterialCustomization::MakeInstance)
	);
	
	// Register Player Weapon Controller customization
	PropertyModule.RegisterCustomClassLayout(
		UEBPlayerWeaponController::StaticClass()->GetFName(),
		FOnGetDetailCustomizationInstance::CreateStatic(&FEBPlayerWeaponControllerCustomization::MakeInstance)
	);

	// Register JSON Import/Export tool - delay this until after editor is fully loaded
	FCoreDelegates::OnFEngineLoopInitComplete.AddLambda([]()
	{
		FEBJsonImportExportTool::RegisterMenus();
	});
}

void FEasyBallisticsEditorModule::ShutdownModule()
{
	// Unregister JSON Import/Export tool
	FEBJsonImportExportTool::UnregisterMenus();

	// Unregister Component Visualizers
	if (GUnrealEd)
	{
		for (const FName& ClassName : RegisteredComponentClassNames)
		{
			GUnrealEd->UnregisterComponentVisualizer(ClassName);
		}
	}
	RegisteredComponentClassNames.Empty();

	// Unregister customizations
	if (FModuleManager::Get().IsModuleLoaded("PropertyEditor"))
	{
		FPropertyEditorModule& PropertyModule = FModuleManager::GetModuleChecked<FPropertyEditorModule>("PropertyEditor");
		PropertyModule.UnregisterCustomClassLayout(UPhysicalMaterial::StaticClass()->GetFName());
		PropertyModule.UnregisterCustomClassLayout(UEBPlayerWeaponController::StaticClass()->GetFName());
	}
	
	// Unregister asset type actions
	if (FModuleManager::Get().IsModuleLoaded("AssetTools"))
	{
		IAssetTools& AssetTools = FModuleManager::GetModuleChecked<FAssetToolsModule>("AssetTools").Get();
		
		for (auto& AssetTypeAction : RegisteredAssetTypeActions)
		{
			if (AssetTypeAction.IsValid())
			{
				AssetTools.UnregisterAssetTypeActions(AssetTypeAction.ToSharedRef());
			}
		}
	}
	
	RegisteredAssetTypeActions.Empty();
}

#undef LOCTEXT_NAMESPACE

IMPLEMENT_MODULE(FEasyBallisticsEditorModule, EasyBallisticsEditor)