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

#pragma once

#include "CoreMinimal.h"
#include "MCPHandler.h"
#include "MCPAssetFinder.h"
#include "MCPUtils.h"
#include "Materials/Material.h"
#include "MaterialDomain.h"
#include "Materials/MaterialInstanceConstant.h"
#include "Materials/MaterialFunction.h"
#include "Materials/MaterialExpression.h"
#include "Materials/MaterialExpressionScalarParameter.h"
#include "Materials/MaterialExpressionVectorParameter.h"
#include "Materials/MaterialExpressionTextureObjectParameter.h"
#include "Materials/MaterialExpressionTextureSampleParameter2D.h"
#include "Materials/MaterialExpressionStaticSwitchParameter.h"
#include "Materials/MaterialExpressionConstant.h"
#include "Materials/MaterialExpressionConstant3Vector.h"
#include "Materials/MaterialExpressionConstant4Vector.h"
#include "Materials/MaterialExpressionTextureSample.h"
#include "Materials/MaterialExpressionTextureCoordinate.h"
#include "Materials/MaterialExpressionComponentMask.h"
#include "Materials/MaterialExpressionCustom.h"
#include "Materials/MaterialExpressionFunctionInput.h"
#include "Materials/MaterialExpressionFunctionOutput.h"
#include "Materials/MaterialExpressionMaterialFunctionCall.h"
#include "MaterialGraph/MaterialGraph.h"
#include "MaterialGraph/MaterialGraphNode.h"
#include "MaterialGraph/MaterialGraphSchema.h"
#include "Factories/MaterialFactoryNew.h"
#include "Factories/MaterialFunctionFactoryNew.h"
#include "AssetToolsModule.h"
#include "IAssetTools.h"
#include "AssetRegistry/AssetRegistryModule.h"
#include "EdGraph/EdGraph.h"
#include "EdGraph/EdGraphNode.h"
#include "Serialization/JsonReader.h"
#include "Serialization/JsonWriter.h"
#include "Serialization/JsonSerializer.h"
#include "Misc/Guid.h"
#include "Misc/FileHelper.h"
#include "Misc/Paths.h"
#include "UObject/SavePackage.h"
#include "UObject/UObjectIterator.h"
#include "Kismet2/BlueprintEditorUtils.h"
#include "UMCPHandler_ConnectMaterialExpressionPins.generated.h"


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

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

