FLESH/Source/FLESHEditor/Private/DismembermentGraph/DismembermentGraphSchema.cpp

#include "DismembermentGraph/DismembermentGraphSchema.h"
#include "DismembermentGraph/DismembermentGraph.h"
#include "DismembermentGraph/DismembermentGraphNode.h"
#include "DismembermentGraph/DismembermentGraphNodeCut.h"
#include "DismembermentGraph/DismembermentGraphNodeBoneSelect.h"
#include "DismembermentGraph/DismembermentGraphNodeBloodEffect.h"
#include "DismembermentGraph/DismembermentGraphNodePhysics.h"
#include "DismembermentGraph/DismembermentGraphNodeOrgan.h"
#include "DismembermentGraph/DismembermentGraphNodeWound.h"
#include "EdGraphNode_Comment.h"
#include "Framework/Commands/GenericCommands.h"
#include "GraphEditorActions.h"
#include "ToolMenus.h"

#define LOCTEXT_NAMESPACE "DismembermentGraphSchema"

// Pin categories
const FName UDismembermentGraphSchema::PC_Exec("Exec");
const FName UDismembermentGraphSchema::PC_Bone("Bone");
const FName UDismembermentGraphSchema::PC_Cut("Cut");
const FName UDismembermentGraphSchema::PC_Blood("Blood");
const FName UDismembermentGraphSchema::PC_Physics("Physics");
const FName UDismembermentGraphSchema::PC_Organ("Organ");
const FName UDismembermentGraphSchema::PC_Wound("Wound");

void UDismembermentGraphSchema::GetGraphContextActions(FGraphContextMenuBuilder& ContextMenuBuilder) const
{
    // Add node actions
    const FText ToolTip = LOCTEXT("NewDismembermentNodeTooltip", "Add node here");
    const FText MenuDesc = LOCTEXT("NewDismembermentNodeMenuDesc", "Add Node");

    TArray<TSubclassOf<UDismembermentGraphNode>> NodeClasses;
    NodeClasses.Add(UDismembermentGraphNodeCut::StaticClass());
    NodeClasses.Add(UDismembermentGraphNodeBoneSelect::StaticClass());
    NodeClasses.Add(UDismembermentGraphNodeBloodEffect::StaticClass());
    NodeClasses.Add(UDismembermentGraphNodePhysics::StaticClass());
    NodeClasses.Add(UDismembermentGraphNodeOrgan::StaticClass());
    NodeClasses.Add(UDismembermentGraphNodeWound::StaticClass());

    for (TSubclassOf<UDismembermentGraphNode> NodeClass : NodeClasses)
    {
        UDismembermentGraphNode* NodeTemplate = NodeClass->GetDefaultObject<UDismembermentGraphNode>();
        FText NodeCategory = NodeTemplate->NodeCategory;
        FText NodeTitle = NodeTemplate->GetNodeTitle(ENodeTitleType::ListView);
        FText NodeTooltip = NodeTemplate->GetTooltipText();

        TSharedPtr<FDismembermentSchemaAction_NewNode> NewNodeAction(new FDismembermentSchemaAction_NewNode(
            NodeCategory,
            NodeTitle,
            NodeTooltip,
            0));

        NewNodeAction->NodeClass = NodeClass;
        ContextMenuBuilder.AddAction(NewNodeAction);
    }

    // Add comment node
    TSharedPtr<FDismembermentSchemaAction_NewNode> NewCommentAction(new FDismembermentSchemaAction_NewNode(
        FText::GetEmpty(),
        LOCTEXT("NewComment", "Comment"),
        LOCTEXT("NewCommentTooltip", "Add a comment node"),
        0));

    NewCommentAction->NodeClass = UEdGraphNode_Comment::StaticClass();
    ContextMenuBuilder.AddAction(NewCommentAction);
}

