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

#pragma once

#include "CoreMinimal.h"
#include "MCPHandler.h"
#include "MCPAssetFinder.h"
#include "MCPUtils.h"
#include "Engine/Blueprint.h"
#include "Engine/SimpleConstructionScript.h"
#include "Engine/SCS_Node.h"
#include "Components/ActorComponent.h"
#include "Kismet2/BlueprintEditorUtils.h"
#include "UObject/UObjectIterator.h"
#include "MCPHandlers_Components.generated.h"

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

UCLASS(meta=(ToolName="list_blueprint_components"))
class UMCPHandler_ListBlueprintComponents : public UObject, public IMCPHandler
{
	GENERATED_BODY()

public:
	UPROPERTY(meta=(Description="Blueprint name or package path"))
	FString Blueprint;

	virtual FString GetDescription() const override
	{
		return TEXT("List all components in a Blueprint's SimpleConstructionScript, "
			"including hierarchy and scene root information.");
	}

	virtual void Handle(const FJsonObject* Json, FJsonObject* Result) override
	{
		MCPAssets<UBlueprint> Assets;
		if (!Assets.Exact(Blueprint).Errors(Result).ENone().ETwo().Load()) return;
		UBlueprint* BP = Assets.Object();

		USimpleConstructionScript* SCS = BP->SimpleConstructionScript;
		if (!SCS)
		{
			return MCPUtils::MakeErrorJson(Result, FString::Printf(
				TEXT("Blueprint '%s' does not have a SimpleConstructionScript (not an Actor Blueprint)"),
				*Blueprint));
		}

		const TArray<USCS_Node*>& AllNodes = SCS->GetAllNodes();

		TArray<TSharedPtr<FJsonValue>> ComponentsArr;
		for (USCS_Node* Node : AllNodes)
		{
			if (!Node)
			{
				continue;
			}

			TSharedRef<FJsonObject> CompObj = MakeShared<FJsonObject>();
			CompObj->SetStringField(TEXT("name"), Node->GetVariableName().ToString());

			if (Node->ComponentClass)
			{
				CompObj->SetStringField(TEXT("componentClass"), Node->ComponentClass->GetName());
			}
			else
			{
				CompObj->SetStringField(TEXT("componentClass"), TEXT("None"));
			}

			// Parent component info
			USCS_Node* ParentNode = nullptr;
			for (USCS_Node* Candidate : AllNodes)
			{
				if (Candidate && Candidate->GetChildNodes().Contains(Node))
				{
					ParentNode = Candidate;
					break;
				}
			}

			if (ParentNode)
			{
				CompObj->SetStringField(TEXT("parentComponent"), ParentNode->GetVariableName().ToString());
			}

			// Check if this is a default scene root (first root node with SceneComponent class)
			bool bIsSceneRoot = false;
			const TArray<USCS_Node*>& RootNodes = SCS->GetRootNodes();
			if (RootNodes.Num() > 0 && RootNodes[0] == Node)
			{
				bIsSceneRoot = true;
			}
			CompObj->SetBoolField(TEXT("isSceneRoot"), bIsSceneRoot);

			// List child count for informational purposes
			CompObj->SetNumberField(TEXT("childCount"), Node->GetChildNodes().Num());

			ComponentsArr.Add(MakeShared<FJsonValueObject>(CompObj));
		}

		Result->SetNumberField(TEXT("count"), ComponentsArr.Num());
		Result->SetArrayField(TEXT("components"), ComponentsArr);
	}
};

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

UCLASS(meta=(ToolName="add_blueprint_component"))
class UMCPHandler_AddBlueprintComponent : public UObject, public IMCPHandler
{
	GENERATED_BODY()

public:
	UPROPERTY(meta=(Description="Blueprint name or package path"))
	FString Blueprint;

	UPROPERTY(meta=(Description="Component class name (e.g. StaticMeshComponent, SceneComponent)"))
	FString ComponentClass;

	UPROPERTY(meta=(Description="Component name for the new component"))
	FString Component;

	UPROPERTY(meta=(Optional, Description="Name of the parent component to attach to"))
	FString ParentComponent;

	virtual FString GetDescription() const override
	{
		return TEXT("Add a component to a Blueprint's SimpleConstructionScript. "
			"Optionally attach it to an existing parent component.");
	}

