/**
 * Copyright 2022-2025 Theai, Inc. dba Inworld AI
 */
#include "GraphEditor/InworldGraphAssetEditor/AssetGraphSchema_InworldGraph.h"

#include "Framework/Commands/GenericCommands.h"
#include "Graph/InworldGraph.h"
#include "Graph/Nodes/InworldNode_Custom.h"
#include "Graph/Nodes/InworldNode_Routing.h"
#include "Graph/Nodes/InworldNode_GetRuntimeData.h"
#include "Graph/Nodes/InworldNode_SubgraphNode.h"
#include "Graph/Assets/InworldGraphAsset.h"
#include "GraphEditor/InworldGraphAssetEditor/ConnectionDrawingPolicy_InworldGraph.h"
#include "GraphEditor/InworldGraphAssetEditor/EdGraph_InworldGraph.h"
#include "GraphEditor/InworldGraphAssetEditor/EdNode_InworldGraphEdge.h"
#include "GraphEditor/InworldGraphAssetEditor/EdNode_InworldGraphNode.h"
#include "GraphEditor/InworldGraphAssetEditor/SEdNode_InworldGraphNode.h"
#include "GraphEditor/InworldGraphNodeInputOutputSlot.h"
#include "Graph/InworldData.h"
#include "GraphEditorActions.h"
#include "ToolMenus.h"
#include "GraphEditor/NodeDescriptionParser.h"
#include "AssetRegistry/AssetRegistryModule.h"
#include "InworldRuntimeSettings.h"
#include "EdGraphNode_Comment.h"
#include "HAL/PlatformApplicationMisc.h"
#include "Widgets/DeclarativeSyntaxSupport.h"

#include "ScopedTransaction.h"
#include "Graph/Edges/InworldEdge_WithCondition.h"
#include "Graph/Edges/InworldEdge_SafetyResult.h"
#include "Graph/Edges/InworldEdge_IsDataType.h"

#define LOCTEXT_NAMESPACE "AssetSchema_InworldGraph"

int32 UAssetGraphSchema_InworldGraph::CurrentCacheRefreshID = 0;

class FNodeVisitorCycleChecker
{
public:
	/** Check whether a loop in the graph would be caused by linking the passed-in nodes */
	bool CheckForLoop(UEdGraphNode* StartNode, UEdGraphNode* EndNode)
	{
		VisitedNodes.Add(StartNode);

		return TraverseNodes(EndNode);
	}

private:
	bool TraverseNodes(UEdGraphNode* Node)
	{
		VisitedNodes.Add(Node);

		for (auto MyPin : Node->Pins)
		{
			if (MyPin->Direction == EGPD_Output)
			{
				for (auto OtherPin : MyPin->LinkedTo)
				{
					UEdGraphNode* OtherNode = OtherPin->GetOwningNode();
					if (VisitedNodes.Contains(OtherNode))
					{
						// Only  an issue if this is a back-edge
						return false;
					}
					else if (!FinishedNodes.Contains(OtherNode))
					{
						// Only should traverse if this node hasn't been traversed
						if (!TraverseNodes(OtherNode))
							return false;
					}
				}
			}
		}

		VisitedNodes.Remove(Node);
		FinishedNodes.Add(Node);
		return true;
	};

	TSet<UEdGraphNode*> VisitedNodes;
	TSet<UEdGraphNode*> FinishedNodes;
};


#if ENGINE_MAJOR_VERSION == 5 && ENGINE_MINOR_VERSION < 6
UEdGraphNode* FAssetSchemaAction_InworldGraph_AddComment::PerformAction(class UEdGraph* ParentGraph, UEdGraphPin* FromPin, const FVector2D Location, bool bSelectNewNode)
{
	UEdGraphNode_Comment* const CommentTemplate = NewObject<UEdGraphNode_Comment>();

	FVector2D SpawnLocation = Location;
	FSlateRect Bounds;

	TSharedPtr<SGraphEditor> GraphEditorPtr = SGraphEditor::FindGraphEditorForGraph(ParentGraph);
	if (GraphEditorPtr.IsValid() && GraphEditorPtr->GetBoundsForSelectedNodes(/*out*/ Bounds, 50.0f))
	{
		CommentTemplate->SetBounds(Bounds);
		SpawnLocation.X = CommentTemplate->NodePosX;
		SpawnLocation.Y = CommentTemplate->NodePosY;
	}

	UEdGraphNode* const NewNode = FEdGraphSchemaAction_NewNode::SpawnNodeFromTemplate<UEdGraphNode_Comment>(
		ParentGraph, CommentTemplate, SpawnLocation, bSelectNewNode);

	return NewNode;
}
#else
UEdGraphNode* FAssetSchemaAction_InworldGraph_AddComment::PerformAction(class UEdGraph* ParentGraph, UEdGraphPin* FromPin, const FVector2f& Location, bool bSelectNewNode)
{
	UEdGraphNode_Comment* const CommentTemplate = NewObject<UEdGraphNode_Comment>();

	FVector2f SpawnLocation = Location;
	FSlateRect Bounds;

	TSharedPtr<SGraphEditor> GraphEditorPtr = SGraphEditor::FindGraphEditorForGraph(ParentGraph);
	if (GraphEditorPtr.IsValid())
	{
		// If they have a selection, build a bounding box around the selection
		if (GraphEditorPtr->GetBoundsForSelectedNodes(/*out*/ Bounds, 50.0f))
		{
			CommentTemplate->SetBounds(Bounds);
			SpawnLocation.X = CommentTemplate->NodePosX;
			SpawnLocation.Y = CommentTemplate->NodePosY;
		}
		else
		{
			// Otherwise initialize a default comment at the user's cursor location.
			SpawnLocation = GraphEditorPtr->GetPasteLocation2f();
		}
	}

	UEdGraphNode* const NewNode = FEdGraphSchemaAction_NewNode::SpawnNodeFromTemplate<UEdGraphNode_Comment>(ParentGraph, CommentTemplate, SpawnLocation, bSelectNewNode);

	return NewNode;
}
#endif

