#include "MCPUtils.h"
#include "Dom/JsonValue.h"
#include "Engine/Blueprint.h"
#include "EdGraph/EdGraph.h"
#include "EdGraph/EdGraphNode.h"
#include "EdGraph/EdGraphPin.h"
#include "K2Node_CallFunction.h"
#include "K2Node_Event.h"
#include "K2Node_CustomEvent.h"
#include "K2Node_FunctionEntry.h"
#include "K2Node_EditablePinBase.h"
#include "K2Node_VariableGet.h"
#include "K2Node_VariableSet.h"
#include "K2Node_MacroInstance.h"
#include "K2Node_DynamicCast.h"
#include "K2Node_CallParentFunction.h"
#include "K2Node_IfThenElse.h"

// Animation Blueprint support
#include "Animation/AnimBlueprint.h"
#include "AnimGraphNode_StateMachine.h"
#include "AnimGraphNode_AssetPlayerBase.h"
#include "AnimGraphNode_SequencePlayer.h"
#include "AnimGraphNode_BlendSpacePlayer.h"
#include "AnimGraphNode_Base.h"
#include "AnimStateNode.h"
#include "AnimStateTransitionNode.h"
#include "AnimStateConduitNode.h"
#include "AnimStateEntryNode.h"
#include "AnimationStateMachineGraph.h"

// Material support
#include "Materials/MaterialExpression.h"
#include "Materials/MaterialExpressionScalarParameter.h"
#include "Materials/MaterialExpressionVectorParameter.h"
#include "Materials/MaterialExpressionTextureSampleParameter2D.h"
#include "Materials/MaterialExpressionStaticSwitchParameter.h"
#include "Materials/MaterialExpressionConstant.h"
#include "Materials/MaterialExpressionConstant3Vector.h"
#include "Materials/MaterialExpressionConstant4Vector.h"
#include "Materials/MaterialExpressionTextureSample.h"
#include "Materials/MaterialExpressionTextureCoordinate.h"
#include "Materials/MaterialExpressionComponentMask.h"
#include "Materials/MaterialExpressionCustom.h"
#include "Materials/MaterialExpressionFunctionInput.h"
#include "Materials/MaterialExpressionFunctionOutput.h"
#include "Materials/MaterialExpressionMaterialFunctionCall.h"
#include "MaterialGraph/MaterialGraphNode.h"

TSharedRef<FJsonObject> MCPUtils::SerializeBlueprint(UBlueprint* BP)
{
	TSharedRef<FJsonObject> J = MakeShared<FJsonObject>();
	J->SetStringField(TEXT("name"), BP->GetName());
	J->SetStringField(TEXT("path"), BP->GetPackage()->GetName());
	J->SetStringField(TEXT("parentClass"), BP->ParentClass ? BP->ParentClass->GetName() : TEXT("None"));
	J->SetStringField(TEXT("blueprintType"),
		StaticEnum<EBlueprintType>()->GetNameStringByValue((int64)BP->BlueprintType));

	// Animation Blueprint detection
	if (UAnimBlueprint* AnimBP = Cast<UAnimBlueprint>(BP))
	{
		J->SetBoolField(TEXT("isAnimBlueprint"), true);
		if (AnimBP->TargetSkeleton)
		{
			J->SetStringField(TEXT("targetSkeleton"), AnimBP->TargetSkeleton->GetName());
			J->SetStringField(TEXT("targetSkeletonPath"), AnimBP->TargetSkeleton->GetPathName());
		}
	}

	// Variables
	TArray<TSharedPtr<FJsonValue>> Vars;
	for (const FBPVariableDescription& V : BP->NewVariables)
	{
		TSharedRef<FJsonObject> VJ = MakeShared<FJsonObject>();
		VJ->SetStringField(TEXT("name"), V.VarName.ToString());
		VJ->SetStringField(TEXT("type"), V.VarType.PinCategory.ToString());
		if (V.VarType.PinSubCategoryObject.IsValid())
			VJ->SetStringField(TEXT("subtype"), V.VarType.PinSubCategoryObject->GetName());
		VJ->SetBoolField(TEXT("isArray"), V.VarType.IsArray());
		VJ->SetBoolField(TEXT("isSet"), V.VarType.IsSet());
		VJ->SetBoolField(TEXT("isMap"), V.VarType.IsMap());
		VJ->SetStringField(TEXT("category"), V.Category.ToString());
		VJ->SetStringField(TEXT("defaultValue"), V.DefaultValue);
		Vars.Add(MakeShared<FJsonValueObject>(VJ));
	}
	J->SetArrayField(TEXT("variables"), Vars);

	// Interfaces
	TArray<TSharedPtr<FJsonValue>> Ifaces;
	for (const FBPInterfaceDescription& I : BP->ImplementedInterfaces)
	{
		if (I.Interface)
			Ifaces.Add(MakeShared<FJsonValueString>(I.Interface->GetName()));
	}
	J->SetArrayField(TEXT("interfaces"), Ifaces);

	return J;
}