	virtual void Handle(const FJsonObject* Json, FJsonObject* Result) override
	{
		MCPAssets<UBlueprint> Assets;
		if (!Assets.Exact(Blueprint).Errors(Result).ENone().ETwo().Load()) return;
		UBlueprint* BP = Assets.Object();

		USimpleConstructionScript* SCS = BP->SimpleConstructionScript;
		if (!SCS)
		{
			return MCPUtils::MakeErrorJson(Result, FString::Printf(
				TEXT("Blueprint '%s' does not have a SimpleConstructionScript (not an Actor Blueprint)"),
				*Blueprint));
		}

		// Check for duplicate component names
		const TArray<USCS_Node*>& ExistingNodes = SCS->GetAllNodes();
		for (USCS_Node* Existing : ExistingNodes)
		{
			if (Existing && Existing->GetVariableName().ToString().Equals(Component, ESearchCase::IgnoreCase))
			{
				return MCPUtils::MakeErrorJson(Result, FString::Printf(
					TEXT("A component named '%s' already exists in Blueprint '%s'"),
					*Component, *Blueprint));
			}
		}

		// Resolve the component class by name
		// Try multiple name variants: exact name, with U prefix, without U prefix
		UClass* ComponentClassObj = nullptr;

		TArray<FString> NamesToTry;
		NamesToTry.Add(ComponentClass);
		if (!ComponentClass.StartsWith(TEXT("U")))
		{
			NamesToTry.Add(FString::Printf(TEXT("U%s"), *ComponentClass));
		}
		else
		{
			// Also try without U prefix
			NamesToTry.Add(ComponentClass.Mid(1));
		}

		for (TObjectIterator<UClass> It; It; ++It)
		{
			if (!It->IsChildOf(UActorComponent::StaticClass()))
			{
				continue;
			}

			FString ClassName = It->GetName();
			for (const FString& NameToTry : NamesToTry)
			{
				if (ClassName.Equals(NameToTry, ESearchCase::IgnoreCase))
				{
					ComponentClassObj = *It;
					break;
				}
			}

			if (ComponentClassObj)
			{
				break;
			}
		}

		if (!ComponentClassObj)
		{
			return MCPUtils::MakeErrorJson(Result, FString::Printf(
				TEXT("Component class '%s' not found or is not a subclass of UActorComponent. "
					"Common classes: StaticMeshComponent, SkeletalMeshComponent, AudioComponent, "
					"SceneComponent, BoxCollisionComponent, SphereCollisionComponent, CapsuleComponent, "
					"ArrowComponent, ChildActorComponent, SpotLightComponent, PointLightComponent, "
					"WidgetComponent, BillboardComponent"),
				*ComponentClass));
		}

		// If parent component specified, find its SCS node
		USCS_Node* ParentSCSNode = nullptr;
		if (!ParentComponent.IsEmpty())
		{
			for (USCS_Node* Node : ExistingNodes)
			{
				if (Node && Node->GetVariableName().ToString().Equals(ParentComponent, ESearchCase::IgnoreCase))
				{
					ParentSCSNode = Node;
					break;
				}
			}

			if (!ParentSCSNode)
			{
				return MCPUtils::MakeErrorJson(Result, FString::Printf(
					TEXT("Parent component '%s' not found in Blueprint '%s'"),
					*ParentComponent, *Blueprint));
			}
		}

		UE_LOG(LogTemp, Display, TEXT("BlueprintMCP: Adding component '%s' (%s) to Blueprint '%s'"),
			*Component, *ComponentClassObj->GetName(), *Blueprint);

		// Create the SCS node
		USCS_Node* NewNode = SCS->CreateNode(ComponentClassObj, FName(*Component));
		if (!NewNode)
		{
			return MCPUtils::MakeErrorJson(Result, FString::Printf(
				TEXT("Failed to create SCS node for component '%s' with class '%s'"),
				*Component, *ComponentClassObj->GetName()));
		}

		// Add to the hierarchy
		if (ParentSCSNode)
		{
			ParentSCSNode->AddChildNode(NewNode);
		}
		else
		{
			SCS->AddNode(NewNode);
		}

		FBlueprintEditorUtils::MarkBlueprintAsStructurallyModified(BP);
		bool bSaved = MCPUtils::SaveBlueprintPackage(BP);

		UE_LOG(LogTemp, Display, TEXT("BlueprintMCP: Added component '%s' (%s) to '%s' (parent: %s, saved: %s)"),
			*Component, *ComponentClassObj->GetName(), *Blueprint,
			ParentSCSNode ? *ParentComponent : TEXT("(root)"),
			bSaved ? TEXT("true") : TEXT("false"));

		Result->SetStringField(TEXT("component"), NewNode->GetVariableName().ToString());
		Result->SetStringField(TEXT("componentClass"), ComponentClassObj->GetName());
		if (ParentSCSNode)
		{
			Result->SetStringField(TEXT("parentComponent"), ParentSCSNode->GetVariableName().ToString());
		}
		Result->SetBoolField(TEXT("saved"), bSaved);
	}
};

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