void UDismembermentGraphSchema::GetContextMenuActions(UToolMenu* Menu, UGraphNodeContextMenuContext* Context) const
{
    if (Context && Context->Node)
    {
        FToolMenuSection& Section = Menu->AddSection("DismembermentGraphNodeActions", LOCTEXT("NodeActionsMenuHeader", "Node Actions"));
        Section.AddMenuEntry(FGenericCommands::Get().Delete);
        Section.AddMenuEntry(FGenericCommands::Get().Cut);
        Section.AddMenuEntry(FGenericCommands::Get().Copy);
        Section.AddMenuEntry(FGenericCommands::Get().Duplicate);

        // Add preview action
        Section.AddMenuEntry(
            "PreviewNode",
            LOCTEXT("PreviewNode", "Preview Node"),
            LOCTEXT("PreviewNodeTooltip", "Preview the effect of this node"),
            FSlateIcon(),
            FUIAction(
                FExecuteAction::CreateLambda([Context]()
                {
                    if (const UDismembermentGraphNode* Node = Cast<UDismembermentGraphNode>(Context->Node))
                    {
                        // TODO: Implement node preview
                    }
                }),
                FCanExecuteAction::CreateLambda([Context]()
                {
                    return Context->Node != nullptr;
                })
            )
        );
    }

    // Add general graph actions
    {
        FToolMenuSection& Section = Menu->AddSection("DismembermentGraphActions", LOCTEXT("GraphActionsMenuHeader", "Graph Actions"));
        Section.AddMenuEntry(FGenericCommands::Get().SelectAll);
        Section.AddMenuEntry(
            "ArrangeNodes",
            LOCTEXT("ArrangeNodes", "Arrange Nodes"),
            LOCTEXT("ArrangeNodesTooltip", "Arrange the nodes in the graph"),
            FSlateIcon(),
            FUIAction(
                FExecuteAction::CreateLambda([Context]()
                {
                    if (Context && Context->Graph)
                    {
                        // TODO: Implement node arrangement
                    }
                }),
                FCanExecuteAction::CreateLambda([Context]()
                {
                    return Context && Context->Graph;
                })
            )
        );
    }
}

const FPinConnectionResponse UDismembermentGraphSchema::CanCreateConnection(const UEdGraphPin* A, const UEdGraphPin* B) const
{
    // Make sure the pins are valid
    if (!A || !B)
    {
        return FPinConnectionResponse(CONNECT_RESPONSE_DISALLOW, LOCTEXT("PinError", "Invalid pins"));
    }

    // Make sure the pins are not on the same node
    if (A->GetOwningNode() == B->GetOwningNode())
    {
        return FPinConnectionResponse(CONNECT_RESPONSE_DISALLOW, LOCTEXT("SameNode", "Can't connect pins on the same node"));
    }

    // Check pin directions
    const UEdGraphPin* InputPin = nullptr;
    const UEdGraphPin* OutputPin = nullptr;

    if (A->Direction == EGPD_Input && B->Direction == EGPD_Output)
    {
        InputPin = A;
        OutputPin = B;
    }
    else if (A->Direction == EGPD_Output && B->Direction == EGPD_Input)
    {
        InputPin = B;
        OutputPin = A;
    }
    else
    {
        return FPinConnectionResponse(CONNECT_RESPONSE_DISALLOW, LOCTEXT("IncompatibleDirections", "Incompatible pin directions"));
    }

    // Check pin types
    FDismembermentGraphPinType InputType = GetPinType(InputPin);
    FDismembermentGraphPinType OutputType = GetPinType(OutputPin);

    if (InputType != OutputType)
    {
        return FPinConnectionResponse(CONNECT_RESPONSE_DISALLOW, LOCTEXT("IncompatibleTypes", "Incompatible pin types"));
    }

    // Check for cycles
    if (WouldCreateCycle(OutputPin, InputPin))
    {
        return FPinConnectionResponse(CONNECT_RESPONSE_DISALLOW, LOCTEXT("CycleDetected", "Connection would create a cycle"));
    }

    // Check if the pins are already connected
    for (int32 i = 0; i < InputPin->LinkedTo.Num(); ++i)
    {
        if (InputPin->LinkedTo[i] == OutputPin)
        {
            return FPinConnectionResponse(CONNECT_RESPONSE_DISALLOW, LOCTEXT("AlreadyConnected", "Pins are already connected"));
        }
    }

    return FPinConnectionResponse(CONNECT_RESPONSE_MAKE, LOCTEXT("CreateConnection", "Create Connection"));
}

bool UDismembermentGraphSchema::TryCreateConnection(UEdGraphPin* A, UEdGraphPin* B) const
{
    // Check if the connection is valid
    const FPinConnectionResponse Response = CanCreateConnection(A, B);
    if (Response.Response != CONNECT_RESPONSE_MAKE)
    {
        return false;
    }

    // Break existing connections on input pins
    if (A->Direction == EGPD_Input)
    {
        A->BreakAllPinLinks();
    }
    else if (B->Direction == EGPD_Input)
    {
        B->BreakAllPinLinks();
    }

    // Make the connection
    A->MakeLinkTo(B);
    return true;
}

bool UDismembermentGraphSchema::ShouldHidePinDefaultValue(UEdGraphPin* Pin) const
{
    return Pin && Pin->LinkedTo.Num() > 0;
}

FLinearColor UDismembermentGraphSchema::GetPinTypeColor(const FEdGraphPinType& PinType) const
{
    FDismembermentGraphPinType DismembermentPinType;
    DismembermentPinType.PinCategory = PinType.PinCategory;
    return GetPinTypeColor(DismembermentPinType);
}