TSharedPtr<FJsonObject> MCPUtils::SerializeNode(UEdGraphNode* Node)
{
	TSharedRef<FJsonObject> NJ = MakeShared<FJsonObject>();
	NJ->SetStringField(TEXT("id"), Node->NodeGuid.ToString());
	NJ->SetStringField(TEXT("class"), Node->GetClass()->GetName());
	NJ->SetStringField(TEXT("title"), Node->GetNodeTitle(ENodeTitleType::FullTitle).ToString());
	if (!Node->NodeComment.IsEmpty())
		NJ->SetStringField(TEXT("comment"), Node->NodeComment);
	NJ->SetNumberField(TEXT("posX"), Node->NodePosX);
	NJ->SetNumberField(TEXT("posY"), Node->NodePosY);

	// Material graph node — extract UMaterialExpression data
	if (UMaterialGraphNode* MatNode = Cast<UMaterialGraphNode>(Node))
	{
		NJ->SetStringField(TEXT("nodeType"), TEXT("MaterialExpression"));
		if (MatNode->MaterialExpression)
		{
			TSharedPtr<FJsonObject> ExprJson = SerializeMaterialExpression(MatNode->MaterialExpression);
			if (ExprJson.IsValid())
			{
				NJ->SetObjectField(TEXT("expression"), ExprJson);
			}
		}
	}
	// Animation Blueprint node types
	else if (auto* SMNode = Cast<UAnimGraphNode_StateMachine>(Node))
	{
		NJ->SetStringField(TEXT("nodeType"), TEXT("AnimStateMachine"));
		if (SMNode->EditorStateMachineGraph)
		{
			NJ->SetStringField(TEXT("stateMachineName"), SMNode->EditorStateMachineGraph->GetName());
			int32 StateCount = 0, TransitionCount = 0;
			for (UEdGraphNode* SubNode : SMNode->EditorStateMachineGraph->Nodes)
			{
				if (Cast<UAnimStateNode>(SubNode)) StateCount++;
				else if (Cast<UAnimStateTransitionNode>(SubNode)) TransitionCount++;
			}
			NJ->SetNumberField(TEXT("stateCount"), StateCount);
			NJ->SetNumberField(TEXT("transitionCount"), TransitionCount);
		}
	}
	else if (auto* SeqPlayer = Cast<UAnimGraphNode_SequencePlayer>(Node))
	{
		NJ->SetStringField(TEXT("nodeType"), TEXT("AnimSequencePlayer"));
		if (UAnimationAsset* Asset = SeqPlayer->GetAnimationAsset())
		{
			NJ->SetStringField(TEXT("animationAsset"), Asset->GetName());
			NJ->SetStringField(TEXT("animationAssetPath"), Asset->GetPathName());
		}
	}
	else if (auto* BSPlayer = Cast<UAnimGraphNode_BlendSpacePlayer>(Node))
	{
		NJ->SetStringField(TEXT("nodeType"), TEXT("AnimBlendSpacePlayer"));
		if (UAnimationAsset* Asset = BSPlayer->GetAnimationAsset())
		{
			NJ->SetStringField(TEXT("blendSpaceAsset"), Asset->GetName());
			NJ->SetStringField(TEXT("blendSpaceAssetPath"), Asset->GetPathName());
		}
	}
	else if (auto* AssetPlayer = Cast<UAnimGraphNode_AssetPlayerBase>(Node))
	{
		NJ->SetStringField(TEXT("nodeType"), TEXT("AnimAssetPlayer"));
		if (UAnimationAsset* Asset = AssetPlayer->GetAnimationAsset())
		{
			NJ->SetStringField(TEXT("animationAsset"), Asset->GetName());
			NJ->SetStringField(TEXT("animationAssetPath"), Asset->GetPathName());
		}
	}
	else if (Cast<UAnimGraphNode_Base>(Node))
	{
		NJ->SetStringField(TEXT("nodeType"), TEXT("AnimNode"));
	}
	else if (auto* StateNode = Cast<UAnimStateNode>(Node))
	{
		NJ->SetStringField(TEXT("nodeType"), TEXT("AnimState"));
		NJ->SetStringField(TEXT("stateName"), StateNode->GetStateName());
		NJ->SetBoolField(TEXT("bAlwaysResetOnEntry"), StateNode->bAlwaysResetOnEntry);
	}
	else if (auto* TransNode = Cast<UAnimStateTransitionNode>(Node))
	{
		NJ->SetStringField(TEXT("nodeType"), TEXT("AnimTransition"));
		if (UAnimStateNode* FromState = Cast<UAnimStateNode>(TransNode->GetPreviousState()))
		{
			NJ->SetStringField(TEXT("fromState"), FromState->GetStateName());
		}
		if (UAnimStateNode* ToState = Cast<UAnimStateNode>(TransNode->GetNextState()))
		{
			NJ->SetStringField(TEXT("toState"), ToState->GetStateName());
		}
		NJ->SetNumberField(TEXT("crossfadeDuration"), TransNode->CrossfadeDuration);
		NJ->SetNumberField(TEXT("blendMode"), (int32)TransNode->BlendMode);
		NJ->SetNumberField(TEXT("priorityOrder"), TransNode->PriorityOrder);
		NJ->SetNumberField(TEXT("logicType"), (int32)TransNode->LogicType.GetValue());
		NJ->SetBoolField(TEXT("bBidirectional"), TransNode->Bidirectional);
	}
	else if (Cast<UAnimStateConduitNode>(Node))
	{
		NJ->SetStringField(TEXT("nodeType"), TEXT("AnimConduit"));
	}
	else if (Cast<UAnimStateEntryNode>(Node))
	{
		NJ->SetStringField(TEXT("nodeType"), TEXT("AnimStateEntry"));
	}
	// K2Node specifics — check CallParentFunction before CallFunction (inheritance)
	else if (auto* CPF = Cast<UK2Node_CallParentFunction>(Node))
	{
		NJ->SetStringField(TEXT("functionName"), CPF->FunctionReference.GetMemberName().ToString());
		if (CPF->FunctionReference.GetMemberParentClass())
			NJ->SetStringField(TEXT("targetClass"), CPF->FunctionReference.GetMemberParentClass()->GetName());
		NJ->SetStringField(TEXT("nodeType"), TEXT("CallParentFunction"));
	}
	else if (auto* CF = Cast<UK2Node_CallFunction>(Node))
	{
		NJ->SetStringField(TEXT("functionName"), CF->FunctionReference.GetMemberName().ToString());
		if (CF->FunctionReference.GetMemberParentClass())
			NJ->SetStringField(TEXT("targetClass"), CF->FunctionReference.GetMemberParentClass()->GetName());
	}
	else if (auto* FE = Cast<UK2Node_FunctionEntry>(Node))
	{
		NJ->SetStringField(TEXT("nodeType"), TEXT("FunctionEntry"));

		// Serialize UserDefinedPins (parameter names and types)
		TArray<TSharedPtr<FJsonValue>> ParamArr;
		for (const TSharedPtr<FUserPinInfo>& PinInfo : FE->UserDefinedPins)
		{
			if (!PinInfo.IsValid()) continue;
			TSharedRef<FJsonObject> ParamJ = MakeShared<FJsonObject>();
			ParamJ->SetStringField(TEXT("name"), PinInfo->PinName.ToString());
			FString ParamType = PinInfo->PinType.PinCategory.ToString();
			ParamJ->SetStringField(TEXT("type"), ParamType);
			if (PinInfo->PinType.PinSubCategoryObject.IsValid())
				ParamJ->SetStringField(TEXT("subtype"), PinInfo->PinType.PinSubCategoryObject->GetName());
			else if (ParamType == TEXT("None") || ParamType.IsEmpty())
				ParamJ->SetBoolField(TEXT("typeUnknown"), true);
			ParamArr.Add(MakeShared<FJsonValueObject>(ParamJ));
		}
		NJ->SetArrayField(TEXT("parameters"), ParamArr);
	}
	else if (auto* Ev = Cast<UK2Node_Event>(Node))
	{
		NJ->SetStringField(TEXT("eventName"), Ev->EventReference.GetMemberName().ToString());
		NJ->SetStringField(TEXT("nodeType"), Ev->bOverrideFunction ? TEXT("OverrideEvent") : TEXT("Event"));
	}
	else if (auto* CE = Cast<UK2Node_CustomEvent>(Node))
	{
		NJ->SetStringField(TEXT("eventName"), CE->CustomFunctionName.ToString());
		NJ->SetStringField(TEXT("nodeType"), TEXT("CustomEvent"));

		// Serialize UserDefinedPins (parameter names and types)
		TArray<TSharedPtr<FJsonValue>> ParamArr;
		for (const TSharedPtr<FUserPinInfo>& PinInfo : CE->UserDefinedPins)
		{
			if (!PinInfo.IsValid()) continue;
			TSharedRef<FJsonObject> ParamJ = MakeShared<FJsonObject>();
			ParamJ->SetStringField(TEXT("name"), PinInfo->PinName.ToString());
			FString ParamType = PinInfo->PinType.PinCategory.ToString();
			ParamJ->SetStringField(TEXT("type"), ParamType);
			if (PinInfo->PinType.PinSubCategoryObject.IsValid())
				ParamJ->SetStringField(TEXT("subtype"), PinInfo->PinType.PinSubCategoryObject->GetName());
			else if (ParamType == TEXT("None") || ParamType.IsEmpty())
				ParamJ->SetBoolField(TEXT("typeUnknown"), true);
			ParamArr.Add(MakeShared<FJsonValueObject>(ParamJ));
		}
		NJ->SetArrayField(TEXT("parameters"), ParamArr);
	}
	else if (auto* VG = Cast<UK2Node_VariableGet>(Node))
	{
		NJ->SetStringField(TEXT("variableName"), VG->GetVarName().ToString());
		NJ->SetStringField(TEXT("nodeType"), TEXT("VariableGet"));
	}
	else if (auto* VS = Cast<UK2Node_VariableSet>(Node))
	{
		NJ->SetStringField(TEXT("variableName"), VS->GetVarName().ToString());
		NJ->SetStringField(TEXT("nodeType"), TEXT("VariableSet"));
	}
	else if (auto* MI = Cast<UK2Node_MacroInstance>(Node))
	{
		if (MI->GetMacroGraph())
			NJ->SetStringField(TEXT("macroName"), MI->GetMacroGraph()->GetName());
		NJ->SetStringField(TEXT("nodeType"), TEXT("MacroInstance"));
	}
	else if (auto* DC = Cast<UK2Node_DynamicCast>(Node))
	{
		if (DC->TargetType)
			NJ->SetStringField(TEXT("castTarget"), DC->TargetType->GetName());
		NJ->SetStringField(TEXT("nodeType"), TEXT("DynamicCast"));
	}
	else if (Cast<UK2Node_IfThenElse>(Node))
	{
		NJ->SetStringField(TEXT("nodeType"), TEXT("Branch"));
	}

	// Pins
	TArray<TSharedPtr<FJsonValue>> Pins;
	for (UEdGraphPin* Pin : Node->Pins)
	{
		if (!Pin || Pin->bHidden) continue;
		TSharedPtr<FJsonObject> PJ = SerializePin(Pin);
		if (PJ.IsValid())
			Pins.Add(MakeShared<FJsonValueObject>(PJ.ToSharedRef()));
	}
	NJ->SetArrayField(TEXT("pins"), Pins);
	return NJ;
}

