LuckyWorld/Plugins/LuckyMujoco/Source/LuckyMujocoEditor/Private/MujocoExporter.cpp

#include "MujocoExporter.h"
#include "Misc/FileHelper.h"
#include "Misc/Paths.h"
#include "Components/MujocoBodyComponent.h"
#include "XmlFile.h"
#include "Engine/SCS_Node.h"
#include "StaticMeshOperations.h"
#include "Widgets/SWindow.h"
#include "Interfaces/IMainFrameModule.h"
#include "Widgets/Text/STextBlock.h"
#include "Widgets/Notifications/SProgressBar.h"
#include "Widgets/Input/SButton.h"
#include "Framework/Application/SlateApplication.h"
#include "Widgets/Input/SSpinBox.h"
#include "tinyxml2.h"
#include <functional>
#include "Actors/MujocoStaticMeshActor.h"
#include "Engine/LevelScriptActor.h"
#include "Components/MujocoGeomComponent.h"
#include "IDetailsView.h"
#include "Components/MujocoJointComponent.h"
#include "Components/MujocoSiteComponent.h"
#include "Components/MujocoActuatorComponent.h"
#include "Components/MujocoEqualityComponent.h"
#include "Components/MujocoTendonComponent.h"
#include "UObject/PropertyOptional.h"
#include "Structs/MujocoInertial.h"
#include "Structs/MujocoSite.h"
#include "Structs/MujocoGeom.h"
#include "Enums/MujocoEnums.h"
#include "Camera/CameraComponent.h"
#include "GameFramework/SpringArmComponent.h"
#include "Misc/MujocoOptions.h"
#include "Misc/MujocoXmlGenerator.h"

class SActorToMujocoExportWindow : public SWindow
{
	TArray<UObject*> Objects;
	FString ExportFilename;
	TSharedPtr<SButton> ExportButton;
	TSharedPtr<SButton> CancelButton;
	TSharedPtr<IDetailsView> DetailsView;
	TObjectPtr<UMujocoExportOptions> ExportOptions;

public:
	SLATE_BEGIN_ARGS(SActorToMujocoExportWindow) {}
	SLATE_END_ARGS()
	void Construct(const FArguments& InArgs, const TArray<UObject*>& InObjects, FString InFilename)
	{
		Objects = InObjects;
		ExportFilename = InFilename;
		ExportButton = SNew(SButton).Text(FText::FromString(TEXT("Export"))).OnClicked(this, &SActorToMujocoExportWindow::OnExportClicked);
		CancelButton = SNew(SButton).Text(FText::FromString(TEXT("Cancel"))).OnClicked(this, &SActorToMujocoExportWindow::OnCancelClicked);
		FPropertyEditorModule& PropertyModule = FModuleManager::LoadModuleChecked<FPropertyEditorModule>("PropertyEditor");
		FDetailsViewArgs DetailsViewArgs;
		DetailsViewArgs.bAllowSearch = false;
		DetailsViewArgs.bHideSelectionTip = true;
		DetailsViewArgs.bShowOptions = false;
		DetailsViewArgs.bShowModifiedPropertiesOption = false;
		DetailsViewArgs.bShowScrollBar = true;
		DetailsViewArgs.bShowOptions = false;
		DetailsViewArgs.bShowPropertyMatrixButton = false;
		DetailsView = PropertyModule.CreateDetailView(DetailsViewArgs);
		ExportOptions = NewObject<UMujocoExportOptions>();
		ExportOptions->LoadConfig();
		ExportOptions->ModelName = FPaths::GetBaseFilename(ExportFilename);
		DetailsView->SetObject(ExportOptions);
		
		// clang-format off
		TSharedRef<SVerticalBox> MainVerticalBox =
			SNew(SVerticalBox)
			+ SVerticalBox::Slot()
			.HAlign(HAlign_Fill)
			.FillHeight(1.0f)
			.Padding(0, 0)
			[
				DetailsView.ToSharedRef()
			]
			+ SVerticalBox::Slot()
			.AutoHeight()
			.Padding(6, 4)
			.HAlign(HAlign_Right)
			[
				SNew(SHorizontalBox)
				+ SHorizontalBox::Slot()
				.AutoWidth()
				.Padding(5)
				[
					ExportButton.ToSharedRef()
				]
				+ SHorizontalBox::Slot()
				.AutoWidth()
				.Padding(5)
				[
					CancelButton.ToSharedRef()
				]
			];
		SWindow::Construct(SWindow::FArguments()
			.Title(FText::FromString(TEXT("Export to Mujoco")))
			.ClientSize(FVector2D(500, 300))
			.SupportsMaximize(false)
			.SupportsMinimize(false)
			.CreateTitleBar(true)
			.SizingRule(ESizingRule::UserSized)
			.FocusWhenFirstShown(true)
			.ActivationPolicy(EWindowActivationPolicy::Always)
			[
				MainVerticalBox
			]);
		// clang-format on
	}
	FReply OnCancelClicked()
	{
		RequestDestroyWindow();
		return FReply::Handled();
	}
	static void ShowDialog(const TArray<UObject*>& InObjects, FString InFilename)
	{
		if (InObjects.IsEmpty())
		{
			FMessageDialog::Open(EAppMsgType::Ok, FText::FromString(TEXT("No objects selected")));
			return;
		}
		TSharedRef<SActorToMujocoExportWindow> Window = SNew(SActorToMujocoExportWindow, InObjects, InFilename);
		TSharedPtr<SWindow> ParentWindow;
		if (FModuleManager::Get().IsModuleLoaded("MainFrame"))
		{
			IMainFrameModule& MainFrame = FModuleManager::LoadModuleChecked<IMainFrameModule>("MainFrame");
			ParentWindow = MainFrame.GetParentWindow();
		}
		if (ParentWindow.IsValid())
		{
			FSlateApplication::Get().AddModalWindow(Window, ParentWindow.ToSharedRef());
		}
	}
	FReply OnExportClicked()
	{
		ExportOptions->SaveConfig();
		AsyncTask(ENamedThreads::GameThread, [this]() {
			ExportButton->SetEnabled(false);
			CancelButton->SetEnabled(false);
			RequestDestroyWindow();

			TUniquePtr<FMujocoXmlGenerator> Generator = MakeUnique<FMujocoXmlGenerator>();			
			TUniquePtr<tinyxml2::XMLDocument> Doc = Generator->GenerateMujocoXml(ExportOptions, Objects, ExportFilename);

			FString XmlString;
			tinyxml2::XMLPrinter Printer;
			Doc->Print(&Printer);
			FFileHelper::SaveStringToFile(UTF8_TO_TCHAR(Printer.CStr()), *ExportFilename);

			FMessageDialog::Open(EAppMsgType::Ok, FText::FromString(TEXT("Model exported successfully")));
		});
		return FReply::Handled();
	}
};

