integration/Plugins/BlueprintMCP/Source/BlueprintMCP/Private/MCPHandlers_AnimMutation.cpp

#include "MCPAssetFinder.h"
#include "MCPServer.h"
#include "MCPUtils.h"
#include "Engine/Blueprint.h"
#include "EdGraph/EdGraph.h"
#include "EdGraph/EdGraphNode.h"
#include "EdGraph/EdGraphPin.h"
#include "Kismet2/BlueprintEditorUtils.h"
#include "Kismet2/KismetEditorUtilities.h"
#include "Dom/JsonValue.h"
#include "Serialization/JsonReader.h"
#include "Serialization/JsonWriter.h"
#include "Serialization/JsonSerializer.h"
#include "AssetRegistry/AssetRegistryModule.h"
#include "UObject/SavePackage.h"

// Animation Blueprint includes
#include "Animation/AnimBlueprint.h"
#include "Animation/AnimBlueprintGeneratedClass.h"
#include "Animation/Skeleton.h"
#include "AnimGraphNode_StateMachine.h"
#include "AnimGraphNode_AssetPlayerBase.h"
#include "AnimGraphNode_SequencePlayer.h"
#include "AnimGraphNode_BlendSpacePlayer.h"
#include "AnimGraphNode_Base.h"
#include "Animation/AnimSequence.h"
#include "Animation/BlendSpace.h"

// ============================================================
// HandleCreateAnimBlueprint — create a new Animation Blueprint
// ============================================================

void FBlueprintMCPServer::HandleCreateAnimBlueprint(const FJsonObject* Json, FJsonObject* Result)
{
	FString Name = Json->GetStringField(TEXT("name"));
	FString PackagePath = Json->GetStringField(TEXT("packagePath"));
	FString SkeletonName = Json->GetStringField(TEXT("skeleton"));
	FString ParentClassName = Json->GetStringField(TEXT("parentClass"));

	if (Name.IsEmpty() || PackagePath.IsEmpty() || SkeletonName.IsEmpty())
	{
		return MCPUtils::MakeErrorJson(Result, TEXT("Missing required fields: name, packagePath, skeleton"));
	}

	if (!PackagePath.StartsWith(TEXT("/Game")))
	{
		return MCPUtils::MakeErrorJson(Result, TEXT("packagePath must start with '/Game'"));
	}

	// Check if asset already exists
	FString FullAssetPath = PackagePath / Name;
	if (UMCPAssetFinder::FindAsset(UBlueprint::StaticClass(), Name))
	{
		return MCPUtils::MakeErrorJson(Result, FString::Printf(
			TEXT("Blueprint '%s' already exists. Use a different name or delete the existing asset first."),
			*Name));
	}

	// Resolve skeleton
	USkeleton* Skeleton = UMCPAssetFinder::LoadAsset<USkeleton>(SkeletonName, Result);
	if (!Skeleton) return;

	// Resolve parent class (default: UAnimInstance)
	UClass* ParentClass = UAnimInstance::StaticClass();
	if (!ParentClassName.IsEmpty() && ParentClassName != TEXT("AnimInstance"))
	{
		for (TObjectIterator<UClass> It; It; ++It)
		{
			if (It->GetName() == ParentClassName && It->IsChildOf(UAnimInstance::StaticClass()))
			{
				ParentClass = *It;
				break;
			}
		}
	}

	UE_LOG(LogTemp, Display, TEXT("BlueprintMCP: Creating AnimBlueprint '%s' in '%s' with skeleton '%s'"),
		*Name, *PackagePath, *Skeleton->GetName());

	// Create the package
	FString FullPackagePath = PackagePath / Name;
	UPackage* Package = CreatePackage(*FullPackagePath);
	if (!Package)
	{
		return MCPUtils::MakeErrorJson(Result, FString::Printf(TEXT("Failed to create package at '%s'"), *FullPackagePath));
	}

	// Create the Animation Blueprint
	UAnimBlueprint* NewAnimBP = CastChecked<UAnimBlueprint>(
		FKismetEditorUtilities::CreateBlueprint(
			ParentClass,
			Package,
			FName(*Name),
			BPTYPE_Normal,
			UAnimBlueprint::StaticClass(),
			UAnimBlueprintGeneratedClass::StaticClass()
		));

	if (!NewAnimBP)
	{
		return MCPUtils::MakeErrorJson(Result, TEXT("FKismetEditorUtilities::CreateBlueprint returned null for AnimBlueprint"));
	}

	// Set target skeleton
	NewAnimBP->TargetSkeleton = Skeleton;

	// Compile
	FKismetEditorUtilities::CompileBlueprint(NewAnimBP);

	// Save
	bool bSaved = MCPUtils::SaveBlueprintPackage(NewAnimBP);


	// Collect graph names
	TArray<TSharedPtr<FJsonValue>> GraphNames;
	TArray<UEdGraph*> AllGraphs;
	NewAnimBP->GetAllGraphs(AllGraphs);
	for (UEdGraph* Graph : AllGraphs)
	{
		if (Graph)
		{
			GraphNames.Add(MakeShared<FJsonValueString>(Graph->GetName()));
		}
	}

	UE_LOG(LogTemp, Display, TEXT("BlueprintMCP: Created AnimBlueprint '%s' with %d graphs (saved: %s)"),
		*Name, GraphNames.Num(), bSaved ? TEXT("true") : TEXT("false"));

	Result->SetBoolField(TEXT("success"), true);
	Result->SetStringField(TEXT("blueprintName"), Name);
	Result->SetStringField(TEXT("packagePath"), PackagePath);
	Result->SetStringField(TEXT("assetPath"), FullAssetPath);
	Result->SetStringField(TEXT("targetSkeleton"), Skeleton->GetName());
	Result->SetStringField(TEXT("parentClass"), ParentClass->GetName());
	Result->SetBoolField(TEXT("isAnimBlueprint"), true);
	Result->SetBoolField(TEXT("saved"), bSaved);
	Result->SetArrayField(TEXT("graphs"), GraphNames);
}

