#pragma once

#include "CoreMinimal.h"
#include "MCPHandler.h"
#include "MCPAssets.h"
#include "MCPUtils.h"
#include "Materials/Material.h"
#include "Materials/MaterialFunction.h"
#include "Materials/MaterialExpression.h"
#include "MaterialGraph/MaterialGraph.h"
#include "EdGraph/EdGraph.h"
#include "EdGraph/EdGraphNode.h"
#include "UMCPHandler_ConnectMaterialExpressionPins.generated.h"


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

USTRUCT()
struct FConnectMaterialPinsEntry
{
	GENERATED_BODY()

	UPROPERTY()
	FString SourceNode;

	UPROPERTY()
	FString SourcePin;

	UPROPERTY()
	FString TargetNode;

	UPROPERTY()
	FString TargetPin;
};


UCLASS(meta=(Group="Unclassified"))
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="Array of {sourceNode, sourcePin, targetNode, targetPin} objects"))
	FMCPJsonArray Connections;

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

	virtual void Handle(const FJsonObject* Json, FStringBuilderBase& Result) override
	{
		if (Material.IsEmpty() && MaterialFunction.IsEmpty())
		{
			Result.Append(TEXT("ERROR: Specify 'material' or 'materialFunction'.\n"));
			return;
		}

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

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

		if (MaterialObj) MCPUtils::EnsureMaterialGraph(MaterialObj);
		UEdGraph* Graph = MaterialObj ? (UEdGraph*)MaterialObj->MaterialGraph : (MatFunc ? MatFunc->MaterialGraph : nullptr);
		if (!Graph)
		{
			Result.Appendf(TEXT("ERROR: %s has no material graph.\n"),
				MaterialObj ? *MCPUtils::FormatName(MaterialObj) : *MCPUtils::FormatName(MatFunc));
			return;
		}

		int32 SuccessCount = 0;
		const UEdGraphSchema* Schema = Graph->GetSchema();

		for (const TSharedPtr<FJsonValue>& ConnVal : Connections.Array)
		{
			FConnectMaterialPinsEntry Entry;
			if (!MCPUtils::PopulateFromJson(FConnectMaterialPinsEntry::StaticStruct(), &Entry, ConnVal, MCPErrorCallback(Result)))
				continue;

			// Find source node
			UEdGraphNode* SrcNode = FindNodeByName(Graph, Entry.SourceNode);
			if (!SrcNode)
			{
				Result.Appendf(TEXT("ERROR: Source node '%s' not found.\n"), *Entry.SourceNode);
				continue;
			}

			// Find target node
			UEdGraphNode* TgtNode = FindNodeByName(Graph, Entry.TargetNode);
			if (!TgtNode)
			{
				Result.Appendf(TEXT("ERROR: Target node '%s' not found.\n"), *Entry.TargetNode);
				continue;
			}

			// Find source pin
			UEdGraphPin* SrcPin = FindPinByName(SrcNode, Entry.SourcePin);
			if (!SrcPin)
			{
				Result.Appendf(TEXT("ERROR: Pin '%s' not found on %s. Available:"),
					*Entry.SourcePin, *MCPUtils::FormatName(SrcNode));
				for (UEdGraphPin* P : SrcNode->Pins)
					if (P) Result.Appendf(TEXT(" %s"), *MCPUtils::FormatName(P));
				Result.Append(TEXT("\n"));
				continue;
			}

			// Find target pin
			UEdGraphPin* TgtPin = FindPinByName(TgtNode, Entry.TargetPin);
			if (!TgtPin)
			{
				Result.Appendf(TEXT("ERROR: Pin '%s' not found on %s. Available:"),
					*Entry.TargetPin, *MCPUtils::FormatName(TgtNode));
				for (UEdGraphPin* P : TgtNode->Pins)
					if (P) Result.Appendf(TEXT(" %s"), *MCPUtils::FormatName(P));
				Result.Append(TEXT("\n"));
				continue;
			}

			// Check compatibility
			const FPinConnectionResponse Response = Schema->CanCreateConnection(SrcPin, TgtPin);
			if (Response.Response == CONNECT_RESPONSE_DISALLOW)
			{
				Result.Appendf(TEXT("ERROR: Cannot connect %s.%s -> %s.%s: %s\n"),
					*MCPUtils::FormatName(SrcNode), *MCPUtils::FormatName(SrcPin),
					*MCPUtils::FormatName(TgtNode), *MCPUtils::FormatName(TgtPin),
					*Response.Message.ToString());
				continue;
			}

			Schema->TryCreateConnection(SrcPin, TgtPin);
			SuccessCount++;
		}

		if (SuccessCount > 0)
		{
			UObject* Asset = MaterialObj ? (UObject*)MaterialObj : (UObject*)MatFunc;
			Asset->PreEditChange(nullptr);
			Asset->PostEditChange();
			bool bSaved = MaterialObj ? MCPUtils::SaveMaterialPackage(MaterialObj) : MCPUtils::SaveGenericPackage(MatFunc);
			Result.Appendf(TEXT("Connected %d/%d. Saved: %s\n"),
				SuccessCount, Connections.Array.Num(), bSaved ? TEXT("yes") : TEXT("no"));
		}
		else
		{
			Result.Appendf(TEXT("0/%d connections succeeded.\n"), Connections.Array.Num());
		}
	}

private:
	UEdGraphNode* FindNodeByName(UEdGraph* Graph, const FString& Name)
	{
		for (UEdGraphNode* Node : Graph->Nodes)
			if (Node && MCPUtils::Identifies(Name, Node))
				return Node;
		return nullptr;
	}

	UEdGraphPin* FindPinByName(UEdGraphNode* Node, const FString& Name)
	{
		for (UEdGraphPin* Pin : Node->Pins)
			if (Pin && MCPUtils::Identifies(Name, Pin))
				return Pin;
		return nullptr;
	}
};