UMujocoExporter::UMujocoExporter()
{
	SupportedClass = UBlueprint::StaticClass();
	FormatExtension.Add(TEXT("xml"));
	FormatDescription.Add(TEXT("Mujoco Model"));
	bText = true;
}

bool UMujocoExporter::ExportText(const FExportObjectInnerContext* InContext, UObject* InObject, const TCHAR* InType, FOutputDevice& OutAr, FFeedbackContext* InWarn, uint32 InPortFlags)
{
	TArray<UObject*> Objects;
	Objects.Add(InObject);
	SActorToMujocoExportWindow::ShowDialog(Objects, CurrentFilename);
	return true;
}

UMujocoLevelExporter::UMujocoLevelExporter()
{
	SupportedClass = UWorld::StaticClass();
	FormatExtension.Add(TEXT("xml"));
	FormatDescription.Add(TEXT("Mujoco Model"));
	bText = true;
}

bool UMujocoLevelExporter::ExportText(const class FExportObjectInnerContext* Context, UObject* Object, const TCHAR* Type, FOutputDevice& Ar, FFeedbackContext* Warn, uint32 PortFlags)
{
	UWorld* World = Cast<UWorld>(Object);
	if (!World)
	{
		FMessageDialog::Open(EAppMsgType::Ok, FText::FromString(TEXT("No world found")));
		return false;
	}
	ULevel* Level = World->PersistentLevel;
	if (!Level)
	{
		FMessageDialog::Open(EAppMsgType::Ok, FText::FromString(TEXT("No level found")));
		return false;
	}
	TArray<UObject*> Objects;
	for (AActor* Actor : Level->Actors)
	{
		if (!Actor || Actor->IsPendingKillPending())
		{
			continue;
		}
		TInlineComponentArray<UMujocoBodyComponent*> BodyComponents;
		Actor->GetComponents<UMujocoBodyComponent>(BodyComponents);
		if (BodyComponents.Num() == 0)
		{
			continue;
		}
		Objects.Add(Actor);
	}
	SActorToMujocoExportWindow::ShowDialog(Objects, CurrentFilename);
	return true;
}