// ============================================================
// Tier 3: AnimGraph Blend Tree Mutation
// ============================================================

void FBlueprintMCPServer::HandleAddAnimNode(const FJsonObject* Json, FJsonObject* Result)
{
	FString BlueprintName = Json->GetStringField(TEXT("blueprint"));
	FString GraphName = Json->GetStringField(TEXT("graph"));
	FString NodeType = Json->GetStringField(TEXT("nodeType"));

	if (BlueprintName.IsEmpty() || NodeType.IsEmpty())
	{
		return MCPUtils::MakeErrorJson(Result, TEXT("Missing required fields: blueprint, nodeType"));
	}

	UAnimBlueprint* AnimBP = UMCPAssetFinder::LoadAsset<UAnimBlueprint>(BlueprintName, Result);
	if (!AnimBP) return;

	// Find target graph (default to AnimGraph if not specified)
	UEdGraph* TargetGraph = nullptr;
	if (GraphName.IsEmpty())
	{
		GraphName = TEXT("AnimGraph");
	}

	TArray<UEdGraph*> AllGraphs;
	AnimBP->GetAllGraphs(AllGraphs);
	for (UEdGraph* Graph : AllGraphs)
	{
		if (Graph->GetName() == GraphName)
		{
			TargetGraph = Graph;
			break;
		}
	}

	if (!TargetGraph)
	{
		return MCPUtils::MakeErrorJson(Result, FString::Printf(TEXT("Graph '%s' not found"), *GraphName));
	}

	int32 PosX = Json->HasField(TEXT("posX")) ? (int32)Json->GetNumberField(TEXT("posX")) : 0;
	int32 PosY = Json->HasField(TEXT("posY")) ? (int32)Json->GetNumberField(TEXT("posY")) : 0;

	UAnimGraphNode_Base* NewNode = nullptr;

	if (NodeType == TEXT("SequencePlayer"))
	{
		UAnimGraphNode_SequencePlayer* SeqNode = NewObject<UAnimGraphNode_SequencePlayer>(TargetGraph);
		SeqNode->CreateNewGuid();
		SeqNode->PostPlacedNewNode();
		SeqNode->AllocateDefaultPins();

		// Optionally set animation asset
		FString AnimAssetName = Json->GetStringField(TEXT("animationAsset"));
		if (!AnimAssetName.IsEmpty())
		{
			FAssetData* FoundAnimAsset = UMCPAssetFinder::FindAsset(UAnimSequence::StaticClass(), AnimAssetName);
			if (FoundAnimAsset)
			{
				UAnimSequence* AnimSeq = Cast<UAnimSequence>(FoundAnimAsset->GetAsset());
				if (AnimSeq)
				{
					SeqNode->SetAnimationAsset(AnimSeq);
				}
			}
		}

		NewNode = SeqNode;
	}
	else if (NodeType == TEXT("BlendSpacePlayer"))
	{
		UAnimGraphNode_BlendSpacePlayer* BSNode = NewObject<UAnimGraphNode_BlendSpacePlayer>(TargetGraph);
		BSNode->CreateNewGuid();
		BSNode->PostPlacedNewNode();
		BSNode->AllocateDefaultPins();

		// Optionally set blend space asset
		FString BSAssetName = Json->GetStringField(TEXT("animationAsset"));
		if (!BSAssetName.IsEmpty())
		{
			FAssetData* FoundBSAsset = UMCPAssetFinder::FindAsset(UBlendSpace::StaticClass(), BSAssetName);
			if (FoundBSAsset)
			{
				UBlendSpace* BS = Cast<UBlendSpace>(FoundBSAsset->GetAsset());
				if (BS)
				{
					BSNode->SetAnimationAsset(BS);
				}
			}
		}

		NewNode = BSNode;
	}
	else if (NodeType == TEXT("StateMachine"))
	{
		UAnimGraphNode_StateMachine* SMNode = NewObject<UAnimGraphNode_StateMachine>(TargetGraph);
		SMNode->CreateNewGuid();
		SMNode->PostPlacedNewNode();
		SMNode->AllocateDefaultPins();
		NewNode = SMNode;
	}
	else
	{
		return MCPUtils::MakeErrorJson(Result, FString::Printf(
			TEXT("Unsupported nodeType '%s'. Supported: SequencePlayer, BlendSpacePlayer, StateMachine"),
			*NodeType));
	}

	if (!NewNode)
	{
		return MCPUtils::MakeErrorJson(Result, TEXT("Failed to create anim node"));
	}

	NewNode->NodePosX = PosX;
	NewNode->NodePosY = PosY;
	TargetGraph->AddNode(NewNode, false, false);
	NewNode->SetFlags(RF_Transactional);

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

	Result->SetBoolField(TEXT("success"), true);
	Result->SetStringField(TEXT("nodeType"), NodeType);
	Result->SetStringField(TEXT("nodeId"), NewNode->NodeGuid.ToString());
	Result->SetStringField(TEXT("graph"), GraphName);
	Result->SetBoolField(TEXT("saved"), bSaved);

	// For StateMachine, report the sub-graph name
	if (NodeType == TEXT("StateMachine"))
	{
		if (UAnimGraphNode_StateMachine* SMNode = Cast<UAnimGraphNode_StateMachine>(NewNode))
		{
			if (SMNode->EditorStateMachineGraph)
			{
				Result->SetStringField(TEXT("stateMachineGraph"), SMNode->EditorStateMachineGraph->GetName());
			}
		}
	}
}