UEdGraphNode* FAssetSchemaAction_InworldGraph_NewNode::PerformAction(class UEdGraph* ParentGraph, UEdGraphPin* FromPin, const FVector2D Location, bool bSelectNewNode /*= true*/)
{
	UEdGraphNode* ResultNode = nullptr;
	if (NodeTemplate != nullptr)
	{
		const FScopedTransaction Transaction(LOCTEXT("InworldGraphEditorNewNode", "Inworld Graph Editor: New Node"));
		ParentGraph->Modify();
		if (FromPin != nullptr)
			FromPin->Modify();

		NodeTemplate->Rename(nullptr, ParentGraph);
		ParentGraph->AddNode(NodeTemplate, true, bSelectNewNode);

		NodeTemplate->CreateNewGuid();
		NodeTemplate->CreateRuntimeNode(SubgraphAsset);
		NodeTemplate->PostPlacedNewNode();
		NodeTemplate->AllocateDefaultPins();
		NodeTemplate->AutowireNewNode(FromPin);

		NodeTemplate->NodePosX = Location.X;
		NodeTemplate->NodePosY = Location.Y;

		NodeTemplate->SetFlags(RF_Transactional);

		ResultNode = NodeTemplate;

		if (UEdGraph_InworldGraph* EdInworldGraph = Cast<UEdGraph_InworldGraph>(ParentGraph))
		{
			EdInworldGraph->NotifyGraphChanged();
			EdInworldGraph->RebuildEditorGraph();
		}
	}

	return ResultNode;
}

void FAssetSchemaAction_InworldGraph_NewNode::AddReferencedObjects(FReferenceCollector& Collector)
{
	FEdGraphSchemaAction::AddReferencedObjects(Collector);
	Collector.AddReferencedObject(NodeTemplate);
	Collector.AddReferencedObject(SubgraphAsset);
}

UEdGraphNode* FAssetSchemaAction_InworldGraph_NewEdge::PerformAction(class UEdGraph* ParentGraph, UEdGraphPin* FromPin, const FVector2D Location, bool bSelectNewNode /*= true*/)
{
	UEdGraphNode* ResultNode = nullptr;

	if (NodeTemplate != nullptr)
	{
		// creating transaction for undo/redo
		const FScopedTransaction Transaction(LOCTEXT("InworldGraphEditorNewEdge", "Inworld Graph Editor: New Edge"));
		ParentGraph->Modify();
		if (FromPin != nullptr)
			FromPin->Modify();

		// renaming node and add to graph
		NodeTemplate->Rename(nullptr, ParentGraph);
		ParentGraph->AddNode(NodeTemplate, true, bSelectNewNode);

		// node initialisation
		NodeTemplate->CreateNewGuid();
		NodeTemplate->PostPlacedNewNode();
		NodeTemplate->AllocateDefaultPins();
		NodeTemplate->AutowireNewNode(FromPin);

		// node positioning
		NodeTemplate->NodePosX = Location.X;
		NodeTemplate->NodePosY = Location.Y;

		// set flag for transaction
		NodeTemplate->InworldGraphEdge->SetFlags(RF_Transactional);
		NodeTemplate->SetFlags(RF_Transactional);

		ResultNode = NodeTemplate;

		if (UEdGraph_InworldGraph* EdInworldGraph = Cast<UEdGraph_InworldGraph>(ParentGraph))
		{
			EdInworldGraph->NotifyGraphChanged();
			EdInworldGraph->RebuildEditorGraph();
		}
	}

	return ResultNode;
}

void FAssetSchemaAction_InworldGraph_NewEdge::AddReferencedObjects(FReferenceCollector& Collector)
{
	FEdGraphSchemaAction::AddReferencedObjects(Collector);
	Collector.AddReferencedObject(NodeTemplate);
}

void UAssetGraphSchema_InworldGraph::GetBreakLinkToSubMenuActions(UToolMenu* Menu, UEdGraphPin* InGraphPin)
{
	// Make sure we have a unique name for every entry in the list
	TMap<FString, uint32> LinkTitleCount;

	FToolMenuSection& Section = Menu->FindOrAddSection("InworldGraphAssetGraphSchemaPinActions");

	// Add all the links we could break from
	for (TArray<class UEdGraphPin*>::TConstIterator Links(InGraphPin->LinkedTo); Links; ++Links)
	{
		UEdGraphPin* Pin = *Links;
		UEdNode_InworldGraphEdge* Edge = Cast<UEdNode_InworldGraphEdge>(Pin->GetOwningNode());
		if (IsValid(Edge))
		{
			FString TitleString;
			FFormatNamedArguments Args;
			if (Pin->Direction == EEdGraphPinDirection::EGPD_Input)
			{
				Args.Add(TEXT("Direction"), FText::FromString("to"));
				TitleString = Edge->InworldGraphEdge->GetEndNodeName();
			}
			else
			{
				Args.Add(TEXT("Direction"), FText::FromString("from"));
				TitleString = Edge->InworldGraphEdge->GetStartNodeName();
			}

			uint32& Count = LinkTitleCount.FindOrAdd(TitleString);

			FText Description;
			Args.Add(TEXT("NodeTitle"), FText::FromString(TitleString));
			Args.Add(TEXT("NumberOfNodes"), Count);

			if (Count == 0)
			{
				Description = FText::Format(LOCTEXT("BreakDesc", "Break link {Direction} {NodeTitle}"), Args);
			}
			else
			{
				Description = FText::Format(LOCTEXT("BreakDescMulti", "Break link {Direction} {NodeTitle}"), Args);
			}
			++Count;

			Section.AddMenuEntry(NAME_None, Description, Description, FSlateIcon(),
				FUIAction(FExecuteAction::CreateUObject(this, &UAssetGraphSchema_InworldGraph::BreakSinglePinLink, const_cast<UEdGraphPin*>(InGraphPin), *Links)));
		}
	}
}

