integration/Plugins/UEWingman/Source/UEWingman/Private/WingFetcher.cpp

#include "WingFetcher.h"
#include "WingServer.h"
#include "WingUtils.h"
#include "WingActorComponent.h"
#include "Engine/Blueprint.h"
#include "EdGraph/EdGraph.h"
#include "EdGraph/EdGraphNode.h"
#include "EdGraph/EdGraphPin.h"
#include "Engine/SimpleConstructionScript.h"
#include "Engine/SCS_Node.h"
#include "Engine/World.h"
#include "Materials/Material.h"
#include "MaterialGraph/MaterialGraph.h"
#include "MaterialGraph/MaterialGraphNode.h"
#include "IMaterialEditor.h"
#include "Engine/LevelScriptBlueprint.h"
#include "WidgetBlueprint.h"
#include "Blueprint/WidgetTree.h"
#include "Components/Widget.h"
#include "Subsystems/AssetEditorSubsystem.h"
#include "WingServer.h"
#include "WingManual.h"

WingFetcher::WalkFunc WingFetcher::GetWalker(const FString& Step)
{
	if (Step.Equals(TEXT("graph"),          ESearchCase::IgnoreCase)) return &WingFetcher::Graph;
	if (Step.Equals(TEXT("node"),           ESearchCase::IgnoreCase)) return &WingFetcher::Node;
	if (Step.Equals(TEXT("pin"),            ESearchCase::IgnoreCase)) return &WingFetcher::Pin;
	if (Step.Equals(TEXT("component"),      ESearchCase::IgnoreCase)) return &WingFetcher::Component;
	if (Step.Equals(TEXT("widget"),         ESearchCase::IgnoreCase)) return &WingFetcher::Widget;
	if (Step.Equals(TEXT("levelblueprint"), ESearchCase::IgnoreCase)) return &WingFetcher::LevelBlueprint;
	return nullptr;
}

void WingFetcher::SetObj(UObject* InObj)
{
	if (!bSkipNotify)
		UWingServer::AddTouchedObject(InObj);
	Obj = InObj;
	ResultPin = nullptr;
}

void WingFetcher::SetPin(UEdGraphPin* InPin)
{
	ResultPin = InPin;
	Obj = nullptr;
}

WingFetcher& WingFetcher::SetError()
{
	bError = true;
	Obj = nullptr;
	ResultPin = nullptr;
	OriginalAsset = nullptr;
	Editor = nullptr;
	return *this;
}

void WingFetcher::PathFailed(const TCHAR* Expected)
{
	if (ResultPin)
		UWingServer::Printf(TEXT("ERROR: Path specifies a pin, but expected %s\n"), Expected);
	else if (Obj)
		UWingServer::Printf(TEXT("ERROR: Path specifies a %s, but expected %s\n"), *Obj->GetClass()->GetName(), Expected);
	else
		UWingServer::Printf(TEXT("ERROR: Path led to a null pointer\n"));
	UWingServer::SuggestManual(WingManual::Section::Paths);
	SetError();
}

WingFetcher& WingFetcher::TypeMismatch(const TCHAR* Walker, const TCHAR* Expected)
{
	if (ResultPin)
		UWingServer::Printf(TEXT("ERROR: Input to '%s' is a pin, but expected %s\n"), Walker, Expected);
	else if (Obj)
		UWingServer::Printf(TEXT("ERROR: Input to '%s' is %s, but expected %s\n"), Walker, *Obj->GetClass()->GetName(), Expected);
	else
		UWingServer::Printf(TEXT("ERROR: Path led to a null pointer\n"));
	UWingServer::SuggestManual(WingManual::Section::Paths);
	SetError();
	return *this;
}

WingFetcher& WingFetcher::Walk(const FString& Path)
{
	if (bError) return *this;

	if (Path.Contains(TEXT(" ")))
	{
		UWingServer::Printf(TEXT("ERROR: Paths may not contain whitespace."));
		UWingServer::SuggestManual(WingManual::Section::Paths);
		UWingServer::SuggestManual(WingManual::Section::IdentifierSanitization);
		return SetError();
	}
	TArray<FString> Segments;
	Path.ParseIntoArray(Segments, TEXT(","));
	if (Segments.Num() == 0)
	{
		UWingServer::Print(TEXT("ERROR: Empty path\n"));
		UWingServer::SuggestManual(WingManual::Section::Paths);
		return SetError();
	}

	for (int32 i = 0; i < Segments.Num(); i++)
	{
		if (!Obj && !ResultPin)
		{
			Asset(Segments[i]);
			if (bError) return *this;
			continue;
		}

		FString Key, Value;
		if (!Segments[i].Split(TEXT(":"), &Key, &Value))
			Key = Segments[i];

		WalkFunc Func = GetWalker(Key);
		if (!Func)
		{
			UWingServer::Printf(TEXT("ERROR: Unknown path step '%s'\n"), *Key);
			UWingServer::SuggestManual(WingManual::Section::Paths);
			return SetError();
		}
		(this->*Func)(Value);
		if (bError) return *this;
	}

	return *this;
}