// ============================================================
// Tier 4: Advanced Operations
// ============================================================

void FBlueprintMCPServer::HandleAddStateMachine(const FJsonObject* Json, FJsonObject* Result)
{
	FString BlueprintName = Json->GetStringField(TEXT("blueprint"));
	FString MachineName = Json->GetStringField(TEXT("name"));

	if (BlueprintName.IsEmpty())
	{
		return MCPUtils::MakeErrorJson(Result, TEXT("Missing required field: blueprint"));
	}

	// Default name
	if (MachineName.IsEmpty())
	{
		MachineName = TEXT("NewStateMachine");
	}

	// Delegate to HandleAddAnimNode with StateMachine type and AnimGraph as target
	TSharedRef<FJsonObject> ForwardJson = MakeShared<FJsonObject>();
	ForwardJson->SetStringField(TEXT("blueprint"), BlueprintName);
	ForwardJson->SetStringField(TEXT("graph"), TEXT("AnimGraph"));
	ForwardJson->SetStringField(TEXT("nodeType"), TEXT("StateMachine"));
	if (Json->HasField(TEXT("posX")))
		ForwardJson->SetNumberField(TEXT("posX"), Json->GetNumberField(TEXT("posX")));
	if (Json->HasField(TEXT("posY")))
		ForwardJson->SetNumberField(TEXT("posY"), Json->GetNumberField(TEXT("posY")));

	HandleAddAnimNode(&*ForwardJson, Result);
}

