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

#pragma once

#include "CoreMinimal.h"
#include "MCPHandler.h"
#include "MCPAssetFinder.h"
#include "MCPUtils.h"
#include "Engine/Blueprint.h"
#include "EdGraph/EdGraph.h"
#include "EdGraph/EdGraphNode.h"
#include "Kismet2/KismetEditorUtilities.h"
#include "Dom/JsonValue.h"
#include "Animation/AnimBlueprint.h"
#include "Animation/AnimBlueprintGeneratedClass.h"
#include "Animation/Skeleton.h"
#include "AnimGraphNode_Base.h"
#include "Animation/AnimSequence.h"
#include "Animation/BlendSpace.h"
#include "MCPHandlers_AnimMutation.generated.h"

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

UCLASS(meta=(ToolName="create_anim_blueprint_asset"))
class UMCPHandler_CreateAnimBlueprint : public UObject, public IMCPHandler
{
	GENERATED_BODY()

public:
	UPROPERTY(meta=(Description="Name for the new Animation Blueprint asset"))
	FString Name;

	UPROPERTY(meta=(Description="Package path where the asset will be created (must start with /Game)"))
	FString PackagePath;

	UPROPERTY(meta=(Description="Name or path of the skeleton asset to use"))
	FString Skeleton;

	UPROPERTY(meta=(Optional, Description="Parent class name (default: AnimInstance)"))
	FString ParentClass;

	virtual FString GetDescription() const override
	{
		return TEXT("Create a new Animation Blueprint asset with a specified skeleton.");
	}

