integration/Plugins/BlueprintMCP/Source/BlueprintMCP/Private/MCPFetcher.cpp

#include "MCPFetcher.h"
#include "MCPUtils.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 "Engine/LevelScriptBlueprint.h"

MCPFetcher& MCPFetcher::SetError(const FString& Msg)
{
	bError = true;
	ErrorCB.SetError(Msg);
	return *this;
}

MCPFetcher& MCPFetcher::TypeMismatch(const TCHAR* Walker, const TCHAR* Expected)
{
	bError = true;
	if (ResultPin)
		ErrorCB.SetError(FString::Printf(TEXT("Input to '%s' is a pin, expected %s"), Walker, Expected));
	else if (Obj)
		ErrorCB.SetError(FString::Printf(TEXT("Input to '%s' is %s, expected %s"), Walker, *Obj->GetClass()->GetName(), Expected));
	else
		ErrorCB.SetError(FString::Printf(TEXT("Input to '%s' is null, expected %s"), Walker, Expected));
	return *this;
}

const TArray<MCPFetcher::FWalker>& MCPFetcher::GetWalkerTable()
{
	static const TArray<FWalker> Table = {
		{ TEXT("graph"),          TEXT("Find a named UEdGraph (blank name for material graphs)"), &MCPFetcher::Graph },
		{ TEXT("node"),           TEXT("Find a named UEdGraphNode within a graph or blueprint"),  &MCPFetcher::Node },
		{ TEXT("pin"),            TEXT("Find a named UEdGraphPin on a node"),                     &MCPFetcher::Pin },
		{ TEXT("component"),      TEXT("Find a named component in a Blueprint's SCS"),            &MCPFetcher::Component },
		{ TEXT("levelblueprint"), TEXT("Get the level blueprint from a UWorld"),                  &MCPFetcher::LevelBlueprint },
		// { TEXT("matexp"),         TEXT("Get the UMaterialExpression from a UMaterialGraphNode"),  &MCPFetcher::MatExp },
	};
	return Table;
}

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

	TArray<FString> Segments;
	Path.ParseIntoArray(Segments, TEXT(","));
	if (Segments.Num() == 0)
	{
		SetError(TEXT("Empty path"));
		return *this;
	}

	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];

		const FWalker* W = GetWalker(Key);
		if (!W)
		{
			SetError(FString::Printf(TEXT("Unknown path step '%s'"), *Key));
			return *this;
		}
		(this->*W->Func)(Value);
		if (bError) return *this;
	}

	return *this;
}

MCPFetcher& MCPFetcher::Asset(const FString& PackagePath)
{
	SetObj(LoadObject<UObject>(nullptr, *PackagePath));
	if (!Obj)
		SetError(FString::Printf(TEXT("Could not load asset '%s'"), *PackagePath));

	// If this is a material open in the editor, use the editor's transient copy.
	if (UMaterial* Mat = ::Cast<UMaterial>(Obj))
		SetObj(MCPUtils::ReplaceMaterialWithTransientCopy(Mat));
	return *this;
}