void UDismembermentGraphSchema::BreakNodeLinks(UEdGraphNode& TargetNode) const
{
    Super::BreakNodeLinks(TargetNode);
}

void UDismembermentGraphSchema::BreakPinLinks(UEdGraphPin& TargetPin, bool bSendsNodeNotification) const
{
    Super::BreakPinLinks(TargetPin, bSendsNodeNotification);
}

void UDismembermentGraphSchema::BreakSinglePinLink(UEdGraphPin* SourcePin, UEdGraphPin* TargetPin) const
{
    Super::BreakSinglePinLink(SourcePin, TargetPin);
}

void UDismembermentGraphSchema::DroppedAssetsOnGraph(const TArray<FAssetData>& Assets, const FVector2D& GraphPosition, UEdGraph* Graph) const
{
    // TODO: Implement asset dropping
}

void UDismembermentGraphSchema::DroppedAssetsOnNode(const TArray<FAssetData>& Assets, const FVector2D& GraphPosition, UEdGraphNode* Node) const
{
    // TODO: Implement asset dropping on nodes
}

void UDismembermentGraphSchema::DroppedAssetsOnPin(const TArray<FAssetData>& Assets, const FVector2D& GraphPosition, UEdGraphPin* Pin) const
{
    // TODO: Implement asset dropping on pins
}

void UDismembermentGraphSchema::GetAssetsNodeHoverMessage(const TArray<FAssetData>& Assets, const UEdGraphNode* HoverNode, FString& OutTooltipText, bool& OutOkIcon) const
{
    // TODO: Implement asset hover message
    OutTooltipText = TEXT("Drop to create a reference");
    OutOkIcon = true;
}

void UDismembermentGraphSchema::GetAssetsPinHoverMessage(const TArray<FAssetData>& Assets, const UEdGraphPin* HoverPin, FString& OutTooltipText, bool& OutOkIcon) const
{
    // TODO: Implement asset hover message
    OutTooltipText = TEXT("Drop to create a reference");
    OutOkIcon = true;
}

bool UDismembermentGraphSchema::CanConnectPins(const UEdGraphPin* PinA, const UEdGraphPin* PinB, FDismembermentGraphConnectionResponse& OutResponse) const
{
    // Check if the pins are valid
    if (!PinA || !PinB)
    {
        OutResponse.Response = FDismembermentGraphConnectionResponse_K2::ERROR_INCOMPATIBLE;
        OutResponse.Message = LOCTEXT("PinError", "Invalid pins");
        return false;
    }

    // Check if the pins are on the same node
    if (PinA->GetOwningNode() == PinB->GetOwningNode())
    {
        OutResponse.Response = FDismembermentGraphConnectionResponse_K2::ERROR_SELF_CONNECTION;
        OutResponse.Message = LOCTEXT("SameNode", "Can't connect pins on the same node");
        return false;
    }

    // Check pin directions
    const UEdGraphPin* InputPin = nullptr;
    const UEdGraphPin* OutputPin = nullptr;

    if (PinA->Direction == EGPD_Input && PinB->Direction == EGPD_Output)
    {
        InputPin = PinA;
        OutputPin = PinB;
    }
    else if (PinA->Direction == EGPD_Output && PinB->Direction == EGPD_Input)
    {
        InputPin = PinB;
        OutputPin = PinA;
    }
    else
    {
        OutResponse.Response = FDismembermentGraphConnectionResponse_K2::ERROR_INCOMPATIBLE;
        OutResponse.Message = LOCTEXT("IncompatibleDirections", "Incompatible pin directions");
        return false;
    }

    // Check pin types
    FDismembermentGraphPinType InputType = GetPinType(InputPin);
    FDismembermentGraphPinType OutputType = GetPinType(OutputPin);

    if (InputType != OutputType)
    {
        OutResponse.Response = FDismembermentGraphConnectionResponse_K2::ERROR_INCOMPATIBLE;
        OutResponse.Message = LOCTEXT("IncompatibleTypes", "Incompatible pin types");
        return false;
    }

    // Check for cycles
    if (WouldCreateCycle(OutputPin, InputPin))
    {
        OutResponse.Response = FDismembermentGraphConnectionResponse_K2::ERROR_CYCLE;
        OutResponse.Message = LOCTEXT("CycleDetected", "Connection would create a cycle");
        return false;
    }

    // Check if the pins are already connected
    for (int32 i = 0; i < InputPin->LinkedTo.Num(); ++i)
    {
        if (InputPin->LinkedTo[i] == OutputPin)
        {
            OutResponse.Response = FDismembermentGraphConnectionResponse_K2::ERROR_DISALLOWED;
            OutResponse.Message = LOCTEXT("AlreadyConnected", "Pins are already connected");
            return false;
        }
    }

    OutResponse.Response = FDismembermentGraphConnectionResponse_K2::OK;
    return true;
}