void UAssetGraphSchema_InworldGraph::SetNodeAsStartEnd(UEdNode_InworldGraphNode* Node)
{
	if (IsValid(Node))
	{
		Node->SetNodeAs(EInworldNodeType::START_END_NODE);
	}
}

void UAssetGraphSchema_InworldGraph::CopyNodeId(UEdNode_InworldGraphNode* Node)
{
	if (IsValid(Node) && Node->InworldGraphNode && Node->InworldGraphNode->RuntimeNode)
	{
		FPlatformApplicationMisc::ClipboardCopy(*Node->InworldGraphNode->RuntimeNode->GetNodeNameNative());
	}
}

void UAssetGraphSchema_InworldGraph::SetNodeAsStart(UEdNode_InworldGraphNode* Node)
{
	if (IsValid(Node))
	{
		Node->SetNodeAs(EInworldNodeType::START_NODE);
	}
}

void UAssetGraphSchema_InworldGraph::SetNodeAsEnd(UEdNode_InworldGraphNode* Node)
{
	if (IsValid(Node))
	{
		Node->SetNodeAs(EInworldNodeType::END_NODE);
	}
}

void UAssetGraphSchema_InworldGraph::RebuildNode(UEdNode_InworldGraphNode* Node)
{
	if (!IsValid(Node))
	{
		return;
	}

	if (UEdGraph* EdGraph = Node->GetGraph())
	{
		Node->Modify();

		if (Node->SEdNode)
		{
			Node->SEdNode->ForceUpdateGraphNode();
		}
	}
}

void UAssetGraphSchema_InworldGraph::SetEdgeType(UEdGraphNode* Edge, UClass* EdgeType)
{
	if (IsValid(Edge))
	{
		if (UEdNode_InworldGraphEdge* InworldEdge = Cast<UEdNode_InworldGraphEdge>(Edge))
		{
			InworldEdge->SetEdgeType(EdgeType);
		}
	}
}

EGraphType UAssetGraphSchema_InworldGraph::GetGraphType(const UEdGraph* TestEdGraph) const
{
	return GT_StateMachine;
}

