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

#include "LuckyMujocoEditor.h"
#include "Misc/MessageDialog.h"
#include "Modules/ModuleManager.h"
#include "Interfaces/IPluginManager.h"
#include "Misc/Paths.h"
#include "HAL/PlatformProcess.h"
#include "AssetToolsModule.h"
#include "MujocoComponentDetails.h"
#include "Selection.h"
#include "Engine/StaticMeshActor.h"
#include "Actors/MujocoStaticMeshActor.h"
#include "Components/MujocoBodyComponent.h"
#include "Components/MujocoGeomComponent.h"
#include "LevelEditor.h"

DEFINE_LOG_CATEGORY(LogMujoco);

#define LOCTEXT_NAMESPACE "FLuckyMujocoEditorModule"

void RegisterMujocoMeshActorConversion()
{
	if (UToolMenus* ToolMenus = UToolMenus::Get())
	{
		if (UToolMenu* Menu = ToolMenus->ExtendMenu("LevelEditor.ActorContextMenu"))
		{

			FToolMenuSection& Section = Menu->FindOrAddSection("ActorTypeTools");

			Section.AddMenuEntry("ConvertToMujocoStaticMeshActor",									 //
				FText::FromString("Convert to Mujoco Body"),										 //
				FText::FromString("Convert the selected StaticMeshActor to MujocoStaticMeshActor."), //
				FSlateIcon(FAppStyle::GetAppStyleSetName(), "Icons.Import"),						 //
				FUIAction(FExecuteAction::CreateLambda([]() {
					for (FSelectionIterator It(GEditor->GetSelectedActorIterator()); It; ++It)
					{
						AStaticMeshActor* OldActor = Cast<AStaticMeshActor>(*It);
						if (OldActor)
						{
							UWorld* World = OldActor->GetWorld();
							FActorSpawnParameters Params;
							Params.SpawnCollisionHandlingOverride = ESpawnActorCollisionHandlingMethod::AlwaysSpawn;

							if (AMujocoStaticMeshActor* NewActor = World->SpawnActor<AMujocoStaticMeshActor>(AMujocoStaticMeshActor::StaticClass(), OldActor->GetTransform(), Params))
							{
								TInlineComponentArray<UStaticMeshComponent*> StaticMeshComponents;
								OldActor->GetComponents<UStaticMeshComponent>(StaticMeshComponents);

								UMujocoBodyComponent* MujocoBodyComponent = NewObject<UMujocoBodyComponent>(NewActor, FName(*(OldActor->GetActorLabel() + TEXT("_Body"))));

								NewActor->AddInstanceComponent(MujocoBodyComponent);
								NewActor->SetRootComponent(MujocoBodyComponent);							

								if (StaticMeshComponents.Num() > 0)
								{
									for (UStaticMeshComponent* StaticMeshComponent : StaticMeshComponents)
									{
										if (UStaticMesh* StaticMesh = StaticMeshComponent->GetStaticMesh())
										{
											UMujocoGeomComponent* Geom = NewObject<UMujocoGeomComponent>(NewActor, FName(*(StaticMesh->GetName() + TEXT("_Geom"))));
											Geom->SetupAttachment(MujocoBodyComponent);

											Geom->SetStaticMesh(StaticMesh);
											Geom->SetWorldScale3D(StaticMeshComponent->GetComponentScale());
											Geom->SetCollisionEnabled(StaticMeshComponent->GetCollisionEnabled());

											int32 NumMaterials = StaticMeshComponent->GetNumMaterials();
											for (int32 i = 0; i < NumMaterials; i++)
											{
												Geom->SetMaterial(i, StaticMeshComponent->GetMaterial(i));
											}
											
											FProperty* OldProperty = Geom->GetClass()->FindPropertyByName(GET_MEMBER_NAME_CHECKED(UMujocoGeomComponent, MujocoStaticMesh));
											Geom->PreEditChange(OldProperty);
											Geom->MujocoStaticMesh = StaticMesh;											
											FPropertyChangedEvent PropertyChangedEvent(OldProperty, EPropertyChangeType::ValueSet);
											Geom->PostEditChangeProperty(PropertyChangedEvent);

											NewActor->AddInstanceComponent(Geom);
										}
									}
								}

								FTransform NewTransform = OldActor->GetTransform();
								NewTransform.SetScale3D(FVector::OneVector);
								NewActor->SetActorTransform(NewTransform);
								NewActor->SetActorLabel(OldActor->GetActorLabel());
								NewActor->RegisterAllComponents();
								
								World->DestroyActor(OldActor);								

								if (FLevelEditorModule* LevelEditorModule = FModuleManager::GetModulePtr<FLevelEditorModule>("LevelEditor"))
								{
									LevelEditorModule->BroadcastComponentsEdited();
								}

								GEditor->SelectActor(NewActor, true, true);
							}
						}
					}
				})));
		}
	}
}