TSharedPtr<FJsonObject> MCPUtils::SerializePin(UEdGraphPin* Pin)
{
	TSharedRef<FJsonObject> PJ = MakeShared<FJsonObject>();
	PJ->SetStringField(TEXT("name"), Pin->PinName.ToString());
	PJ->SetStringField(TEXT("direction"), Pin->Direction == EGPD_Input ? TEXT("Input") : TEXT("Output"));
	PJ->SetStringField(TEXT("type"), Pin->PinType.PinCategory.ToString());
	if (Pin->PinType.PinSubCategoryObject.IsValid())
		PJ->SetStringField(TEXT("subtype"), Pin->PinType.PinSubCategoryObject->GetName());
	if (!Pin->DefaultValue.IsEmpty())
		PJ->SetStringField(TEXT("defaultValue"), Pin->DefaultValue);

	if (Pin->LinkedTo.Num() > 0)
	{
		TArray<TSharedPtr<FJsonValue>> Conns;
		for (UEdGraphPin* Linked : Pin->LinkedTo)
		{
			if (!Linked || !Linked->GetOwningNode()) continue;
			TSharedRef<FJsonObject> CJ = MakeShared<FJsonObject>();
			CJ->SetStringField(TEXT("nodeId"), Linked->GetOwningNode()->NodeGuid.ToString());
			CJ->SetStringField(TEXT("pinName"), Linked->PinName.ToString());
			Conns.Add(MakeShared<FJsonValueObject>(CJ));
		}
		PJ->SetArrayField(TEXT("connections"), Conns);
	}
	return PJ;
}

