#pragma once

#include "CoreMinimal.h"
#include "WingServer.h"
#include "WingHandler.h"
#include "WingFetcher.h"
#include "WingUtils.h"
#include "EdGraph/EdGraph.h"
#include "EdGraph/EdGraphNode.h"
#include "EdGraph/EdGraphPin.h"
#include "Kismet2/KismetEditorUtilities.h"
#include "Animation/AnimBlueprint.h"
#include "Animation/BlendSpace.h"
#include "AnimGraphNode_BlendSpacePlayer.h"
#include "EdGraphSchema_K2.h"
#include "AnimStateNode.h"
#include "AnimationStateMachineGraph.h"
#include "K2Node_VariableGet.h"
#include "StateMachine_SetBlendSpace.generated.h"


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

UCLASS()
class UWing_StateMachine_SetBlendSpace : public UObject, public IWingHandler
{
	GENERATED_BODY()

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

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

	UPROPERTY(meta=(Description="Name of the state to modify"))
	FString StateName;

	UPROPERTY(meta=(Description="Blend Space asset package path"))
	FString BlendSpace;

	UPROPERTY(meta=(Optional, Description="Blueprint variable name to wire to the X axis input"))
	FString XVariable;

	UPROPERTY(meta=(Optional, Description="Blueprint variable name to wire to the Y axis input"))
	FString YVariable;

	virtual FString GetDescription() const override
	{
		return TEXT("Place a BlendSpacePlayer in a state's inner graph, connect it to the output pose, "
			"and optionally wire blueprint variables to the X and Y axis inputs.");
	}

	virtual void Handle() override
	{
		// Load the anim blueprint
		WingFetcher F;
		UAnimBlueprint* AnimBP = F.Asset(Blueprint).Cast<UAnimBlueprint>();
		if (!AnimBP) return;

		// Find the state machine graph and state
		UAnimationStateMachineGraph* SMGraph = WingUtils::FindStateMachineGraph(AnimBP, Graph);
		if (!SMGraph) { UWingServer::Printf(TEXT("ERROR: State machine graph '%s' not found\n"), *Graph); return; }

		UAnimStateNode* StateNode = WingUtils::FindStateByName(SMGraph, StateName);
		if (!StateNode) return;

		UEdGraph* InnerGraph = StateNode->GetBoundGraph();
		if (!InnerGraph) { UWingServer::Printf(TEXT("ERROR: State '%s' has no bound graph\n"), *StateName); return; }

		// Load the blend space asset
		WingFetcher F2;
		UBlendSpace* BlendSpaceAsset = F2.Asset(BlendSpace).Cast<UBlendSpace>();
		if (!BlendSpaceAsset) return;

		// Find existing BlendSpacePlayer or create one
		UAnimGraphNode_BlendSpacePlayer* BSNode = nullptr;
		for (UEdGraphNode* Node : InnerGraph->Nodes)
		{
			BSNode = Cast<UAnimGraphNode_BlendSpacePlayer>(Node);
			if (BSNode) break;
		}

		if (!BSNode)
		{
			BSNode = NewObject<UAnimGraphNode_BlendSpacePlayer>(InnerGraph);
			BSNode->CreateNewGuid();
			BSNode->PostPlacedNewNode();
			BSNode->AllocateDefaultPins();
			BSNode->NodePosX = 0;
			BSNode->NodePosY = 0;
			InnerGraph->AddNode(BSNode, false, false);
		}

		BSNode->SetAnimationAsset(BlendSpaceAsset);

		// Connect BlendSpacePlayer output to the Output Animation Pose node
		ConnectToOutputPose(BSNode, InnerGraph);

		// Wire X and Y variables if provided
		WireVariable(AnimBP, InnerGraph, BSNode, XVariable, TEXT("X"));
		WireVariable(AnimBP, InnerGraph, BSNode, YVariable, TEXT("Y"));

		// Compile
		FKismetEditorUtilities::CompileBlueprint(AnimBP);

		UWingServer::Printf(TEXT("BlendSpacePlayer %s placed in state %s\n"),
			*WingUtils::FormatName(BSNode), *StateName);
	}

private:
	void ConnectToOutputPose(UAnimGraphNode_BlendSpacePlayer* BSNode, UEdGraph* InnerGraph)
	{
		// Find the result node (AnimGraphNode_Root or AnimGraphNode_StateResult)
		UEdGraphNode* ResultNode = nullptr;
		for (UEdGraphNode* Node : InnerGraph->Nodes)
		{
			if (Node->GetClass()->GetName().Contains(TEXT("AnimGraphNode_Root")) ||
				Node->GetClass()->GetName().Contains(TEXT("AnimGraphNode_StateResult")))
			{
				ResultNode = Node;
				break;
			}
		}
		if (!ResultNode) return;

		// Find the pose output pin on BlendSpacePlayer and input pin on result node
		UEdGraphPin* BSOutputPin = nullptr;
		for (UEdGraphPin* Pin : BSNode->Pins)
		{
			if (Pin && Pin->Direction == EGPD_Output && Pin->PinType.PinCategory == UEdGraphSchema_K2::PC_Struct)
			{
				BSOutputPin = Pin;
				break;
			}
		}

		UEdGraphPin* ResultInputPin = nullptr;
		for (UEdGraphPin* Pin : ResultNode->Pins)
		{
			if (Pin && Pin->Direction == EGPD_Input && Pin->PinType.PinCategory == UEdGraphSchema_K2::PC_Struct)
			{
				ResultInputPin = Pin;
				break;
			}
		}

		if (!BSOutputPin || !ResultInputPin) return;

		ResultInputPin->BreakAllPinLinks();
		const UEdGraphSchema* Schema = InnerGraph->GetSchema();
		if (Schema)
			Schema->TryCreateConnection(BSOutputPin, ResultInputPin);
	}