void FLuckyMujocoEditorModule::StartupModule()
{
	FPropertyEditorModule& PropertyModule = FModuleManager::LoadModuleChecked<FPropertyEditorModule>(TEXT("PropertyEditor"));

	PropertyModule.RegisterCustomPropertyTypeLayout("MujocoActuatorV2", FOnGetPropertyTypeCustomizationInstance::CreateStatic(&FMujocoPropertyCustomization::MakeInstance));	
	PropertyModule.RegisterCustomPropertyTypeLayout("MujocoEquality", FOnGetPropertyTypeCustomizationInstance::CreateStatic(&FMujocoPropertyCustomization::MakeInstance));
	PropertyModule.RegisterCustomPropertyTypeLayout("MujocoTendonSpatialSite", FOnGetPropertyTypeCustomizationInstance::CreateStatic(&FMujocoPropertyCustomization::MakeInstance));
	PropertyModule.RegisterCustomPropertyTypeLayout("MujocoTendonSpatialGeom", FOnGetPropertyTypeCustomizationInstance::CreateStatic(&FMujocoPropertyCustomization::MakeInstance));
	PropertyModule.RegisterCustomPropertyTypeLayout("MujocoTendonPulley", FOnGetPropertyTypeCustomizationInstance::CreateStatic(&FMujocoPropertyCustomization::MakeInstance));
	PropertyModule.RegisterCustomPropertyTypeLayout("MujocoTendonFixedJoint", FOnGetPropertyTypeCustomizationInstance::CreateStatic(&FMujocoPropertyCustomization::MakeInstance));	

	PropertyModule.NotifyCustomizationModuleChanged();

	IAssetTools& AssetTools = FModuleManager::LoadModuleChecked<FAssetToolsModule>("AssetTools").Get();
	
	RegisterMujocoMeshActorConversion();
}

void FLuckyMujocoEditorModule::ShutdownModule()
{
	if (FModuleManager::Get().IsModuleLoaded("AssetTools"))
	{
		IAssetTools& AssetTools = FModuleManager::GetModuleChecked<FAssetToolsModule>("AssetTools").Get();		
	}

	FPropertyEditorModule& PropertyModule = FModuleManager::LoadModuleChecked<FPropertyEditorModule>(TEXT("PropertyEditor"));
	
	PropertyModule.UnregisterCustomPropertyTypeLayout("MujocoActuatorV2");	
	PropertyModule.UnregisterCustomPropertyTypeLayout("MujocoEquality");
	PropertyModule.UnregisterCustomPropertyTypeLayout("MujocoTendonSpatialSite");
	PropertyModule.UnregisterCustomPropertyTypeLayout("MujocoTendonSpatialGeom");
	PropertyModule.UnregisterCustomPropertyTypeLayout("MujocoTendonPulley");
	PropertyModule.UnregisterCustomPropertyTypeLayout("MujocoTendonFixedJoint");	

	PropertyModule.NotifyCustomizationModuleChanged();
}

#undef LOCTEXT_NAMESPACE

IMPLEMENT_MODULE(FLuckyMujocoEditorModule, LuckyMujocoEditor)