UCLASS(meta=(ToolName="remove_blueprint_component"))
class UMCPHandler_RemoveBlueprintComponent : public UObject, public IMCPHandler
{
	GENERATED_BODY()

public:
	UPROPERTY(meta=(Description="Blueprint name or package path"))
	FString Blueprint;

	UPROPERTY(meta=(Description="Component to remove"))
	FString Component;

	virtual FString GetDescription() const override
	{
		return TEXT("Remove a component from a Blueprint's SimpleConstructionScript.");
	}

	virtual void Handle(const FJsonObject* Json, FJsonObject* Result) override
	{
		MCPAssets<UBlueprint> Assets;
		if (!Assets.Exact(Blueprint).Errors(Result).ENone().ETwo().Load()) return;
		UBlueprint* BP = Assets.Object();

		USimpleConstructionScript* SCS = BP->SimpleConstructionScript;
		if (!SCS)
		{
			return MCPUtils::MakeErrorJson(Result, FString::Printf(
				TEXT("Blueprint '%s' does not have a SimpleConstructionScript (not an Actor Blueprint)"),
				*Blueprint));
		}

		// Find the node to remove
		USCS_Node* NodeToRemove = nullptr;
		const TArray<USCS_Node*>& AllNodes = SCS->GetAllNodes();
		for (USCS_Node* Node : AllNodes)
		{
			if (Node && Node->GetVariableName().ToString().Equals(Component, ESearchCase::IgnoreCase))
			{
				NodeToRemove = Node;
				break;
			}
		}

		if (!NodeToRemove)
		{
			// Build list of component names for the error message
			TArray<TSharedPtr<FJsonValue>> CompList;
			for (USCS_Node* Node : AllNodes)
			{
				if (Node)
				{
					CompList.Add(MakeShared<FJsonValueString>(Node->GetVariableName().ToString()));
				}
			}

			MCPUtils::MakeErrorJson(Result, FString::Printf(
				TEXT("Component '%s' not found in Blueprint '%s'"),
				*Component, *Blueprint));
			Result->SetArrayField(TEXT("existingComponents"), CompList);
			return;
		}

		// Prevent removing the root scene component if it has children
		const TArray<USCS_Node*>& RootNodes = SCS->GetRootNodes();
		if (RootNodes.Contains(NodeToRemove) && NodeToRemove->GetChildNodes().Num() > 0)
		{
			return MCPUtils::MakeErrorJson(Result, FString::Printf(
				TEXT("Cannot remove component '%s' because it is a root component with %d child(ren). "
					"Remove or re-parent the children first."),
				*Component, NodeToRemove->GetChildNodes().Num()));
		}

		UE_LOG(LogTemp, Display, TEXT("BlueprintMCP: Removing component '%s' from Blueprint '%s'"),
			*Component, *Blueprint);

		// Remove the node (promotes children to parent if it has any — but we've guarded root above)
		SCS->RemoveNodeAndPromoteChildren(NodeToRemove);

		FBlueprintEditorUtils::MarkBlueprintAsStructurallyModified(BP);
		bool bSaved = MCPUtils::SaveBlueprintPackage(BP);

		UE_LOG(LogTemp, Display, TEXT("BlueprintMCP: Removed component '%s' from '%s' (saved: %s)"),
			*Component, *Blueprint, bSaved ? TEXT("true") : TEXT("false"));

		Result->SetBoolField(TEXT("saved"), bSaved);
	}
};