const MCPFetcher::FWalker* MCPFetcher::GetWalker(const FString& Key)
{
	for (const FWalker& W : GetWalkerTable())
	{
		if (StrEq(Key, W.Key))
			return &W;
	}
	return nullptr;
}

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

	// Material with blank graph name → navigate to the material graph.
	if (UMaterial* Mat = ::Cast<UMaterial>(Obj))
	{
		if (!Value.IsEmpty())
			return SetError(FString::Printf(TEXT("Materials do not have named graphs (got '%s')"), *Value));
		MCPUtils::EnsureMaterialGraph(Mat);
		if (!Mat->MaterialGraph)
			return SetError(FString::Printf(TEXT("Material '%s' has no material graph"), *Mat->GetName()));
		SetObj(Mat->MaterialGraph);
		return *this;
	}

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

	TArray<UEdGraph*> Matches = MCPUtils::AllGraphsNamed(BP, Value);
	if (Matches.Num() == 0)
		return SetError(FString::Printf(TEXT("Graph '%s' not found in %s"), *Value, *BP->GetName()));
	if (Matches.Num() > 1)
		return SetError(FString::Printf(TEXT("Ambiguous graph '%s' in %s — %d matches"), *Value, *BP->GetName(), Matches.Num()));

	SetObj(Matches[0]);
	return *this;
}

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

	// If current object is a graph, search that graph
	if (UEdGraph* G = ::Cast<UEdGraph>(Obj))
	{
		UEdGraphNode* Found = nullptr;
		for (UEdGraphNode* N : G->Nodes)
		{
			if (!N || !MCPUtils::Identifies(Value, N))
				continue;
			if (Found)
				return SetError(FString::Printf(TEXT("Ambiguous node '%s' in graph %s"), *Value, *G->GetName()));
			Found = N;
		}
		if (!Found)
			return SetError(FString::Printf(TEXT("Node '%s' not found in graph %s"), *Value, *G->GetName()));
		SetObj(Found);
		return *this;
	}

	// If current object is a blueprint, search all graphs
	if (UBlueprint* BP = ::Cast<UBlueprint>(Obj))
	{
		UEdGraphNode* Found = nullptr;
		for (UEdGraph* G : MCPUtils::AllGraphs(BP))
		{
			for (UEdGraphNode* N : G->Nodes)
			{
				if (!N || !MCPUtils::Identifies(Value, N))
					continue;
				if (Found)
					return SetError(FString::Printf(TEXT("Ambiguous node '%s' in %s"), *Value, *BP->GetName()));
				Found = N;
			}
		}
		if (!Found)
			return SetError(FString::Printf(TEXT("Node '%s' not found in %s"), *Value, *BP->GetName()));
		SetObj(Found);
		return *this;
	}

	return TypeMismatch(TEXT("node"), TEXT("graph or Blueprint"));
}

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

	UEdGraphNode* N = ::Cast<UEdGraphNode>(Obj);
	if (!N)
		return TypeMismatch(TEXT("pin"), TEXT("node"));
	UEdGraphPin* Found = nullptr;
	for (UEdGraphPin *P : N->Pins)
	{
		if (!MCPUtils::Identifies(Value, P))
			continue;
		if (Found)
			return SetError(FString::Printf(TEXT("Ambiguous pin '%s' on node %s"),
					*Value, *MCPUtils::FormatName(N)));
		Found = P;
	}
	if (!Found)
		return SetError(FString::Printf(TEXT("Pin '%s' not found on node %s"),
				*Value, *MCPUtils::FormatName(N)));

	SetPin(Found);
	return *this;
}

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

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

	USimpleConstructionScript* SCS = BP->SimpleConstructionScript;
	if (!SCS)
		return SetError(FString::Printf(TEXT("Blueprint %s has no SimpleConstructionScript (not an Actor Blueprint)"), *BP->GetName()));

	FName SearchName(*Value);
	for (USCS_Node* SCSNode : SCS->GetAllNodes())
	{
		if (SCSNode && SCSNode->GetVariableName() == SearchName)
		{
			SetObj(SCSNode->ComponentTemplate);
			return *this;
		}
	}

	return SetError(FString::Printf(TEXT("Component '%s' not found in %s"), *Value, *BP->GetName()));
}

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

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

	if (!World->PersistentLevel)
		return SetError(TEXT("World has no PersistentLevel"));

	ULevelScriptBlueprint* LevelBP = World->PersistentLevel->GetLevelScriptBlueprint(true);
	if (!LevelBP)
		return SetError(TEXT("World has no level blueprint"));

	SetObj(LevelBP);
	return *this;
}

MCPFetcher& MCPFetcher::Template()
{
	if (bError) return *this;
	if (!Obj)
		return SetError(TEXT("Template: object is null"));

	if (UBlueprint* BP = ::Cast<UBlueprint>(Obj))
	{
		if (!BP->GeneratedClass)
			return SetError(FString::Printf(TEXT("Blueprint '%s' has no GeneratedClass"), *Obj->GetName()));
		SetObj(BP->GeneratedClass->GetDefaultObject());
		return *this;
	}

	// Everything else is its own template — no navigation needed.
	return *this;
}

MCPFetcher& MCPFetcher::MatExp(const FString& Value)
{
	if (bError) return *this;

	UMaterialGraphNode* MatNode = ::Cast<UMaterialGraphNode>(Obj);
	if (!MatNode)
		return TypeMismatch(TEXT("matexp"), TEXT("UMaterialGraphNode"));

	if (!MatNode->MaterialExpression)
		return SetError(FString::Printf(TEXT("Node '%s' has no MaterialExpression"), *MCPUtils::FormatName(MatNode)));

	SetObj(MatNode->MaterialExpression);
	return *this;
}