void UAssetGraphSchema_InworldGraph::GetGraphContextActions(FGraphContextMenuBuilder& ContextMenuBuilder) const
{
	TArray<UClass*> AllNodeClasses;

	IAssetRegistry& AssetRegistry = FModuleManager::LoadModuleChecked<FAssetRegistryModule>("AssetRegistry").Get();
	TArray<FAssetData> BlueprintAssetData;
	FARFilter Filter;
	Filter.ClassPaths.Add(UBlueprint::StaticClass()->GetClassPathName());
	Filter.TagsAndValues.Emplace(
		FBlueprintTags::NativeParentClassPath, FAssetData(UInworldNode_Custom::StaticClass()).GetExportTextName());
	AssetRegistry.GetAssets(Filter, BlueprintAssetData);
	for (const FAssetData& AssetData : BlueprintAssetData)
	{
		if (UObject* Asset = AssetData.GetAsset())
		{
			if (UBlueprint* Blueprint = Cast<UBlueprint>(Asset))
			{
				if (Blueprint->GeneratedClass && Blueprint->GeneratedClass->IsChildOf(UInworldNode::StaticClass()))
				{
					AllNodeClasses.AddUnique(Blueprint->GetClass());
				}
			}
		}
	}

	for (TObjectIterator<UClass> It; It; ++It)
	{
		UClass* CurrentClass = *It;
		if (CurrentClass->IsChildOf(UInworldNode::StaticClass()))
		{
			AllNodeClasses.Add(CurrentClass);
		}
	}

	// Add Comment Context Menu
	if (!ContextMenuBuilder.FromPin)
	{
		TSharedPtr<SGraphEditor> GraphEditorPtr = SGraphEditor::FindGraphEditorForGraph(ContextMenuBuilder.CurrentGraph);
		const bool bIsManyNodesSelected = GraphEditorPtr.IsValid() && (GraphEditorPtr->GetNumberOfSelectedNodes() > 0);
		const FText MenuDescription = bIsManyNodesSelected ? LOCTEXT("CreateCommentSelection", "Add Comment around Selection")
														   : LOCTEXT("AddComment", "Add Comment");
		const FText ToolTip = LOCTEXT("CreateCommentSelectionTooltip", "Create a resizeable comment box around selected nodes.");

		TSharedPtr<FAssetSchemaAction_InworldGraph_AddComment> NewComment = TSharedPtr<FAssetSchemaAction_InworldGraph_AddComment>(
			new FAssetSchemaAction_InworldGraph_AddComment(FText::GetEmpty(), MenuDescription, ToolTip, 2));
		ContextMenuBuilder.AddAction(NewComment);
	}

	UInworldGraphAsset* EditorGraphAsset = CastChecked<UInworldGraphAsset>(ContextMenuBuilder.CurrentGraph->GetOuter());
	for (UClass* NodeClass : AllNodeClasses)
	{
		// skipping REINST and SKEL classes (temporary classes created during compilation)
		const bool bIsChild = NodeClass->IsChildOf(UInworldNode::StaticClass());
		const bool bIsAbstract = NodeClass->HasAnyClassFlags(CLASS_Abstract | CLASS_Deprecated | CLASS_NewerVersionExists);
		const bool bIsSkeleton = NodeClass->GetName().StartsWith(TEXT("SKEL_")) || NodeClass->GetName().StartsWith("REINST");
		// skip special nodes
		const bool bIsExcludeClass = NodeClass->IsChildOf(UInworldNode_SubgraphNode::StaticClass());
		const bool bExcludeTestPackage = !GetDefault<UInworldRuntimeEditorSettings>()->bAllowTestPluginAssets &&
								 NodeClass->GetClassPathName().GetPackageName().IsEqual("/Script/InworldTest");
		
		if (!bIsChild || bIsAbstract || bIsSkeleton || bIsExcludeClass || bExcludeTestPackage)
			continue;

		const UInworldNode* TempNode = NodeClass->GetDefaultObject<UInworldNode>();
		FString NodeName = TempNode->GetNodeName();
		if (NodeName.IsEmpty())
		{
			NodeName = NodeClass->GetDisplayNameText().ToString();
		}

		FText CategoryLabel = FText::GetEmpty();
		int ActionGroup = 0;
		if (NodeClass->IsChildOf(UInworldNode_Routing::StaticClass()) ||
			NodeClass->IsChildOf(UInworldNode_GetRuntimeData::StaticClass()))
		{
			ActionGroup = 1;
		}
		else if (NodeClass->IsChildOf(UInworldNode_Custom::StaticClass()))
		{
			CategoryLabel = LOCTEXT("InworldGraphNodeAction", "Custom Nodes");
		}
		else
		{
			CategoryLabel = LOCTEXT("InworldGraphNodeAction", "Built-In Nodes");
		}

		FNodeDescriptionParser Parser(NodeClass);
		TSharedPtr<FAssetSchemaAction_InworldGraph_NewNode> Action(new FAssetSchemaAction_InworldGraph_NewNode(
			CategoryLabel, FText::FromString(NodeName), FText::FromString(Parser.GetMainDescription()), ActionGroup));

		// Creating a node template
		Action->NodeTemplate = NewObject<UEdNode_InworldGraphNode>(ContextMenuBuilder.OwnerOfTemporaries);
		Action->NodeTemplate->InworldGraphNode = NewObject<UInworldNodeAsset>(EditorGraphAsset, UInworldNodeAsset::StaticClass());
		Action->NodeTemplate->InworldGraphNode->NodeClass = NodeClass;
		const TArray<FInworldNodeInputOutputSlot> InputSlots{
			FInworldNodeInputOutputSlot::GetInputOutputSlotsForNodeInput(TempNode)
		};
		const FInworldNodeInputOutputSlot OutputSlot{FInworldNodeInputOutputSlot::GetInputOutputSlotForNodeOutput(TempNode)};

		// If user is dragging from a pin, optionally filter actions by type-compatibility with that pin
		bool bAddAction = true;
		const bool bFilter = GetDefault<UInworldRuntimeEditorSettings>()->bEnableCompatibleNodesFilter;
		if (ContextMenuBuilder.FromPin && bFilter)
		{
			const UEdGraphPin* FromPin = ContextMenuBuilder.FromPin;
			UEdNode_InworldGraphNode* FromNodeEd = Cast<UEdNode_InworldGraphNode>(FromPin->GetOwningNode());
			if (FromNodeEd && Action->NodeTemplate && Action->NodeTemplate->InworldGraphNode)
			{
				if (FromPin->Direction == EGPD_Output)
				{
					const FInworldNodeInputOutputSlot& FromOutSlot = FromNodeEd->CachedOutputSlot;
					bAddAction = false;
					if (InputSlots.Num() == 0)
					{
						bAddAction = true;
					}
					else
					{
						for (const FInworldNodeInputOutputSlot& InputSlot : InputSlots)
						{
							// Empty ValidTypes on input means accept any
							if (InputSlot.ValidTypes.Num() == 0 || FromOutSlot.AreCompatibleWith(InputSlot))
							{
								bAddAction = true;
								break;
							}
						}
					}
				}
				else
				{
					int32 PinIndexOnNode = FromNodeEd->Pins.Find(const_cast<UEdGraphPin*>(FromPin));
					int32 InputSlotIndex = (PinIndexOnNode > 0) ? (PinIndexOnNode - 1) : 0;
					const TArray<FInworldNodeInputOutputSlot>& FromInputSlots = FromNodeEd->CachedInputSlots;

					if (FromInputSlots.IsValidIndex(InputSlotIndex))
					{
						// If our input slot is empty (wildcard), accept any candidate
						if (FromInputSlots[InputSlotIndex].ValidTypes.Num() == 0)
						{
							bAddAction = true;
						}
						else
						{
							bAddAction = FromInputSlots[InputSlotIndex].AreCompatibleWith(OutputSlot);
						}
					}
					else
					{
						bAddAction = false;
					}
				}
			}
		}

		// Adding an action to the context menu if compatible or not filtering
		if (bAddAction)
		{
			ContextMenuBuilder.AddAction(Action);
		}
	}

	// Add Subgraphs
	{
		TArray<FAssetData> SubgraphAssets;
		FARFilter SubgraphFilter;
		SubgraphFilter.ClassPaths.Add(UInworldGraphAsset::StaticClass()->GetClassPathName());
		AssetRegistry.GetAssets(SubgraphFilter, SubgraphAssets);

		const FText SubgraphsCategory = LOCTEXT("InworldGraphNodeAction", "Subgraphs");
		int32 ActionGroup = 0;
		for (const FAssetData& AssetData : SubgraphAssets)
		{
			UInworldGraphAsset* InworldGraphAsset = Cast<UInworldGraphAsset>(AssetData.GetAsset());
			FTopLevelAssetPath InworldGraphAssetPath = InworldGraphAsset->GetClass()->GetClassPathName();
			const bool bExcludeTestPackage = !GetDefault<UInworldRuntimeEditorSettings>()->bAllowTestPluginAssets &&
											 InworldGraphAssetPath.GetPackageName().IsEqual("/Script/InworldTest");

			if (InworldGraphAsset->IsValidAsSubgraph() || bExcludeTestPackage)
			{
				const FText MenuDesc = FText::FromName(AssetData.AssetName);
				const FText ToolTip = FText::FromString(InworldGraphAsset->Description);
			
				TSharedPtr<FAssetSchemaAction_InworldGraph_NewNode> Action(
					new FAssetSchemaAction_InworldGraph_NewNode(SubgraphsCategory, MenuDesc, ToolTip, ActionGroup));
				
				// Creating a node template
				Action->NodeTemplate = NewObject<UEdNode_InworldGraphNode>(ContextMenuBuilder.OwnerOfTemporaries);
				Action->NodeTemplate->InworldGraphNode =
					NewObject<UInworldNodeAsset>(EditorGraphAsset, UInworldNodeAsset::StaticClass());
				Action->NodeTemplate->InworldGraphNode->NodeClass = UInworldNode_SubgraphNode::StaticClass();
				Action->NodeTemplate->InworldGraphNode->NodeTitle = AssetData.AssetName.ToString();
				Action->SubgraphAsset = InworldGraphAsset;

				ContextMenuBuilder.AddAction(Action);
			}
		}
	}
}