	virtual void Handle(const FJsonObject* Json, FJsonObject* Result) override
	{

		if (Name.IsEmpty() || PackagePath.IsEmpty() || Skeleton.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;
		MCPAssets<UBlueprint> ExistCheck;
		if (!ExistCheck.Exact(Name).Errors(Result).EAny().Info()) return;

		// Resolve skeleton
		MCPAssets<USkeleton> SkeletonAssets;
		if (!SkeletonAssets.Exact(Skeleton).Errors(Result).ENone().ETwo().Load()) return;
		USkeleton* SkeletonObj = SkeletonAssets.Object();

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

		UE_LOG(LogTemp, Display, TEXT("BlueprintMCP: Creating AnimBlueprint '%s' in '%s' with skeleton '%s'"),
			*Name, *PackagePath, *SkeletonObj->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(
				ParentClassObj,
				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 = SkeletonObj;

		// Compile
		FKismetEditorUtilities::CompileBlueprint(NewAnimBP);

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


		TArray<TSharedPtr<FJsonValue>> GraphNames = MCPUtils::AllGraphNamesJson(NewAnimBP);

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

		Result->SetStringField(TEXT("assetPath"), FullAssetPath);
		Result->SetStringField(TEXT("targetSkeleton"), SkeletonObj->GetName());
		Result->SetStringField(TEXT("parentClass"), ParentClassObj->GetName());
		Result->SetBoolField(TEXT("saved"), bSaved);
		Result->SetArrayField(TEXT("graphs"), GraphNames);
	}
};

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

UCLASS(meta=(ToolName="list_anim_slot_names"))
class UMCPHandler_ListAnimSlots : public UObject, public IMCPHandler
{
	GENERATED_BODY()

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

	virtual FString GetDescription() const override
	{
		return TEXT("List all animation slot names used in an Animation Blueprint.");
	}

	virtual void Handle(const FJsonObject* Json, FJsonObject* Result) override
	{
		if (Blueprint.IsEmpty())
		{
			return MCPUtils::MakeErrorJson(Result, TEXT("Missing required field: blueprint"));
		}

		MCPAssets<UAnimBlueprint> Assets;
		if (!Assets.Exact(Blueprint).Errors(Result).ENone().ETwo().Load()) return;
		UAnimBlueprint* AnimBP = Assets.Object();

		// Walk all anim nodes to collect slot names
		TSet<FString> SlotNames;
		for (UAnimGraphNode_Base* AnimNode : MCPUtils::AllNodes<UAnimGraphNode_Base>(AnimBP))
		{
			// 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->SetArrayField(TEXT("slots"), SlotsArr);
		Result->SetNumberField(TEXT("count"), SlotsArr.Num());
	}
};

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

UCLASS(meta=(ToolName="list_anim_sync_groups"))
class UMCPHandler_ListSyncGroups : public UObject, public IMCPHandler
{
	GENERATED_BODY()

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

	virtual FString GetDescription() const override
	{
		return TEXT("List all sync group names used in an Animation Blueprint.");
	}

	virtual void Handle(const FJsonObject* Json, FJsonObject* Result) override
	{
		if (Blueprint.IsEmpty())
		{
			return MCPUtils::MakeErrorJson(Result, TEXT("Missing required field: blueprint"));
		}

		MCPAssets<UAnimBlueprint> Assets;
		if (!Assets.Exact(Blueprint).Errors(Result).ENone().ETwo().Load()) return;
		UAnimBlueprint* AnimBP = Assets.Object();

		// Walk all anim nodes to collect sync group names
		TSet<FString> SyncGroupNames;
		for (UAnimGraphNode_Base* AnimNode : MCPUtils::AllNodes<UAnimGraphNode_Base>(AnimBP))
		{
			// 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->SetArrayField(TEXT("syncGroups"), GroupsArr);
		Result->SetNumberField(TEXT("count"), GroupsArr.Num());
	}
};

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

UCLASS(meta=(ToolName="create_blend_space_asset"))
class UMCPHandler_CreateBlendSpace : public UObject, public IMCPHandler
{
	GENERATED_BODY()

public:
	UPROPERTY(meta=(Description="Name for the new Blend Space asset"))
	FString Name;

	UPROPERTY(meta=(Description="Package path where the asset will be created (must start with /Game)"))
	FString PackagePath;

	UPROPERTY(meta=(Description="Name or path of the skeleton asset to use"))
	FString Skeleton;

	virtual FString GetDescription() const override
	{
		return TEXT("Create a new 2D Blend Space asset with a specified skeleton.");
	}

	virtual void Handle(const FJsonObject* Json, FJsonObject* Result) override
	{

		if (Name.IsEmpty() || PackagePath.IsEmpty() || Skeleton.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;
		MCPAssets<UBlendSpace> ExistCheck;
		if (!ExistCheck.Exact(Name).Errors(Result).EAny().Info()) return;

		// Resolve skeleton
		MCPAssets<USkeleton> SkeletonAssets;
		if (!SkeletonAssets.Exact(Skeleton).Errors(Result).ENone().ETwo().Load()) return;
		USkeleton* SkeletonObj = SkeletonAssets.Object();

		UE_LOG(LogTemp, Display, TEXT("BlueprintMCP: Creating Blend Space '%s' in '%s' with skeleton '%s'"),
			*Name, *PackagePath, *SkeletonObj->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(SkeletonObj);

		// 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->SetStringField(TEXT("assetPath"), FullAssetPath);
		Result->SetStringField(TEXT("skeleton"), SkeletonObj->GetName());
		Result->SetBoolField(TEXT("saved"), bSaved);
	}
};

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

USTRUCT()
struct FBlendSpaceSampleEntry
{
	GENERATED_BODY()

	UPROPERTY()
	FString AnimationAsset;

	UPROPERTY()
	float X = 0.0f;

	UPROPERTY()
	float Y = 0.0f;
};

UCLASS(meta=(ToolName="set_blend_space_sample_points"))
class UMCPHandler_SetBlendSpaceSamples : public UObject, public IMCPHandler
{
	GENERATED_BODY()

public:
	UPROPERTY(meta=(Description="Blend Space asset name or package path"))
	FString BlendSpace;

	UPROPERTY(meta=(Optional, Description="Display name for the X axis"))
	FString AxisXName;

	UPROPERTY(meta=(Optional, Description="Display name for the Y axis"))
	FString AxisYName;

	UPROPERTY(meta=(Optional, Description="Minimum value for X axis"))
	float AxisXMin = 0.0f;

	UPROPERTY(meta=(Optional, Description="Maximum value for X axis"))
	float AxisXMax = 0.0f;

	UPROPERTY(meta=(Optional, Description="Minimum value for Y axis"))
	float AxisYMin = 0.0f;

	UPROPERTY(meta=(Optional, Description="Maximum value for Y axis"))
	float AxisYMax = 0.0f;

	UPROPERTY(meta=(Optional, Description="Array of sample points, each with animationAsset, x, y"))
	FMCPJsonArray Samples;

	virtual FString GetDescription() const override
	{
		return TEXT("Set axis parameters and animation sample points on a Blend Space. "
			"Replaces all existing samples.");
	}

	virtual void Handle(const FJsonObject* Json, FJsonObject* Result) override
	{
		if (BlendSpace.IsEmpty())
		{
			return MCPUtils::MakeErrorJson(Result, TEXT("Missing required field: blendSpace"));
		}

		// Load the blend space
		MCPAssets<UBlendSpace> Assets;
		if (!Assets.Exact(BlendSpace).Errors(Result).ENone().ETwo().Load()) return;
		UBlendSpace* BS = Assets.Object();

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

		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 = AxisXMin;
		if (Json->HasField(TEXT("axisXMax"))) MutableParamX.Max = AxisXMax;

		if (!AxisYName.IsEmpty()) MutableParamY.DisplayName = AxisYName;
		if (Json->HasField(TEXT("axisYMin"))) MutableParamY.Min = AxisYMin;
		if (Json->HasField(TEXT("axisYMax"))) MutableParamY.Max = 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
		int32 SamplesSet = 0;

		for (const TSharedPtr<FJsonValue>& SampleVal : Samples.Array)
		{
			FBlendSpaceSampleEntry Entry;
			if (!MCPUtils::PopulateFromJson(FBlendSpaceSampleEntry::StaticStruct(), &Entry, SampleVal, Result)) return;

			UAnimSequence* AnimSeq = nullptr;
			if (!Entry.AnimationAsset.IsEmpty())
			{
				MCPAssets<UAnimSequence> AnimAssets;
				if (AnimAssets.Exact(Entry.AnimationAsset).Load())
					AnimSeq = AnimAssets.Object();
			}

			FVector SampleValue(Entry.X, Entry.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->SetStringField(TEXT("blendSpace"), BS->GetPathName());
		Result->SetNumberField(TEXT("samplesSet"), SamplesSet);
		Result->SetBoolField(TEXT("saved"), bSaved);
	}
};