bool UDismembermentGraphSchema::WouldCreateCycle(const UEdGraphPin* PinA, const UEdGraphPin* PinB) const
{
    // Check if connecting these pins would create a cycle
    const UEdGraphNode* NodeA = PinA->GetOwningNode();
    const UEdGraphNode* NodeB = PinB->GetOwningNode();

    // If the nodes are the same, there's a cycle
    if (NodeA == NodeB)
    {
        return true;
    }

    // Depth-first search to check for cycles
    TSet<const UEdGraphNode*> VisitedNodes;
    TArray<const UEdGraphNode*> NodesToVisit;
    NodesToVisit.Push(NodeB);

    while (NodesToVisit.Num() > 0)
    {
        const UEdGraphNode* CurrentNode = NodesToVisit.Pop();
        if (VisitedNodes.Contains(CurrentNode))
        {
            continue;
        }

        VisitedNodes.Add(CurrentNode);

        // Check all output pins
        for (int32 PinIndex = 0; PinIndex < CurrentNode->Pins.Num(); ++PinIndex)
        {
            const UEdGraphPin* Pin = CurrentNode->Pins[PinIndex];
            if (Pin->Direction == EGPD_Output)
            {
                // Check all connections
                for (int32 LinkIndex = 0; LinkIndex < Pin->LinkedTo.Num(); ++LinkIndex)
                {
                    const UEdGraphPin* LinkedPin = Pin->LinkedTo[LinkIndex];
                    const UEdGraphNode* LinkedNode = LinkedPin->GetOwningNode();

                    // If we found NodeA, there's a cycle
                    if (LinkedNode == NodeA)
                    {
                        return true;
                    }

                    // Add the linked node to the nodes to visit
                    NodesToVisit.Push(LinkedNode);
                }
            }
        }
    }

    return false;
}

FDismembermentGraphPinType UDismembermentGraphSchema::GetPinType(const UEdGraphPin* Pin)
{
    return FDismembermentGraphPinType(Pin->PinType.PinCategory);
}

FLinearColor UDismembermentGraphSchema::GetPinTypeColor(const FDismembermentGraphPinType& PinType)
{
    if (PinType.PinCategory == PC_Exec)
    {
        return FLinearColor::White;
    }
    else if (PinType.PinCategory == PC_Bone)
    {
        return FLinearColor(0.9f, 0.9f, 0.2f);
    }
    else if (PinType.PinCategory == PC_Cut)
    {
        return FLinearColor(1.0f, 0.0f, 0.0f);
    }
    else if (PinType.PinCategory == PC_Blood)
    {
        return FLinearColor(0.8f, 0.0f, 0.0f);
    }
    else if (PinType.PinCategory == PC_Physics)
    {
        return FLinearColor(0.0f, 0.8f, 0.8f);
    }
    else if (PinType.PinCategory == PC_Organ)
    {
        return FLinearColor(0.8f, 0.4f, 0.4f);
    }
    else if (PinType.PinCategory == PC_Wound)
    {
        return FLinearColor(0.6f, 0.0f, 0.0f);
    }

    return FLinearColor::Black;
}

UDismembermentGraphNode* UDismembermentGraphSchema::CreateNode(TSubclassOf<UDismembermentGraphNode> NodeClass, UEdGraph* ParentGraph, float NodePosX, float NodePosY, bool bSelectNewNode)
{
    UDismembermentGraphNode* NewNode = NewObject<UDismembermentGraphNode>(ParentGraph, NodeClass);
    NewNode->CreateNewGuid();
    NewNode->PostPlacedNewNode();
    NewNode->NodePosX = NodePosX;
    NewNode->NodePosY = NodePosY;
    NewNode->AllocateDefaultPins();

    ParentGraph->AddNode(NewNode, true, bSelectNewNode);

    return NewNode;
}

// Schema action for creating a new node
FDismembermentSchemaAction_NewNode::FDismembermentSchemaAction_NewNode(const FText& InNodeCategory, const FText& InMenuDesc, const FText& InToolTip, const int32 InGrouping)
    : FEdGraphSchemaAction(InNodeCategory, InMenuDesc, InToolTip, InGrouping)
{
}

UEdGraphNode* FDismembermentSchemaAction_NewNode::PerformAction(UEdGraph* ParentGraph, UEdGraphPin* FromPin, const FVector2D Location, bool bSelectNewNode)
{
    UDismembermentGraphNode* NewNode = UDismembermentGraphSchema::CreateNode(NodeClass, ParentGraph, Location.X, Location.Y, bSelectNewNode);
    return NewNode;
}

#undef LOCTEXT_NAMESPACE