void UAssetGraphSchema_InworldGraph::GetContextMenuActions(UToolMenu* Menu, UGraphNodeContextMenuContext* Context) const
{
	if (Context->Pin)
	{
		{
			FToolMenuSection& Section = Menu->AddSection("InworldGraphAssetGraphSchemaNodeActions", LOCTEXT("PinActionsMenuHeader", "Pin Actions"));
			// Only display the 'Break Links' option if there is a link to break!
			if (Context->Pin->LinkedTo.Num() > 0)
			{
				// add sub menu for break link to
				if (Context->Pin->LinkedTo.Num() > 1)
				{
					Section.AddSubMenu("BreakLinkTo", LOCTEXT("BreakLinkTo", "Break Link To..."),
						LOCTEXT("BreakSpecificLinks", "Break a specific link..."),
						FNewToolMenuDelegate::CreateUObject((UAssetGraphSchema_InworldGraph* const) this,
							&UAssetGraphSchema_InworldGraph::GetBreakLinkToSubMenuActions, const_cast<UEdGraphPin*>(Context->Pin)));
				}
				else
				{
					((UAssetGraphSchema_InworldGraph* const) this)->GetBreakLinkToSubMenuActions(Menu, const_cast<UEdGraphPin*>(Context->Pin));
				}
			}
		}
	}
	else if (Context->Node)
	{
		{
			if (UEdNode_InworldGraphNode* EdInworldNode =
					Cast<UEdNode_InworldGraphNode>(const_cast<UEdGraphNode*>(Context->Node.Get())))
			{
				FToolMenuSection& Section =
					Menu->AddSection("InworldGraphAssetGraphSchemaNodeActions", LOCTEXT("ClassActionsMenuHeader", "Node Actions"));

				if (!EdInworldNode->ErrorMsg.IsEmpty())
				{
					Section.AddMenuEntry("RebuildNode", LOCTEXT("RebuildNode", "Rebuild Node"),
						LOCTEXT("RebuildNode_Tooltip", "Rebuild node and pins"), FSlateIcon(),
						FUIAction(FExecuteAction::CreateUObject((UAssetGraphSchema_InworldGraph* const) this,
							&UAssetGraphSchema_InworldGraph::RebuildNode, EdInworldNode)));
				}

				Section.AddMenuEntry("CopyNodeId", LOCTEXT("CopyNodeId", "Copy Node Id"), LOCTEXT("CopyNodeId_Tooltip", ""),
					FSlateIcon(), FUIAction(FExecuteAction::CreateUObject((UAssetGraphSchema_InworldGraph* const) this, 
						&UAssetGraphSchema_InworldGraph::CopyNodeId, EdInworldNode)));

				if (EdInworldNode->IsRootNode() || EdInworldNode->IsLeafNode())
				{
					Section.AddSeparator("SpecialCommands");

					if (EdInworldNode->IsLeafNode())
					{
						Section.AddMenuEntry("SetNodeAsEnd", LOCTEXT("SetNodeAsEnd", "Set As End"),
							LOCTEXT("SetNodeAsEnd_Tooltip", ""), FSlateIcon(),
							FUIAction(FExecuteAction::CreateUObject((UAssetGraphSchema_InworldGraph* const) this,
								&UAssetGraphSchema_InworldGraph::SetNodeAsEnd, EdInworldNode)));
					}
					if (EdInworldNode->IsRootNode())
					{
						Section.AddMenuEntry("SetNodeAsStart", LOCTEXT("SetNodeAsStart", "Set As Start"),
							LOCTEXT("SetNodeAsStart_Tooltip", ""), FSlateIcon(),
							FUIAction(FExecuteAction::CreateUObject((UAssetGraphSchema_InworldGraph* const) this,
								&UAssetGraphSchema_InworldGraph::SetNodeAsStart, EdInworldNode)));
					}
					if (EdInworldNode->IsRootNode() && EdInworldNode->IsLeafNode())
					{
						Section.AddMenuEntry("SetNodeAsStartEnd", LOCTEXT("SetNodeAsStart", "Set As Start and End"),
							LOCTEXT("SetNodeAsStart_Tooltip", ""), FSlateIcon(),
							FUIAction(FExecuteAction::CreateUObject((UAssetGraphSchema_InworldGraph* const) this,
								&UAssetGraphSchema_InworldGraph::SetNodeAsStartEnd, EdInworldNode)));
					}
				}

				// Section.AddMenuEntry(FGenericCommands::Get().Cut);
				// Section.AddMenuEntry(FGenericCommands::Get().Copy);
				// Section.AddMenuEntry(FGenericCommands::Get().Duplicate);
				// Section.AddMenuEntry(FGraphEditorCommands::Get().BreakNodeLinks);
			}
			else if (UEdNode_InworldGraphEdge* EdInworldEdge =
						 Cast<UEdNode_InworldGraphEdge>(const_cast<UEdGraphNode*>(Context->Node.Get())))
			{
				FToolMenuSection& Section = Menu->AddSection("InworldGraphAssetGraphSchemaNodeActions", 
						LOCTEXT("ClassActionsMenuHeader", "Set Edge Type"));
				
				TArray<UClass*> AllEdgeClasses;

				IAssetRegistry& AssetRegistry = FModuleManager::LoadModuleChecked<FAssetRegistryModule>("AssetRegistry").Get();
				TArray<FAssetData> BlueprintAssetData;
				FARFilter Filter;
				Filter.ClassPaths.Add(UBlueprint::StaticClass()->GetClassPathName());
				Filter.TagsAndValues.Emplace(FBlueprintTags::NativeParentClassPath, FAssetData(UInworldEdge_WithCondition::StaticClass()).GetExportTextName());
				AssetRegistry.GetAssets(Filter, BlueprintAssetData);
				for (const FAssetData& AssetData : BlueprintAssetData)
				{
					if (UObject* Asset = AssetData.GetAsset())
					{
						if (UBlueprint* Blueprint = Cast<UBlueprint>(Asset))
						{
							if (Blueprint->GeneratedClass && Blueprint->GeneratedClass->IsChildOf(UInworldEdge::StaticClass()))
							{
								AllEdgeClasses.AddUnique(Blueprint->GetClass());
							}
						}
					}
				}

				for (TObjectIterator<UClass> It; It; ++It)
				{
					UClass* CurrentClass = *It;
					if (CurrentClass->IsChildOf(UInworldEdge::StaticClass()))
					{
						AllEdgeClasses.Add(CurrentClass);
					}
				}

				TArray<TPair<FString, UClass*>> SpecialEdgeTypes;
				TArray<TPair<FString, UClass*>> ConditionalEdgeTypes;

				for (UClass* EdgeClass : AllEdgeClasses)
				{
					// skipping REINST and SKEL classes (temporary classes created during compilation)
					const bool bIsChild = EdgeClass->IsChildOf(UInworldEdge::StaticClass());
					const bool bIsAbstract = EdgeClass->HasAnyClassFlags(CLASS_Abstract | CLASS_Deprecated | CLASS_NewerVersionExists);
					const bool bIsSkeleton = EdgeClass->GetName().StartsWith(TEXT("SKEL_")) || EdgeClass->GetName().StartsWith("REINST");
					const bool bExcludeTestPackage = !GetDefault<UInworldRuntimeEditorSettings>()->bAllowTestPluginAssets &&
												 EdgeClass->GetClassPathName().GetPackageName().IsEqual("/Script/InworldTest");

					if (!bIsChild || bIsAbstract || bIsSkeleton || bExcludeTestPackage)
						continue;

					const UInworldEdge* TempEdge = EdgeClass->GetDefaultObject<UInworldEdge>();
					if (TempEdge == nullptr)
						continue;

					FString EdgeTypeName{};
					if (EdgeClass == UInworldEdge::StaticClass())
					{
						SpecialEdgeTypes.Add(TPair<FString, UClass*>{TEXT("Generic Edge"), EdgeClass});
					}
					else
					{
						if (const UInworldEdge_WithCondition* TempCustomEdge = Cast<UInworldEdge_WithCondition>(TempEdge))
						{
							if (TempCustomEdge->EdgeTitle.IsEmpty())
							{
								EdgeTypeName = EdgeClass->GetName();
								EdgeTypeName.RemoveFromStart("InworldEdge_");
								EdgeTypeName.RemoveFromEnd("_C");
							}
							else
							{
								EdgeTypeName = TempCustomEdge->EdgeTitle;
							}
						}
						else
						{
							EdgeTypeName = EdgeClass->GetName();
							EdgeTypeName.RemoveFromStart("InworldEdge_");
							EdgeTypeName.RemoveFromEnd("_C");
						}

						if (EdgeClass == UInworldEdge_SafetyResult::StaticClass() ||
							EdgeClass == UInworldEdge_IsDataType::StaticClass())
						{
							SpecialEdgeTypes.Add(TPair<FString, UClass*>{EdgeTypeName, EdgeClass});
						}
						else
						{
							ConditionalEdgeTypes.Add(TPair<FString, UClass*>{EdgeTypeName, EdgeClass});
						}
					}
				}

				for (const TPair<FString, UClass*>& Pair : SpecialEdgeTypes)
				{
					const FString& EdgeTypeName = Pair.Key;
					UClass* const EdgeClass = Pair.Value;
					Section.AddMenuEntry(FName(EdgeTypeName), FText::FromString(EdgeTypeName), FText::FromString(EdgeTypeName),
							FSlateIcon(), FUIAction(FExecuteAction::CreateUObject((UAssetGraphSchema_InworldGraph* const) this,
								&UAssetGraphSchema_InworldGraph::SetEdgeType, const_cast<UEdGraphNode*>(Context->Node.Get()), EdgeClass)));
				}

				if (ConditionalEdgeTypes.Num() > 0 && SpecialEdgeTypes.Num() > 0)
				{
					Section.AddSeparator(TEXT("EdgeTypes_GroupSeparator"));
				}

				for (const TPair<FString, UClass*>& Pair : ConditionalEdgeTypes)
				{
					const FString& EdgeTypeName = Pair.Key;
					UClass* const EdgeClass = Pair.Value;
					Section.AddMenuEntry(FName(EdgeTypeName), FText::FromString(EdgeTypeName), FText::FromString(EdgeTypeName),
							FSlateIcon(), FUIAction(FExecuteAction::CreateUObject((UAssetGraphSchema_InworldGraph* const) this,
								&UAssetGraphSchema_InworldGraph::SetEdgeType, const_cast<UEdGraphNode*>(Context->Node.Get()), EdgeClass)));
				}
			}

			//Section.AddMenuEntry(FGenericCommands::Get().Delete);
		}
	}

	Super::GetContextMenuActions(Menu, Context);
}