TSharedPtr<FJsonObject> MCPUtils::SerializeMaterialExpression(UMaterialExpression* Expression)
{
	if (!Expression) return nullptr;

	TSharedRef<FJsonObject> EJ = MakeShared<FJsonObject>();
	EJ->SetStringField(TEXT("class"), Expression->GetClass()->GetName());
	EJ->SetStringField(TEXT("name"), Expression->GetName());
	EJ->SetStringField(TEXT("description"), Expression->GetDescription());
	EJ->SetNumberField(TEXT("posX"), Expression->MaterialExpressionEditorX);
	EJ->SetNumberField(TEXT("posY"), Expression->MaterialExpressionEditorY);

	if (auto* SP = Cast<UMaterialExpressionScalarParameter>(Expression))
	{
		EJ->SetStringField(TEXT("expressionType"), TEXT("ScalarParameter"));
		EJ->SetStringField(TEXT("parameterName"), SP->ParameterName.ToString());
		EJ->SetNumberField(TEXT("defaultValue"), SP->DefaultValue);
		EJ->SetStringField(TEXT("group"), SP->Group.ToString());
	}
	else if (auto* VP = Cast<UMaterialExpressionVectorParameter>(Expression))
	{
		EJ->SetStringField(TEXT("expressionType"), TEXT("VectorParameter"));
		EJ->SetStringField(TEXT("parameterName"), VP->ParameterName.ToString());
		TSharedRef<FJsonObject> DefVal = MakeShared<FJsonObject>();
		DefVal->SetNumberField(TEXT("r"), VP->DefaultValue.R);
		DefVal->SetNumberField(TEXT("g"), VP->DefaultValue.G);
		DefVal->SetNumberField(TEXT("b"), VP->DefaultValue.B);
		DefVal->SetNumberField(TEXT("a"), VP->DefaultValue.A);
		EJ->SetObjectField(TEXT("defaultValue"), DefVal);
		EJ->SetStringField(TEXT("group"), VP->Group.ToString());
	}
	else if (auto* TP = Cast<UMaterialExpressionTextureSampleParameter2D>(Expression))
	{
		EJ->SetStringField(TEXT("expressionType"), TEXT("TextureSampleParameter2D"));
		EJ->SetStringField(TEXT("parameterName"), TP->ParameterName.ToString());
		if (TP->Texture)
			EJ->SetStringField(TEXT("texture"), TP->Texture->GetPathName());
		EJ->SetStringField(TEXT("group"), TP->Group.ToString());
	}
	else if (auto* SSP = Cast<UMaterialExpressionStaticSwitchParameter>(Expression))
	{
		EJ->SetStringField(TEXT("expressionType"), TEXT("StaticSwitchParameter"));
		EJ->SetStringField(TEXT("parameterName"), SSP->ParameterName.ToString());
		EJ->SetBoolField(TEXT("defaultValue"), SSP->DefaultValue);
		EJ->SetStringField(TEXT("group"), SSP->Group.ToString());
	}
	else if (auto* SC = Cast<UMaterialExpressionConstant>(Expression))
	{
		EJ->SetStringField(TEXT("expressionType"), TEXT("Constant"));
		EJ->SetNumberField(TEXT("value"), SC->R);
	}
	else if (auto* C3 = Cast<UMaterialExpressionConstant3Vector>(Expression))
	{
		EJ->SetStringField(TEXT("expressionType"), TEXT("Constant3Vector"));
		TSharedRef<FJsonObject> Val = MakeShared<FJsonObject>();
		Val->SetNumberField(TEXT("r"), C3->Constant.R);
		Val->SetNumberField(TEXT("g"), C3->Constant.G);
		Val->SetNumberField(TEXT("b"), C3->Constant.B);
		EJ->SetObjectField(TEXT("value"), Val);
	}
	else if (auto* C4 = Cast<UMaterialExpressionConstant4Vector>(Expression))
	{
		EJ->SetStringField(TEXT("expressionType"), TEXT("Constant4Vector"));
		TSharedRef<FJsonObject> Val = MakeShared<FJsonObject>();
		Val->SetNumberField(TEXT("r"), C4->Constant.R);
		Val->SetNumberField(TEXT("g"), C4->Constant.G);
		Val->SetNumberField(TEXT("b"), C4->Constant.B);
		Val->SetNumberField(TEXT("a"), C4->Constant.A);
		EJ->SetObjectField(TEXT("value"), Val);
	}
	else if (auto* TS = Cast<UMaterialExpressionTextureSample>(Expression))
	{
		EJ->SetStringField(TEXT("expressionType"), TEXT("TextureSample"));
		if (TS->Texture)
			EJ->SetStringField(TEXT("texture"), TS->Texture->GetPathName());
	}
	else if (auto* TC = Cast<UMaterialExpressionTextureCoordinate>(Expression))
	{
		EJ->SetStringField(TEXT("expressionType"), TEXT("TextureCoordinate"));
		EJ->SetNumberField(TEXT("coordinateIndex"), TC->CoordinateIndex);
		EJ->SetNumberField(TEXT("uTiling"), TC->UTiling);
		EJ->SetNumberField(TEXT("vTiling"), TC->VTiling);
	}
	else if (auto* CM = Cast<UMaterialExpressionComponentMask>(Expression))
	{
		EJ->SetStringField(TEXT("expressionType"), TEXT("ComponentMask"));
		EJ->SetBoolField(TEXT("r"), CM->R != 0);
		EJ->SetBoolField(TEXT("g"), CM->G != 0);
		EJ->SetBoolField(TEXT("b"), CM->B != 0);
		EJ->SetBoolField(TEXT("a"), CM->A != 0);
	}
	else if (auto* Custom = Cast<UMaterialExpressionCustom>(Expression))
	{
		EJ->SetStringField(TEXT("expressionType"), TEXT("Custom"));
		EJ->SetStringField(TEXT("code"), Custom->Code);
		EJ->SetStringField(TEXT("outputType"), StaticEnum<ECustomMaterialOutputType>()->GetNameStringByValue((int64)Custom->OutputType));
	}
	else if (auto* FI = Cast<UMaterialExpressionFunctionInput>(Expression))
	{
		EJ->SetStringField(TEXT("expressionType"), TEXT("FunctionInput"));
		EJ->SetStringField(TEXT("inputName"), FI->InputName.ToString());
	}
	else if (auto* FO = Cast<UMaterialExpressionFunctionOutput>(Expression))
	{
		EJ->SetStringField(TEXT("expressionType"), TEXT("FunctionOutput"));
		EJ->SetStringField(TEXT("outputName"), FO->OutputName.ToString());
	}
	else if (auto* MFC = Cast<UMaterialExpressionMaterialFunctionCall>(Expression))
	{
		EJ->SetStringField(TEXT("expressionType"), TEXT("MaterialFunctionCall"));
		if (MFC->MaterialFunction)
			EJ->SetStringField(TEXT("functionName"), MFC->MaterialFunction->GetName());
	}
	else
	{
		EJ->SetStringField(TEXT("expressionType"), Expression->GetClass()->GetName());
	}

	return EJ;
}