public:
	UPROPERTY(meta=(Optional, Description="Material name or package path (specify this or materialFunction)"))
	FString Material;

	UPROPERTY(meta=(Optional, Description="Material function name or package path (specify this or material)"))
	FString MaterialFunction;

	UPROPERTY(meta=(Description="Node GUID of the source (output) node"))
	FString SourceNode;

	UPROPERTY(meta=(Description="Pin name on the source node"))
	FString SourcePinName;

	UPROPERTY(meta=(Description="Node GUID of the target (input) node"))
	FString TargetNode;

	UPROPERTY(meta=(Description="Pin name on the target node"))
	FString TargetPinName;

	UPROPERTY(meta=(Optional, Description="If true, preview the change without applying it"))
	bool DryRun = false;

	virtual FString GetDescription() const override
	{
		return TEXT("Connect two pins in a material or material function graph.");
	}

	virtual void Handle(const FJsonObject* Json, FJsonObject* Result) override
	{
		if (Material.IsEmpty() && MaterialFunction.IsEmpty())
		{
			return MCPUtils::MakeErrorJson(Result, TEXT("Missing required field: 'material' or 'materialFunction'"));
		}

		// Load material or material function
		UMaterial* MaterialObj = nullptr;
		UMaterialFunction* MatFunc = nullptr;
		FString AssetDisplayName;

		if (!MaterialFunction.IsEmpty())
		{
			MCPAssets<UMaterialFunction> MFAssets;
			if (!MFAssets.Exact(MaterialFunction).Errors(Result).ENone().ETwo().Load()) return;
			MatFunc = MFAssets.Object();
			AssetDisplayName = MatFunc->GetName();
		}
		else
		{
			MCPAssets<UMaterial> MatAssets;
			if (!MatAssets.Exact(Material).Errors(Result).ENone().ETwo().Load()) return;
			MaterialObj = MatAssets.Object();
			AssetDisplayName = MaterialObj->GetName();
		}

		if (MaterialObj) MCPUtils::EnsureMaterialGraph(MaterialObj);
		UEdGraph* Graph = MaterialObj ? (UEdGraph*)MaterialObj->MaterialGraph : (MatFunc ? MatFunc->MaterialGraph : nullptr);
		if (!Graph)
		{
			return MCPUtils::MakeErrorJson(Result, FString::Printf(TEXT("'%s' has no material graph"), *AssetDisplayName));
		}

		// Find source and target nodes by GUID
		UEdGraphNode* SourceGraphNode = nullptr;
		UEdGraphNode* TargetGraphNode = nullptr;

		for (UEdGraphNode* GraphNode : Graph->Nodes)
		{
			if (!GraphNode) continue;
			if (GraphNode->NodeGuid.ToString() == SourceNode)
				SourceGraphNode = GraphNode;
			if (GraphNode->NodeGuid.ToString() == TargetNode)
				TargetGraphNode = GraphNode;
			if (SourceGraphNode && TargetGraphNode)
				break;
		}

		if (!SourceGraphNode)
		{
			return MCPUtils::MakeErrorJson(Result, FString::Printf(TEXT("Source node '%s' not found in material graph"), *SourceNode));
		}
		if (!TargetGraphNode)
		{
			return MCPUtils::MakeErrorJson(Result, FString::Printf(TEXT("Target node '%s' not found in material graph"), *TargetNode));
		}

		// Find pins
		UEdGraphPin* SourcePin = SourceGraphNode->FindPin(FName(*SourcePinName));
		if (!SourcePin)
		{
			// List available pins for debugging
			TArray<TSharedPtr<FJsonValue>> PinNames;
			for (UEdGraphPin* P : SourceGraphNode->Pins)
			{
				if (P) PinNames.Add(MakeShared<FJsonValueString>(
					FString::Printf(TEXT("%s (%s)"), *P->PinName.ToString(),
						P->Direction == EGPD_Input ? TEXT("Input") : TEXT("Output"))));
			}
			MCPUtils::MakeErrorJson(Result, FString::Printf(TEXT("Source pin '%s' not found on node '%s'"),
				*SourcePinName, *SourceNode));
			Result->SetArrayField(TEXT("availablePins"), PinNames);
			return;
		}

		UEdGraphPin* TargetPin = TargetGraphNode->FindPin(FName(*TargetPinName));
		if (!TargetPin)
		{
			TArray<TSharedPtr<FJsonValue>> PinNames;
			for (UEdGraphPin* P : TargetGraphNode->Pins)
			{
				if (P) PinNames.Add(MakeShared<FJsonValueString>(
					FString::Printf(TEXT("%s (%s)"), *P->PinName.ToString(),
						P->Direction == EGPD_Input ? TEXT("Input") : TEXT("Output"))));
			}
			MCPUtils::MakeErrorJson(Result, FString::Printf(TEXT("Target pin '%s' not found on node '%s'"),
				*TargetPinName, *TargetNode));
			Result->SetArrayField(TEXT("availablePins"), PinNames);
			return;
		}

		if (DryRun)
		{
			UE_LOG(LogTemp, Display, TEXT("BlueprintMCP: [DRY RUN] Would connect %s.%s -> %s.%s in '%s'"),
				*SourceNode, *SourcePinName, *TargetNode, *TargetPinName, *AssetDisplayName);

			Result->SetBoolField(TEXT("dryRun"), true);
			Result->SetStringField(TEXT("material"), AssetDisplayName);
			return;
		}

		UE_LOG(LogTemp, Display, TEXT("BlueprintMCP: Connecting %s.%s -> %s.%s in '%s'"),
			*SourceNode, *SourcePinName, *TargetNode, *TargetPinName, *AssetDisplayName);

		// Try to connect via the schema
		const UEdGraphSchema* Schema = Graph->GetSchema();
		if (!Schema)
		{
			return MCPUtils::MakeErrorJson(Result, TEXT("Material graph schema not found"));
		}

		bool bConnected = Schema->TryCreateConnection(SourcePin, TargetPin);

		Result->SetStringField(TEXT("material"), AssetDisplayName);

		if (!bConnected)
		{
			return MCPUtils::MakeErrorJson(Result, FString::Printf(
				TEXT("Cannot connect %s.%s to %s.%s — types may be incompatible"),
				*SourceNode, *SourcePinName, *TargetNode, *TargetPinName));
		}

		// Save
		UObject* Asset = MaterialObj ? (UObject*)MaterialObj : (UObject*)MatFunc;
		Asset->PreEditChange(nullptr);
		Asset->PostEditChange();
		bool bSaved = MaterialObj ? MCPUtils::SaveMaterialPackage(MaterialObj) : MCPUtils::SaveGenericPackage(MatFunc);
		Result->SetBoolField(TEXT("saved"), bSaved);
	}
};