const FPinConnectionResponse UAssetGraphSchema_InworldGraph::CanCreateConnection(const UEdGraphPin* A, const UEdGraphPin* B) const
{
	// Make sure the pins are not same direction
	if (A->Direction == B->Direction)
	{
		return FPinConnectionResponse(
			CONNECT_RESPONSE_DISALLOW, LOCTEXT("PinErrorSameNode", "Can't connect pins of the same direction."));
	}

	const UEdGraphPin* Out = A->Direction == EGPD_Output ? A : B;
	const UEdGraphPin* In = A->Direction == EGPD_Output ? B : A;

	UEdNode_InworldGraphNode* EdNode_Out = Cast<UEdNode_InworldGraphNode>(Out->GetOwningNode());
	UEdNode_InworldGraphNode* EdNode_In = Cast<UEdNode_InworldGraphNode>(In->GetOwningNode());

	if (EdNode_Out == nullptr || EdNode_In == nullptr)
	{
		return FPinConnectionResponse(CONNECT_RESPONSE_DISALLOW, LOCTEXT("PinError", "Not a valid Nodes."));
	}

	UEdGraph_InworldGraph* EdGraph = Cast<UEdGraph_InworldGraph>(Out->GetOwningNode()->GetGraph());

	// Determine if we can have cycles
	bool bAllowCycles = false;
	if (EdGraph != nullptr)
	{
		bAllowCycles = EdGraph->GetInworldGraph()->bCanBeCyclical;
	}

	// Determine if allowed loops
	bool bAllowLoops = false;
	if (EdGraph != nullptr)
	{
		bAllowLoops = EdGraph->GetInworldGraph()->bAllowLoops;
	}

	// check for loops
	if (!bAllowLoops && A->GetOwningNode() == B->GetOwningNode())
	{
		return FPinConnectionResponse(CONNECT_RESPONSE_DISALLOW, LOCTEXT("PinErrorSameNode", "Can't connect node to itself"));
	}

	// check for cycles
	FNodeVisitorCycleChecker CycleChecker;
	if (!bAllowCycles && !CycleChecker.CheckForLoop(Out->GetOwningNode(), In->GetOwningNode()))
	{
		return FPinConnectionResponse(CONNECT_RESPONSE_DISALLOW, LOCTEXT("PinErrorCycle", "Can't create a graph cycle"));
	}

	FText ErrorMessage;
	if (!EdNode_Out->InworldGraphNode->CanCreateConnectionTo(
			EdNode_In->InworldGraphNode, EdNode_Out->GetOutputPin()->LinkedTo.Num(), ErrorMessage))
	{
		return FPinConnectionResponse(CONNECT_RESPONSE_DISALLOW, ErrorMessage);
	}
	if (!EdNode_In->InworldGraphNode->CanCreateConnectionFrom(
			EdNode_Out->InworldGraphNode, EdNode_In->GetInputPin()->LinkedTo.Num(), ErrorMessage))
	{
		return FPinConnectionResponse(CONNECT_RESPONSE_DISALLOW, ErrorMessage);
	}

	// validate data types based
	int32 InputPinIndexOnNode = INDEX_NONE;
	UEdNode_InworldGraphNode* TargetNode = EdNode_In;
	if (TargetNode)
	{
		InputPinIndexOnNode = TargetNode->Pins.Find(const_cast<UEdGraphPin*>(In));
	}

	if (EdNode_Out && EdNode_In)
	{
		int32 InputSlotIndex = (InputPinIndexOnNode > 0) ? (InputPinIndexOnNode - 1) : 0;
		if (EdNode_In->CachedInputSlots.IsValidIndex(InputSlotIndex))
		{
			const FInworldNodeInputOutputSlot& InputSlot = EdNode_In->CachedInputSlots[InputSlotIndex];
			if (!EdNode_Out->CachedOutputSlot.AreCompatibleWith(InputSlot))
			{
				const FString OutTypesStr = EdNode_Out->CachedOutputSlot.TypesToString();
				const FString InTypesStr = InputSlot.TypesToString();
				const FText Error = FText::FromString(
					FString::Printf(TEXT("Type mismatch\n%s -> %s"), *OutTypesStr, *InTypesStr));
				if (GetDefault<UInworldRuntimeEditorSettings>()->bStrictDataTypesCompatibleCheck)
				{
					return FPinConnectionResponse(CONNECT_RESPONSE_DISALLOW, Error);
				}
				else
				{
					return FPinConnectionResponse(CONNECT_RESPONSE_MAKE_WITH_CONVERSION_NODE, Error);
				}
			}
		}
	}

	return FPinConnectionResponse(CONNECT_RESPONSE_MAKE_WITH_CONVERSION_NODE, LOCTEXT("PinConnect", "Connect nodes with edge"));
}