void FBlueprintMCPServer::HandleListAnimSlots(const FJsonObject* Json, FJsonObject* Result)
{
	FString BlueprintName = Json->GetStringField(TEXT("blueprint"));
	if (BlueprintName.IsEmpty())
	{
		return MCPUtils::MakeErrorJson(Result, TEXT("Missing required field: blueprint"));
	}

	UAnimBlueprint* AnimBP = UMCPAssetFinder::LoadAsset<UAnimBlueprint>(BlueprintName, Result);
	if (!AnimBP) return;

	// Walk all anim nodes to collect slot names
	TSet<FString> SlotNames;
	TArray<UEdGraph*> AllGraphs;
	AnimBP->GetAllGraphs(AllGraphs);
	for (UEdGraph* Graph : AllGraphs)
	{
		for (UEdGraphNode* Node : Graph->Nodes)
		{
			if (UAnimGraphNode_Base* AnimNode = Cast<UAnimGraphNode_Base>(Node))
			{
				// Check for SlotName property via reflection
				for (TFieldIterator<FNameProperty> PropIt(AnimNode->GetClass()); PropIt; ++PropIt)
				{
					if (PropIt->GetName().Contains(TEXT("SlotName")) || PropIt->GetName().Contains(TEXT("Slot")))
					{
						FName SlotValue = PropIt->GetPropertyValue_InContainer(AnimNode);
						if (!SlotValue.IsNone())
						{
							SlotNames.Add(SlotValue.ToString());
						}
					}
				}
			}
		}
	}

	TArray<TSharedPtr<FJsonValue>> SlotsArr;
	for (const FString& Slot : SlotNames)
	{
		SlotsArr.Add(MakeShared<FJsonValueString>(Slot));
	}

	Result->SetBoolField(TEXT("success"), true);
	Result->SetStringField(TEXT("blueprint"), BlueprintName);
	Result->SetArrayField(TEXT("slots"), SlotsArr);
	Result->SetNumberField(TEXT("count"), SlotsArr.Num());
}

void FBlueprintMCPServer::HandleListSyncGroups(const FJsonObject* Json, FJsonObject* Result)
{
	FString BlueprintName = Json->GetStringField(TEXT("blueprint"));
	if (BlueprintName.IsEmpty())
	{
		return MCPUtils::MakeErrorJson(Result, TEXT("Missing required field: blueprint"));
	}

	UAnimBlueprint* AnimBP = UMCPAssetFinder::LoadAsset<UAnimBlueprint>(BlueprintName, Result);
	if (!AnimBP) return;

	// Walk all anim nodes to collect sync group names
	TSet<FString> SyncGroupNames;
	TArray<UEdGraph*> AllGraphs;
	AnimBP->GetAllGraphs(AllGraphs);
	for (UEdGraph* Graph : AllGraphs)
	{
		for (UEdGraphNode* Node : Graph->Nodes)
		{
			if (UAnimGraphNode_Base* AnimNode = Cast<UAnimGraphNode_Base>(Node))
			{
				// Check for SyncGroup/GroupName property via reflection
				for (TFieldIterator<FNameProperty> PropIt(AnimNode->GetClass()); PropIt; ++PropIt)
				{
					if (PropIt->GetName().Contains(TEXT("SyncGroup")) || PropIt->GetName().Contains(TEXT("GroupName")))
					{
						FName GroupValue = PropIt->GetPropertyValue_InContainer(AnimNode);
						if (!GroupValue.IsNone())
						{
							SyncGroupNames.Add(GroupValue.ToString());
						}
					}
				}
			}
		}
	}

	TArray<TSharedPtr<FJsonValue>> GroupsArr;
	for (const FString& Group : SyncGroupNames)
	{
		GroupsArr.Add(MakeShared<FJsonValueString>(Group));
	}

	Result->SetBoolField(TEXT("success"), true);
	Result->SetStringField(TEXT("blueprint"), BlueprintName);
	Result->SetArrayField(TEXT("syncGroups"), GroupsArr);
	Result->SetNumberField(TEXT("count"), GroupsArr.Num());
}