	void WireVariable(UAnimBlueprint* AnimBP, UEdGraph* InnerGraph,
		UAnimGraphNode_BlendSpacePlayer* BSNode, const FString& VarName,
		const TCHAR* PinName)
	{
		if (VarName.IsEmpty()) return;

		// Verify the variable exists in the blueprint
		FName VarFName(*VarName);
		bool bVarFound = false;
		for (FBPVariableDescription& Var : AnimBP->NewVariables)
		{
			if (Var.VarName == VarFName)
			{
				bVarFound = true;
				break;
			}
		}
		if (!bVarFound)
		{
			if (UClass* GenClass = AnimBP->SkeletonGeneratedClass)
			{
				if (GenClass->FindPropertyByName(VarFName))
					bVarFound = true;
			}
		}
		if (!bVarFound)
		{
			UWingServer::Printf(TEXT("WARNING: Variable '%s' not found, skipping %s wire\n"), *VarName, PinName);
			return;
		}

		// Create a VariableGet node
		UK2Node_VariableGet* GetNode = NewObject<UK2Node_VariableGet>(InnerGraph);
		GetNode->VariableReference.SetSelfMember(VarFName);
		GetNode->NodePosX = BSNode->NodePosX - 250;
		GetNode->NodePosY = BSNode->NodePosY;
		InnerGraph->AddNode(GetNode, false, false);
		GetNode->AllocateDefaultPins();

		// Find the variable output pin
		UEdGraphPin* VarOutPin = nullptr;
		for (UEdGraphPin* Pin : GetNode->Pins)
		{
			if (Pin && Pin->Direction == EGPD_Output && Pin->PinName == VarFName)
			{
				VarOutPin = Pin;
				break;
			}
		}

		UEdGraphPin* TargetPin = BSNode->FindPin(FName(PinName));

		if (VarOutPin && TargetPin)
		{
			const UEdGraphSchema* Schema = InnerGraph->GetSchema();
			if (Schema)
				Schema->TryCreateConnection(VarOutPin, TargetPin);
		}
	}
};