bool UAssetGraphSchema_InworldGraph::TryCreateConnection(UEdGraphPin* A, UEdGraphPin* B) const
{
	// check that this is not a connection between the pins of one node.
	//if (A->GetOwningNode() == B->GetOwningNode())
	//	return false;

	// define the output and input pins
	UEdGraphPin* OutputPin = (A->Direction == EGPD_Output) ? A : B;
	UEdGraphPin* InputPin = (A->Direction == EGPD_Output) ? B : A;

	// Check that the pins have right directions
	if (OutputPin->Direction != EGPD_Output || InputPin->Direction != EGPD_Input)
		return false;

	UEdNode_InworldGraphNode* NodeA = Cast<UEdNode_InworldGraphNode>(OutputPin->GetOwningNode());
	UEdNode_InworldGraphNode* NodeB = Cast<UEdNode_InworldGraphNode>(InputPin->GetOwningNode());

	// check that this is not a duplicate connection.
	for (UEdGraphPin* TestPin : OutputPin->LinkedTo)
	{
		UEdGraphNode* ChildNode = TestPin->GetOwningNode();
		if (UEdNode_InworldGraphEdge* EdNode_Edge = Cast<UEdNode_InworldGraphEdge>(ChildNode))
		{
			UEdGraphPin* EndPin = EdNode_Edge->GetEndPin();
			if (EndPin && EndPin->GetOwningNode() == NodeB && EndPin == InputPin)
				return false;
		}
	}

	if (NodeA && NodeB)
	{
		// Schema validation result
		const FPinConnectionResponse& Response = CanCreateConnection(OutputPin, InputPin);
		if (Response.Response == ECanCreateConnectionResponse::CONNECT_RESPONSE_DISALLOW)
		{
			return false;
		}
		// create a direct connection between pins (schema validation above already passed)
		Super::TryCreateConnection(OutputPin, InputPin);
		return true;
	}

	return false;
}