WingFetcher& WingFetcher::Asset(const FString& PackagePath)
{
	if (bError) return *this;

	if (!PackagePath.StartsWith(TEXT("/")))
	{
		UWingServer::Printf(TEXT("ERROR: Path must start with '/', got '%s'\n"), *PackagePath);
		UWingServer::SuggestManual(WingManual::Section::Paths);
		return SetError();
	}

	// Check if the package exists before calling LoadObject, because
	// LoadObject logs its own errors when the package doesn't exist.
	FString PackageName = FPackageName::ObjectPathToPackageName(PackagePath);
	if (!FPackageName::DoesPackageExist(PackageName))
	{
		UWingServer::Printf(TEXT("ERROR: Asset '%s' does not exist.\n"), *PackagePath);
		return SetError();
	}
	SetObj(LoadObject<UObject>(nullptr, *PackagePath));
	if (!Obj)
	{
		UWingServer::Printf(TEXT("ERROR: Could not load asset '%s'\n"), *PackagePath);
		return SetError();
	}

	OriginalAsset = Obj;

	// Open the editor for this asset (or bring it to front if already open).
	UAssetEditorSubsystem* Sub = GEditor->GetEditorSubsystem<UAssetEditorSubsystem>();
	if (!Sub || !Sub->OpenEditorForAsset(Obj))
	{
		UWingServer::Printf(TEXT("ERROR: Could not open editor for '%s'\n"), *PackagePath);
		return SetError();
	}
	Editor = Sub->FindEditorForAsset(OriginalAsset, false);
	if (!Editor)
	{
		UWingServer::Printf(TEXT("ERROR: Could not find editor instance for '%s'\n"), *PackagePath);
		return SetError();
	}

	// If this is a material, use the editor's transient copy.
	if (UMaterial* Mat = ::Cast<UMaterial>(Obj))
	{
		IMaterialEditor *MatEditor = static_cast<IMaterialEditor*>(Editor);
		SetObj(MatEditor->GetMaterialInterface()->GetBaseMaterial());
	}
	return *this;
}

bool WingFetcher::CheckAssetIsA(UClass* StaticClass)
{
	if (bError) return false;
	if (!OriginalAsset || !OriginalAsset->IsA(StaticClass))
	{
		UWingServer::Printf(TEXT("ERROR: Asset is %s, expected %s\n"),
			OriginalAsset ? *OriginalAsset->GetClass()->GetName() : TEXT("null"),
			*StaticClass->GetName());
		SetError();
		return false;
	}
	return true;
}

WingFetcher& WingFetcher::Graph(const FString& Value)
{
	if (bError) return *this;

	// Material with blank graph name
	if (UMaterial* Mat = ::Cast<UMaterial>(Obj))
	{
		if (!Value.IsEmpty())
		{
			UWingServer::Printf(TEXT("ERROR: Materials have only one graph, with a blank name.\n\n"));
			UWingServer::SuggestManual(WingManual::Section::Paths);
			return SetError();
		}
		WingUtils::EnsureMaterialGraph(Mat);
		if (!Mat->MaterialGraph)
		{
			UWingServer::Printf(TEXT("ERROR: Material '%s' has no material graph\n"), *Mat->GetName());
			return SetError();
		}
		SetObj(Mat->MaterialGraph);
		return *this;
	}

	UBlueprint* BP = ::Cast<UBlueprint>(Obj);
	if (!BP)
	{
		TypeMismatch(TEXT("graph"), TEXT("Blueprint or Material"));
		return SetError();
	}

	TArray<UEdGraph*> Graphs = WingUtils::AllGraphs(BP);
	UEdGraph* Found = WingUtils::FindExactlyOneNamed(Value, Graphs, TEXT("graph"));
	if (!Found)
	{
		UWingServer::Printf(TEXT("Graphs that exist in blueprint:\n"));
		for (const UEdGraph *G : Graphs)
		{
			UWingServer::Printf(TEXT("  %s\n"), *WingUtils::FormatName(G));
		}
		return SetError();
	}

	SetObj(Found);
	return *this;
}