// ============================================================
// HandleCreateBlendSpace — create a new 2D Blend Space asset
// ============================================================

void FBlueprintMCPServer::HandleCreateBlendSpace(const FJsonObject* Json, FJsonObject* Result)
{
	FString Name = Json->GetStringField(TEXT("name"));
	FString PackagePath = Json->GetStringField(TEXT("packagePath"));
	FString SkeletonName = Json->GetStringField(TEXT("skeleton"));

	if (Name.IsEmpty() || PackagePath.IsEmpty() || SkeletonName.IsEmpty())
	{
		return MCPUtils::MakeErrorJson(Result, TEXT("Missing required fields: name, packagePath, skeleton"));
	}

	if (!PackagePath.StartsWith(TEXT("/Game")))
	{
		return MCPUtils::MakeErrorJson(Result, TEXT("packagePath must start with '/Game'"));
	}

	// Check if asset already exists
	FString FullAssetPath = PackagePath / Name;
	if (UMCPAssetFinder::FindAsset(UBlendSpace::StaticClass(), Name))
	{
		return MCPUtils::MakeErrorJson(Result, FString::Printf(
			TEXT("Blend Space '%s' already exists. Use a different name or delete the existing asset first."),
			*Name));
	}

	// Resolve skeleton
	USkeleton* Skeleton = UMCPAssetFinder::LoadAsset<USkeleton>(SkeletonName, Result);
	if (!Skeleton) return;

	UE_LOG(LogTemp, Display, TEXT("BlueprintMCP: Creating Blend Space '%s' in '%s' with skeleton '%s'"),
		*Name, *PackagePath, *Skeleton->GetName());

	// Create the package
	FString FullPackagePath = PackagePath / Name;
	UPackage* Package = CreatePackage(*FullPackagePath);
	if (!Package)
	{
		return MCPUtils::MakeErrorJson(Result, FString::Printf(TEXT("Failed to create package at '%s'"), *FullPackagePath));
	}

	// Create the Blend Space
	UBlendSpace* NewBS = NewObject<UBlendSpace>(Package, FName(*Name), RF_Public | RF_Standalone);
	if (!NewBS)
	{
		return MCPUtils::MakeErrorJson(Result, TEXT("Failed to create Blend Space object"));
	}

	// Set skeleton
	NewBS->SetSkeleton(Skeleton);

	// Mark dirty and save
	NewBS->MarkPackageDirty();
	bool bSaved = MCPUtils::SaveGenericPackage(NewBS);

	UE_LOG(LogTemp, Display, TEXT("BlueprintMCP: Created Blend Space '%s' (saved: %s)"),
		*Name, bSaved ? TEXT("true") : TEXT("false"));

	Result->SetBoolField(TEXT("success"), true);
	Result->SetStringField(TEXT("assetPath"), FullAssetPath);
	Result->SetStringField(TEXT("skeleton"), Skeleton->GetName());
	Result->SetBoolField(TEXT("saved"), bSaved);
}

// ============================================================
// HandleSetBlendSpaceSamples — add animation samples to a Blend Space
// ============================================================

