integration/Plugins/BlueprintMCP/Source/BlueprintMCP/Public/MCPFetcher.h

#pragma once

#include "CoreMinimal.h"
#include "MCPUtils.h"

class UEdGraphPin;
class IAssetEditorInstance;
struct FWalker;

// MCPFetcher: Load an Asset and find an object within it.
// To find an object, you use a path.  This is typical:
//
//  F.Walk(TEXT("/Game/Mat/M_Test,graph,node:Param_1"))
//
// A path always starts from an asset name.  The path above
// starts at a material asset, then it walks to the material
// graph, then from there to a specific graph node.
//
// Instead of specifying the path as a string, you can also
// specify it using a sequence of procedural steps, like
// this:
//
//   F.Asset(TEXT("/Game/Materials/M_Test"));
//   F.Graph();
//   F.Node(TEXT("Param_1"));
//
// When you're finally at the object you want, you usually
// use the Cast method to get a pointer to the object.
//
// If any step fails, the MCPFetcher will print an error
// message that can be seen by the MCP's caller.  It will
// also set an error flag.  Once the error flag is set, all
// further ops become no-ops.  At that point, attempting a
// Cast will return nullptr.
//

class MCPFetcher
{
public:	
	// Walk a path from an asset to an object
	// within that asset.  If you call walk a
	// second time, it will walk additional steps.
	//
	MCPFetcher& Walk(const FString& Path);

	// Walk a path using individual path
	// steps instead of a path.  All these steps generate
	// errors if they cannot find the desired element.
	//
	MCPFetcher& Asset(const FString& PackagePath);
	MCPFetcher& Graph(const FString& Value);
	MCPFetcher& Node(const FString& Value);
	MCPFetcher& Pin(const FString& Value);
	MCPFetcher& Component(const FString& Value);
	MCPFetcher& LevelBlueprint(const FString& Value);

	// The following walkers cannot be invoked from
	// paths, only procedurally.  If the current object
	// is already the target type, they do nothing.
	// Otherwise, they attempt to convert (e.g. World
	// to its level blueprint, Material to its graph).
	//
	MCPFetcher& ToBlueprint();
	MCPFetcher& ToGraph();
	
	// Return true if there haven't been any errors.
	// Note that errors always automatically generate
	// output to MCPServer::Printf.
	//
	bool Ok() const { return !bError; }
	
	// Try to fetch the current object as a UObject of
	// the specified type.  If it isn't one, generates an
	// error and returns nullptr.
	//
	template<class T> T *Cast()
	{
		if (bError) return nullptr;
		T* Result = ::Cast<T>(Obj);
		if (Result == nullptr)
			TypeMismatch(TEXT("Cast"), *T::StaticClass()->GetName());
		return Result;
	}

	// Get the current object as a UObject if it is one,
	// otherwise nullptr.  Does not generate errors.
	// 
	UObject* GetObj() const { return Obj; }

	// Get the asset from where it all began: the first
	// step in the walk path.  If the asset couldn't be
	// loaded, returns nullptr.  Does not generate errors.
	// 
	UObject* GetAsset() const { return OriginalAsset; }
	template<class T> T* CastAsset()
	{
		if (!CheckAssetIsA(T::StaticClass())) return nullptr;
		return ::Cast<T>(OriginalAsset);
	}

	// When an asset is loaded, an editor is automatically
	// opened.  Get the editor.  You must specify the type
	// that you expect the asset to be, and the type to cast
	// the editor to.  Does not generate errors.
	// 
	template<class AssetType, class EditorType>
	EditorType* CastEditor()
	{
		if (!CheckAssetIsA(AssetType::StaticClass())) return nullptr;
		return static_cast<EditorType*>(Editor);
	}

	// Initialize empty.  You need to call Asset, or walk
	// a path that starts with an asset.
	//
	MCPFetcher() {}

	// Initialize with an object.  From there, you can walk
	// to sub-objects.
	//
	MCPFetcher(UObject* O) : Obj(O) {}

	// Print out the documentation for paths, for the LLM.
	//
	static void PrintDocs();

private:
	
	// The Current Object.  Only one of these can be non-null.
	UObject* Obj = nullptr;
	UEdGraphPin* ResultPin = nullptr;

	// The Starting Asset and the Editor we Opened.
	UObject* OriginalAsset = nullptr;
	IAssetEditorInstance* Editor = nullptr;

	// True if an error has occurred.
	bool bError = false;

	// Internal methods.
	using WalkFunc = MCPFetcher& (MCPFetcher::*)(const FString&);
	void SetObj(UObject* InObj);
	void SetPin(UEdGraphPin* InPin);
	MCPFetcher& SetError();
	MCPFetcher& TypeMismatch(const TCHAR* Walker, const TCHAR* Expected);
	bool CheckAssetIsA(UClass* StaticClass);
	WalkFunc GetWalker(const FString &Step);
};

template<> inline UEdGraphPin* MCPFetcher::Cast<UEdGraphPin>()
{
	if (bError) return nullptr;
	if (!ResultPin)
		TypeMismatch(TEXT("Cast"), TEXT("pin"));
	return ResultPin;
}