bool UAssetGraphSchema_InworldGraph::CreateAutomaticConversionNodeAndConnections(UEdGraphPin* A, UEdGraphPin* B) const
{
	// define the output and input pins
	UEdGraphPin* OutputPin = (A->Direction == EGPD_Output) ? A : B;
	UEdGraphPin* InputPin = (A->Direction == EGPD_Output) ? B : A;

	UEdNode_InworldGraphNode* OutputNode = Cast<UEdNode_InworldGraphNode>(OutputPin->GetOwningNode());
	UEdNode_InworldGraphNode* InputNode = Cast<UEdNode_InworldGraphNode>(InputPin->GetOwningNode());

	// checking validity
	if (!OutputNode || !InputNode || OutputPin->Direction != EGPD_Output || InputPin->Direction != EGPD_Input)
		return false;

	// find pin indices in your nodes
	int32 OutputPinIndex = OutputNode->Pins.Find(OutputPin);
	int32 InputPinIndex = InputNode->Pins.Find(InputPin);
	UInworldNodeAsset* OutputNodeAsset = OutputNode->InworldGraphNode;

	FVector2D InitPos{double(OutputNode->NodePosX), double(OutputNode->NodePosY)};
	if (InitPos.IsNearlyZero())
	{
		InitPos.X = InputNode->NodePosX;
		InitPos.Y = InputNode->NodePosY;
	}

	FAssetSchemaAction_InworldGraph_NewEdge Action;
	Action.NodeTemplate = NewObject<UEdNode_InworldGraphEdge>(OutputNode->GetGraph());
	Action.NodeTemplate->SetEdge(
		NewObject<UInworldEdgeAsset>(OutputNode->GetInworldGraph(), UInworldEdgeAsset::StaticClass()));
	UEdNode_InworldGraphEdge* EdgeNode =
		Cast<UEdNode_InworldGraphEdge>(Action.PerformAction(OutputNode->GetGraph(), nullptr, InitPos, false));

	// save pin indices in InworldGraphEdge
	EdgeNode->InworldGraphEdge->StartPinIndex = OutputPinIndex;
	EdgeNode->InworldGraphEdge->EndPinIndex = InputPinIndex;

	// create connections between specific pins
	EdgeNode->GetInputPin()->Modify();
	EdgeNode->GetInputPin()->LinkedTo.Empty();
	OutputPin->Modify();
	EdgeNode->GetInputPin()->MakeLinkTo(OutputPin);

	EdgeNode->GetOutputPin()->Modify();
	EdgeNode->GetOutputPin()->LinkedTo.Empty();
	InputPin->Modify();
	EdgeNode->GetOutputPin()->MakeLinkTo(InputPin);

	// Create Runtime Edge
	EdgeNode->InworldGraphEdge->CreateRuntimeEdge();

	UEdGraph_InworldGraph* EdInworldGraph = Cast<UEdGraph_InworldGraph>(OutputNode->GetGraph());
	if (EdInworldGraph)
	{
		EdInworldGraph->RebuildEditorGraph();
	}
	return true;
}

class FConnectionDrawingPolicy* UAssetGraphSchema_InworldGraph::CreateConnectionDrawingPolicy(int32 InBackLayerID,
	int32 InFrontLayerID, float InZoomFactor, const FSlateRect& InClippingRect, class FSlateWindowElementList& InDrawElements,
	class UEdGraph* InGraphObj) const
{
	return new FConnectionDrawingPolicy_InworldGraph(
		InBackLayerID, InFrontLayerID, InZoomFactor, InClippingRect, InDrawElements, InGraphObj);
}

FLinearColor UAssetGraphSchema_InworldGraph::GetPinTypeColor(const FEdGraphPinType& PinType) const
{
	return FColor::White;
}

void UAssetGraphSchema_InworldGraph::BreakNodeLinks(UEdGraphNode& TargetNode) const
{
	const FScopedTransaction Transaction(NSLOCTEXT("UnrealEd", "GraphEd_BreakNodeLinks", "Break Node Links"));

	Super::BreakNodeLinks(TargetNode);
}

void UAssetGraphSchema_InworldGraph::BreakSinglePinLink(UEdGraphPin* SourcePin, UEdGraphPin* TargetPin) const
{
	const FScopedTransaction Transaction(NSLOCTEXT("UnrealEd", "GraphEd_BreakSinglePinLink", "Break Pin Link"));

	Super::BreakSinglePinLink(SourcePin, TargetPin);
}

UEdGraphPin* UAssetGraphSchema_InworldGraph::DropPinOnNode(UEdGraphNode* InTargetNode, const FName& InSourcePinName,
	const FEdGraphPinType& InSourcePinType, EEdGraphPinDirection InSourcePinDirection) const
{
	UEdNode_InworldGraphNode* EdNode = Cast<UEdNode_InworldGraphNode>(InTargetNode);
	switch (InSourcePinDirection)
	{
		case EGPD_Input:
			return EdNode->GetOutputPin();
		case EGPD_Output:
			return EdNode->GetInputPin();
		default:
			return nullptr;
	}
}

bool UAssetGraphSchema_InworldGraph::SupportsDropPinOnNode(UEdGraphNode* InTargetNode, const FEdGraphPinType& InSourcePinType,
	EEdGraphPinDirection InSourcePinDirection, FText& OutErrorMessage) const
{
	return Cast<UEdNode_InworldGraphNode>(InTargetNode) != nullptr;
}

bool UAssetGraphSchema_InworldGraph::IsCacheVisualizationOutOfDate(int32 InVisualizationCacheID) const
{
	return CurrentCacheRefreshID != InVisualizationCacheID;
}

int32 UAssetGraphSchema_InworldGraph::GetCurrentVisualizationCacheID() const
{
	return CurrentCacheRefreshID;
}

void UAssetGraphSchema_InworldGraph::ForceVisualizationCacheClear() const
{
	++CurrentCacheRefreshID;
}

TSharedPtr<FEdGraphSchemaAction> UAssetGraphSchema_InworldGraph::GetCreateCommentAction() const
{
	return TSharedPtr<FEdGraphSchemaAction>(static_cast<FEdGraphSchemaAction*>(new FAssetSchemaAction_InworldGraph_AddComment));
}

#undef LOCTEXT_NAMESPACE
