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

#pragma once

#include "CoreMinimal.h"
#include "MCPHandler.h"
#include "MCPAssetFinder.h"
#include "MCPFetcher.h"
#include "MCPUtils.h"
#include "EdGraph/EdGraph.h"
#include "Kismet2/KismetEditorUtilities.h"
#include "Animation/AnimBlueprint.h"
#include "Animation/AnimSequence.h"
#include "AnimGraphNode_SequencePlayer.h"
#include "AnimStateNode.h"
#include "AnimationStateMachineGraph.h"
#include "UMCPHandler_AddAnimStateToMachine.generated.h"


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

UCLASS(meta=(Group="Unclassified"))
class UMCPHandler_AddAnimStateToMachine : public UObject, public IMCPHandler
{
	GENERATED_BODY()

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

	UPROPERTY(meta=(Description="State machine graph name"))
	FString Graph;

	UPROPERTY(meta=(Description="Name for the new state"))
	FString StateName;

	UPROPERTY(meta=(Optional, Description="X position of the new state node"))
	int32 PosX = 200;

	UPROPERTY(meta=(Optional, Description="Y position of the new state node"))
	int32 PosY = 0;

	UPROPERTY(meta=(Optional, Description="Animation asset name to assign to the state"))
	FString AnimationAsset;

	virtual FString GetDescription() const override
	{
		return TEXT("Add a new state to an animation state machine graph. "
			"Optionally assign an animation asset to the state.");
	}

	virtual void Handle(const FJsonObject* Json, FStringBuilderBase& Result) override
	{
		// Resolve the anim blueprint
		MCPFetcher F(Result);
		UAnimBlueprint* AnimBP = F.Walk(Blueprint).Cast<UAnimBlueprint>();
		if (!AnimBP) return;

		// Find the state machine graph
		UAnimationStateMachineGraph* SMGraph = MCPUtils::FindStateMachineGraph(AnimBP, Graph);
		if (!SMGraph)
		{
			Result.Appendf(TEXT("ERROR: State machine graph '%s' not found in %s\n"), *Graph, *MCPUtils::FormatName(AnimBP));
			return;
		}

		// Check for duplicate state name
		if (MCPUtils::FindStateByName(SMGraph, StateName, nullptr))
		{
			Result.Appendf(TEXT("ERROR: State '%s' already exists in %s\n"), *StateName, *MCPUtils::FormatName(SMGraph));
			return;
		}

		// Create the state node
		UAnimStateNode* NewState = NewObject<UAnimStateNode>(SMGraph);
		NewState->CreateNewGuid();
		NewState->NodePosX = PosX;
		NewState->NodePosY = PosY;

		// Set the state name via the bound graph
		NewState->PostPlacedNewNode();
		NewState->AllocateDefaultPins();

		// Rename the bound graph to set the state name
		if (NewState->GetBoundGraph())
		{
			NewState->GetBoundGraph()->Rename(*StateName, nullptr);
		}

		SMGraph->AddNode(NewState, false, false);
		NewState->SetFlags(RF_Transactional);

		// Optionally set animation asset
		if (!AnimationAsset.IsEmpty() && NewState->GetBoundGraph())
		{
			MCPAssets<UAnimSequence> AnimAssets;
			if (!AnimAssets.Exact(AnimationAsset).Errors(Result).ENone().ETwo().Load()) return;

			UAnimGraphNode_SequencePlayer* SeqNode = NewObject<UAnimGraphNode_SequencePlayer>(NewState->GetBoundGraph());
			SeqNode->CreateNewGuid();
			SeqNode->PostPlacedNewNode();
			SeqNode->AllocateDefaultPins();
			SeqNode->SetAnimationAsset(AnimAssets.Object());
			SeqNode->NodePosX = 0;
			SeqNode->NodePosY = 0;
			NewState->GetBoundGraph()->AddNode(SeqNode, false, false);
		}

		// Compile and save
		FKismetEditorUtilities::CompileBlueprint(AnimBP);
		MCPUtils::SaveBlueprintPackage(AnimBP);

		Result.Appendf(TEXT("Created state '%s' in %s\n"), *StateName, *MCPUtils::FormatName(SMGraph));
		Result.Appendf(TEXT("  node: %s\n"), *MCPUtils::FormatName(NewState));
	}
};