WingFetcher& WingFetcher::Node(const FString& Value)
{
	if (bError) return *this;

	// If current object is not a graph, refuse.
	UEdGraph *Graph = ::Cast<UEdGraph>(Obj);
	if (Graph == nullptr)
	{
		TypeMismatch(TEXT("node"), TEXT("Graph"));
		return SetError();
	}

	// Get the nodes from the graph.
	TArray<UEdGraphNode *> AllNodes = WingUtils::AllNodes(Graph);
	UEdGraphNode *Node = WingUtils::FindExactlyOneNamed(Value, AllNodes, TEXT("node"));
	if (Node == nullptr)
	{
		UWingServer::Printf(TEXT("Nodes that exist in graph:\n"));
		for (const UEdGraphNode *N : AllNodes)
		{
			UWingServer::Printf(TEXT("  %s\n"), *WingUtils::FormatName(N));
		}
		return SetError();
	}
	SetObj(Node);
	return *this;
}

WingFetcher& WingFetcher::Pin(const FString& Value)
{
	if (bError) return *this;

	UEdGraphNode* N = ::Cast<UEdGraphNode>(Obj);
	if (!N)
	{
		TypeMismatch(TEXT("pin"), TEXT("node"));
		return SetError();
	}
	UEdGraphPin *Found = WingUtils::FindExactlyOneNamed(Value, N->Pins, TEXT("pin"));
	if (!Found)
	{
		UWingServer::Printf(TEXT("Pins that exist in the node:\n"));
		for (const UEdGraphPin *P : N->Pins)
		{
			UWingServer::Printf(TEXT("  %s\n"), *WingUtils::FormatName(P));
		}
		return SetError();
	}
	SetPin(Found);
	return *this;
}

WingFetcher& WingFetcher::Component(const FString& Value)
{
	if (bError) return *this;

	UBlueprint* BP = ::Cast<UBlueprint>(Obj);
	if (!BP)
	{
		TypeMismatch(TEXT("component"), TEXT("Blueprint"));
		return SetError();
	}

	TArray<FWingActorComponent> AllComponents = FWingActorComponent::GetAll(BP);
	FWingActorComponent* Comp = WingUtils::FindExactlyOneNamed(Value, AllComponents, TEXT("component"));
	if (!Comp)
	{
		UWingServer::Printf(TEXT("Components that exist in the blueprint:\n"));
		for (const FWingActorComponent &C : AllComponents)
		{
			UWingServer::Printf(TEXT("  %s\n"), *WingUtils::FormatName(C));
		}
		return SetError();
	}

	if (!Comp->IsOwnedBy(BP))
	{
		UWingServer::Printf(TEXT("ERROR: Component '%s' belongs to %s, to edit it, you must go through that blueprint.\n"),
			*Comp->GetName(), *WingUtils::FormatName(Comp->Owner));
		return SetError();
	}
	SetObj(Comp->SCSNode);
	return *this;
}

WingFetcher& WingFetcher::Widget(const FString& Value)
{
	if (bError) return *this;

	UWidgetBlueprint* WidgetBP = ::Cast<UWidgetBlueprint>(Obj);
	if (!WidgetBP)
	{
		TypeMismatch(TEXT("widget"), TEXT("WidgetBlueprint"));
		return SetError();
	}

	TArray<UWidget*> AllWidgets;
	WidgetBP->WidgetTree->GetAllWidgets(AllWidgets);
	UWidget* Found = WingUtils::FindExactlyOneNamed(Value, AllWidgets, TEXT("widget"));
	if (!Found)
	{
		UWingServer::Printf(TEXT("Widgets that exist in the blueprint:\n"));
		for (const UWidget *W : AllWidgets)
		{
			UWingServer::Printf(TEXT("  %s\n"), *WingUtils::FormatName(W));
		}
		return SetError();
	}
	SetObj(Found);
	return *this;
}

WingFetcher& WingFetcher::LevelBlueprint(const FString& Value)
{
	if (bError) return *this;

	UWorld* World = ::Cast<UWorld>(Obj);
	if (!World)
	{
		TypeMismatch(TEXT("levelblueprint"), TEXT("world"));
		return SetError();
	}

	if (!World->PersistentLevel)
	{
		UWingServer::Print(TEXT("ERROR: World has no PersistentLevel\n"));
		return SetError();
	}

	ULevelScriptBlueprint* LevelBP = World->PersistentLevel->GetLevelScriptBlueprint(true);
	if (!LevelBP)
	{
		UWingServer::Print(TEXT("ERROR: World has no level blueprint\n"));
		return SetError();
	}

	SetObj(LevelBP);
	return *this;
}