integration/Plugins/BlueprintMCP/Source/BlueprintMCP/Private/Handlers/UMCPHandler_SearchWithinBlueprints.h

#pragma once

#include "CoreMinimal.h"
#include "MCPHandler.h"
#include "MCPAssetFinder.h"
#include "MCPUtils.h"
#include "Engine/Blueprint.h"
#include "Engine/World.h"
#include "Engine/Level.h"
#include "Engine/LevelScriptBlueprint.h"
#include "EdGraph/EdGraph.h"
#include "EdGraph/EdGraphNode.h"
#include "K2Node_CallFunction.h"
#include "K2Node_Event.h"
#include "K2Node_CustomEvent.h"
#include "K2Node_VariableGet.h"
#include "K2Node_VariableSet.h"
#include "K2Node_BreakStruct.h"
#include "K2Node_MakeStruct.h"
#include "K2Node_FunctionEntry.h"
#include "K2Node_EditablePinBase.h"
#include "AssetRegistry/IAssetRegistry.h"
#include "UMCPHandler_SearchWithinBlueprints.generated.h"


// ---------------------------------------------------------------------------
// ---------------------------------------------------------------------------
// ---------------------------------------------------------------------------

UCLASS()
class UMCPHandler_SearchWithinBlueprints : public UObject, public IMCPHandler
{
	GENERATED_BODY()

public:
	UPROPERTY(meta=(Description="Search query string to match against node titles, function names, event names, and variable names"))
	FString Query;

	UPROPERTY(meta=(Optional, Description="Filter results to blueprints whose path contains this substring"))
	FString Path;

	UPROPERTY(meta=(Optional, Description="Maximum number of results to return (default 50, max 200)"))
	int32 MaxResults = 0;

	virtual FString GetDescription() const override
	{
		return TEXT("Search across all Blueprint graphs for nodes matching a query string.");
	}

	virtual void Handle(const FJsonObject* Json, FJsonObject* Result) override
	{
		int32 EffectiveMaxResults = (MaxResults > 0) ? FMath::Clamp(MaxResults, 1, 200) : 50;

		// Build a combined list of all searchable blueprints (regular + level)
		auto SearchBlueprint = [&](const FString& AssetName, const FString& AssetPath, UBlueprint* BP, TArray<TSharedPtr<FJsonValue>>& OutResults)
		{
			for (UEdGraphNode* Node : MCPUtils::AllNodes(BP))
			{
				if (OutResults.Num() >= EffectiveMaxResults) break;

				FString Title = Node->GetNodeTitle(ENodeTitleType::FullTitle).ToString();

				FString FuncName, EventName, VarName;
				if (auto* CF = Cast<UK2Node_CallFunction>(Node))
				{
					FuncName = CF->FunctionReference.GetMemberName().ToString();
				}
				else if (auto* Ev = Cast<UK2Node_Event>(Node))
				{
					EventName = Ev->EventReference.GetMemberName().ToString();
				}
				else if (auto* CE = Cast<UK2Node_CustomEvent>(Node))
				{
					EventName = CE->CustomFunctionName.ToString();
				}
				else if (auto* VG = Cast<UK2Node_VariableGet>(Node))
				{
					VarName = VG->GetVarName().ToString();
				}
				else if (auto* VS = Cast<UK2Node_VariableSet>(Node))
				{
					VarName = VS->GetVarName().ToString();
				}

				bool bMatch = Title.Contains(Query, ESearchCase::IgnoreCase) ||
					(!FuncName.IsEmpty() && FuncName.Contains(Query, ESearchCase::IgnoreCase)) ||
					(!EventName.IsEmpty() && EventName.Contains(Query, ESearchCase::IgnoreCase)) ||
					(!VarName.IsEmpty() && VarName.Contains(Query, ESearchCase::IgnoreCase));

				if (bMatch)
				{
					TSharedRef<FJsonObject> R = MakeShared<FJsonObject>();
					R->SetStringField(TEXT("blueprint"), AssetName);
					R->SetStringField(TEXT("blueprintPath"), AssetPath);
					R->SetStringField(TEXT("graph"), Node->GetGraph()->GetName());
					R->SetStringField(TEXT("nodeTitle"), Title);
					R->SetStringField(TEXT("nodeClass"), Node->GetClass()->GetName());
					if (!FuncName.IsEmpty()) R->SetStringField(TEXT("functionName"), FuncName);
					if (!EventName.IsEmpty()) R->SetStringField(TEXT("eventName"), EventName);
					if (!VarName.IsEmpty()) R->SetStringField(TEXT("variableName"), VarName);
					OutResults.Add(MakeShared<FJsonValueObject>(R));
				}
			}
		};

		MCPAssets<UBlueprint> AllBlueprints;
		AllBlueprints.Info();
		MCPAssets<UWorld> AllWorlds;
		AllWorlds.Info();

		TArray<TSharedPtr<FJsonValue>> Results;
		for (const FAssetData& Asset : AllBlueprints.AllData())
		{
			if (Results.Num() >= EffectiveMaxResults) break;

			FString AssetPath = Asset.PackageName.ToString();
			if (!Path.IsEmpty() && !AssetPath.Contains(Path, ESearchCase::IgnoreCase))
			{
				continue;
			}

			UBlueprint* BP = Cast<UBlueprint>(const_cast<FAssetData&>(Asset).GetAsset());
			if (!BP) continue;

			SearchBlueprint(Asset.AssetName.ToString(), AssetPath, BP, Results);
		}

		// Also search level blueprints
		for (const FAssetData& MapAsset : AllWorlds.AllData())
		{
			if (Results.Num() >= EffectiveMaxResults) break;

			FString AssetPath = MapAsset.PackageName.ToString();
			if (!Path.IsEmpty() && !AssetPath.Contains(Path, ESearchCase::IgnoreCase))
			{
				continue;
			}

			UWorld* World = Cast<UWorld>(MapAsset.GetAsset());
			if (!World || !World->PersistentLevel) continue;
			ULevelScriptBlueprint* LevelBP = World->PersistentLevel->GetLevelScriptBlueprint(false);
			if (!LevelBP) continue;

			int32 BeforeCount = Results.Num();
			SearchBlueprint(MapAsset.AssetName.ToString(), AssetPath, LevelBP, Results);
			// Tag newly-added entries as level blueprint results
			for (int32 i = BeforeCount; i < Results.Num(); ++i)
			{
				Results[i]->AsObject()->SetBoolField(TEXT("isLevelBlueprint"), true);
			}
		}

		Result->SetNumberField(TEXT("resultCount"), Results.Num());
		Result->SetArrayField(TEXT("results"), Results);
	}
};