void FBlueprintMCPServer::HandleSetBlendSpaceSamples(const FJsonObject* Json, FJsonObject* Result)
{
	FString BlendSpaceName = Json->GetStringField(TEXT("blendSpace"));
	if (BlendSpaceName.IsEmpty())
	{
		return MCPUtils::MakeErrorJson(Result, TEXT("Missing required field: blendSpace"));
	}

	// Load the blend space
	UBlendSpace* BS = UMCPAssetFinder::LoadAsset<UBlendSpace>(BlendSpaceName, Result);
	if (!BS) return;

	// Set axis parameters
	BS->PreEditChange(nullptr);

	FString AxisXName = Json->GetStringField(TEXT("axisXName"));
	FString AxisYName = Json->GetStringField(TEXT("axisYName"));

	const FBlendParameter& ParamX = BS->GetBlendParameter(0);
	const FBlendParameter& ParamY = BS->GetBlendParameter(1);

	// We need to modify BlendParameters directly — use const_cast since there's no setter API
	FBlendParameter& MutableParamX = const_cast<FBlendParameter&>(ParamX);
	FBlendParameter& MutableParamY = const_cast<FBlendParameter&>(ParamY);

	if (!AxisXName.IsEmpty()) MutableParamX.DisplayName = AxisXName;
	if (Json->HasField(TEXT("axisXMin"))) MutableParamX.Min = (float)Json->GetNumberField(TEXT("axisXMin"));
	if (Json->HasField(TEXT("axisXMax"))) MutableParamX.Max = (float)Json->GetNumberField(TEXT("axisXMax"));

	if (!AxisYName.IsEmpty()) MutableParamY.DisplayName = AxisYName;
	if (Json->HasField(TEXT("axisYMin"))) MutableParamY.Min = (float)Json->GetNumberField(TEXT("axisYMin"));
	if (Json->HasField(TEXT("axisYMax"))) MutableParamY.Max = (float)Json->GetNumberField(TEXT("axisYMax"));

	// Clear existing samples (delete from end to start)
	int32 NumExisting = BS->GetNumberOfBlendSamples();
	for (int32 i = NumExisting - 1; i >= 0; --i)
	{
		BS->DeleteSample(i);
	}

	// Add new samples
	const TArray<TSharedPtr<FJsonValue>>* SamplesArray = nullptr;
	int32 SamplesSet = 0;

	if (Json->TryGetArrayField(TEXT("samples"), SamplesArray) && SamplesArray)
	{
		for (const TSharedPtr<FJsonValue>& SampleVal : *SamplesArray)
		{
			const TSharedPtr<FJsonObject>& SampleObj = SampleVal->AsObject();
			if (!SampleObj.IsValid()) continue;

			FString AnimAssetName = SampleObj->GetStringField(TEXT("animationAsset"));
			float X = (float)SampleObj->GetNumberField(TEXT("x"));
			float Y = (float)SampleObj->GetNumberField(TEXT("y"));

			UAnimSequence* AnimSeq = nullptr;
			if (!AnimAssetName.IsEmpty())
			{
				FAssetData* FoundAnimAsset = UMCPAssetFinder::FindAsset(UAnimSequence::StaticClass(), AnimAssetName);
				if (FoundAnimAsset)
				{
					AnimSeq = Cast<UAnimSequence>(FoundAnimAsset->GetAsset());
				}
			}

			FVector SampleValue(X, Y, 0.0f);
			if (AnimSeq)
			{
				BS->AddSample(AnimSeq, SampleValue);
			}
			else
			{
				BS->AddSample(SampleValue);
			}
			SamplesSet++;
		}
	}

	BS->ValidateSampleData();
	BS->PostEditChange();

	// Save
	BS->MarkPackageDirty();
	bool bSaved = MCPUtils::SaveGenericPackage(BS);

	UE_LOG(LogTemp, Display, TEXT("BlueprintMCP: Set %d samples on Blend Space '%s' (saved: %s)"),
		SamplesSet, *BS->GetName(), bSaved ? TEXT("true") : TEXT("false"));

	Result->SetBoolField(TEXT("success"), true);
	Result->SetStringField(TEXT("blendSpace"), BS->GetPathName());
	Result->SetNumberField(TEXT("samplesSet"), SamplesSet);
	Result->SetBoolField(TEXT("saved"), bSaved);
}