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

#include "Graph/Assets/InworldGraphAsset.h"
#include "GraphEditor/InworldGraphAssetEditor/EdNode_InworldGraphEdge.h"
#include "GraphEditor/InworldGraphAssetEditor/EdNode_InworldGraphNode.h"

FConnectionDrawingPolicy_InworldGraph::FConnectionDrawingPolicy_InworldGraph(int32 InBackLayerID, int32 InFrontLayerID,
	float ZoomFactor, const FSlateRect& InClippingRect, FSlateWindowElementList& InDrawElements, UEdGraph* InGraphObj)
	: FConnectionDrawingPolicy(InBackLayerID, InFrontLayerID, ZoomFactor, InClippingRect, InDrawElements), GraphObj(InGraphObj)
{
}

void FConnectionDrawingPolicy_InworldGraph::DetermineWiringStyle(
	UEdGraphPin* OutputPin, UEdGraphPin* InputPin, FConnectionParams& Params)
{
	Params.AssociatedPin1 = OutputPin;
	Params.AssociatedPin2 = InputPin;
	Params.WireThickness = 1.5f;
	Params.bUserFlag2 = true;	  // is required edge - true

	// Default base connection colors
	Params.WireColor = FLinearColor(0.9f, 0.9f, 0.9f, 1.0f);

	// check how many connections connectoed to the input pin
	if (InputPin)
	{
		// find the real pin of the node (not the edge)
		UEdNode_InworldGraphNode* TargetNode = nullptr;
		int32 TargetPinIndex = -1;
		UEdNode_InworldGraphEdge* EdgeNode = Cast<UEdNode_InworldGraphEdge>(InputPin->GetOwningNode());
		if (IsValid(EdgeNode))
		{
			// if it's an edge, use the index of the saved pin
			TargetNode = EdgeNode->GetEndNode();
			TargetPinIndex = EdgeNode->InworldGraphEdge->EndPinIndex;
			Params.WireColor = EdgeNode->InworldGraphEdge->GetEdgeColor();
		}
		else
		{
			// if this is a node, use the current pin
			TargetNode = Cast<UEdNode_InworldGraphNode>(InputPin->GetOwningNode());
			TargetPinIndex = TargetNode->Pins.Find(InputPin);
		}

		// if target node and pin index found
		if (TargetNode && TargetPinIndex >= 0)
		{
			// get the target input pin
			UEdGraphPin* TargetPin = TargetNode->Pins[TargetPinIndex];

			// check the number of incoming connections
			int32 IncomingConnections = 0;

			// count direct connections
			IncomingConnections += TargetPin->LinkedTo.Num();

			// set the color depending on the number of connections
			if (IncomingConnections > 1)
			{
				EdgeNode->InworldGraphEdge->SetRequired(false);
			}
			Params.bUserFlag2 = EdgeNode->InworldGraphEdge->GetIsRequired();
		}
	}

	if (HoveredPins.Num() > 0)
	{
		ApplyHoverDeemphasis(OutputPin, InputPin, Params.WireThickness, Params.WireColor);
	}
}

void FConnectionDrawingPolicy_InworldGraph::Draw(
	TMap<TSharedRef<SWidget>, FArrangedWidget>& InPinGeometries, FArrangedChildren& ArrangedNodes)
{
	// Build an acceleration structure to quickly find geometry for the nodes
	NodeWidgetMap.Empty();
	for (int32 NodeIndex = 0; NodeIndex < ArrangedNodes.Num(); ++NodeIndex)
	{
		FArrangedWidget& CurWidget = ArrangedNodes[NodeIndex];
		TSharedRef<SGraphNode> ChildNode = StaticCastSharedRef<SGraphNode>(CurWidget.Widget);
		NodeWidgetMap.Add(ChildNode->GetNodeObj(), NodeIndex);
	}

	// Now draw
	FConnectionDrawingPolicy::Draw(InPinGeometries, ArrangedNodes);
}

#if ENGINE_MAJOR_VERSION == 5 && ENGINE_MINOR_VERSION < 6
void FConnectionDrawingPolicy_InworldGraph::DrawPreviewConnector(const FGeometry& PinGeometry, const FVector2D& StartPoint, const FVector2D& EndPoint, UEdGraphPin* Pin)
{
	FConnectionParams Params;
	DetermineWiringStyle(Pin, nullptr, /*inout*/ Params);

	if (Pin->Direction == EEdGraphPinDirection::EGPD_Output)
	{
		DrawSplineWithArrow(FGeometryHelper::FindClosestPointOnGeom(PinGeometry, EndPoint), EndPoint, Params);
	}
	else
	{
		DrawSplineWithArrow(FGeometryHelper::FindClosestPointOnGeom(PinGeometry, StartPoint), StartPoint, Params);
	}
}
#else
void FConnectionDrawingPolicy_InworldGraph::DrawPreviewConnector(const FGeometry& PinGeometry, const FVector2f& StartPoint, const FVector2f& EndPoint, UEdGraphPin* Pin)
{
	FConnectionParams Params;
	DetermineWiringStyle(Pin, nullptr, /*inout*/ Params);

	if (Pin->Direction == EEdGraphPinDirection::EGPD_Output)
	{
		DrawSplineWithArrow(FGeometryHelper::FindClosestPointOnGeom(PinGeometry, EndPoint), FVector2D(EndPoint), Params);
	}
	else
	{
		DrawSplineWithArrow(FGeometryHelper::FindClosestPointOnGeom(PinGeometry, StartPoint), FVector2D(StartPoint), Params);
	}
}
#endif

void FConnectionDrawingPolicy_InworldGraph::DrawSplineWithArrow(
	const FVector2D& StartAnchorPoint, const FVector2D& EndAnchorPoint, const FConnectionParams& Params)
{
	// bUserFlag1 indicates that we need to reverse the direction of connection (used by debugger)
	const FVector2D& P0 = Params.bUserFlag1 ? EndAnchorPoint : StartAnchorPoint;
	const FVector2D& P1 = Params.bUserFlag1 ? StartAnchorPoint : EndAnchorPoint;

	Internal_DrawLineWithArrow(P0, P1, Params);
}

void FConnectionDrawingPolicy_InworldGraph::Internal_DrawLineWithArrow(
	const FVector2D& StartAnchorPoint, const FVector2D& EndAnchorPoint, const FConnectionParams& Params)
{
	const FVector2D DeltaPos = EndAnchorPoint - StartAnchorPoint;
	const FVector2D UnitDelta = DeltaPos.GetSafeNormal();

	const FVector2D StartPoint = StartAnchorPoint;
	const FVector2D EndPoint = EndAnchorPoint - ArrowRadius.X * UnitDelta;

	// draw connection
	DrawConnection(WireLayerID, StartPoint, EndPoint, Params);

	// draw arrow
	const FVector2D ArrowDrawPos = EndPoint - ArrowRadius;
	const float AngleInRadians = FMath::Atan2(DeltaPos.Y, DeltaPos.X);

	FSlateDrawElement::MakeRotatedBox(DrawElementsList, ArrowLayerID,
		FPaintGeometry(ArrowDrawPos, ArrowImage->ImageSize * ZoomFactor, ZoomFactor), ArrowImage, ESlateDrawEffect::None,
		AngleInRadians, TOptional<FVector2D>(), FSlateDrawElement::RelativeToElement, Params.WireColor);
}

void FConnectionDrawingPolicy_InworldGraph::DrawSplineWithArrow(
	const FGeometry& StartGeom, const FGeometry& EndGeom, const FConnectionParams& Params)
{
	// get a reasonable seed point (halfway between the boxes)
	const FVector2D StartCenter = FGeometryHelper::CenterOf(StartGeom);
	const FVector2D EndCenter = FGeometryHelper::CenterOf(EndGeom);
	const FVector2D SeedPoint = (StartCenter + EndCenter) * 0.5f;

	// get sizes of pins widgets
	const FVector2D StartSize = StartGeom.GetDrawSize();
	const FVector2D EndSize = EndGeom.GetDrawSize();

	// calculates start/end points with offset related of center of widget
	// const FVector2D StartPoint = FVector2D(StartCenter.X, StartCenter.Y + StartSize.Y * 0.6f);
	const FVector2D StartPoint = FGeometryHelper::FindClosestPointOnGeom(StartGeom, SeedPoint);
	const FVector2D EndPoint = FVector2D(EndCenter.X, EndCenter.Y - EndSize.Y * 0.5f);

	DrawSplineWithArrow(StartPoint, EndPoint, Params);
}

FVector2D FConnectionDrawingPolicy_InworldGraph::ComputeSplineTangent(const FVector2D& Start, const FVector2D& End) const
{
	const FVector2D Delta = End - Start;
	const FVector2D NormDelta = Delta.GetSafeNormal();

	return NormDelta;
}

void FConnectionDrawingPolicy_InworldGraph::DetermineLinkGeometry(FArrangedChildren& ArrangedNodes,
	TSharedRef<SWidget>& OutputPinWidget, UEdGraphPin* OutputPin, UEdGraphPin* InputPin, FArrangedWidget*& StartWidgetGeometry,
	FArrangedWidget*& EndWidgetGeometry)
{
	if (UEdNode_InworldGraphEdge* EdgeNode = Cast<UEdNode_InworldGraphEdge>(InputPin->GetOwningNode()))
	{
		// get the pins that this edge connects
		UEdGraphPin* StartPin = EdgeNode->GetStartPin();
		UEdGraphPin* EndPin = EdgeNode->GetEndPin();

		if (StartPin && EndPin)
		{
			// finding the widgets for pins
			if (TSharedPtr<SGraphPin>* pStartWidget = PinToPinWidgetMap.Find(StartPin))
			{
				TSharedRef<SGraphPin> StartWidget = (*pStartWidget).ToSharedRef();
				StartWidgetGeometry = PinGeometries->Find(StartWidget);
			}

			if (TSharedPtr<SGraphPin>* pEndWidget = PinToPinWidgetMap.Find(EndPin))
			{
				TSharedRef<SGraphPin> EndWidget = (*pEndWidget).ToSharedRef();
				EndWidgetGeometry = PinGeometries->Find(EndWidget);
			}
		}
		else
		{
			// solution if exact pins are not found
			UEdNode_InworldGraphNode* Start = EdgeNode->GetStartNode();
			UEdNode_InworldGraphNode* End = EdgeNode->GetEndNode();

			if (Start && End)
			{
				int32* StartNodeIndex = NodeWidgetMap.Find(Start);
				int32* EndNodeIndex = NodeWidgetMap.Find(End);

				if (StartNodeIndex && EndNodeIndex)
				{
					StartWidgetGeometry = &(ArrangedNodes[*StartNodeIndex]);
					EndWidgetGeometry = &(ArrangedNodes[*EndNodeIndex]);
				}
			}
		}
	}
	else
	{
		// standard processing for normal pin connections
		StartWidgetGeometry = PinGeometries->Find(OutputPinWidget);

		if (TSharedPtr<SGraphPin>* pTargetWidget = PinToPinWidgetMap.Find(InputPin))
		{
			TSharedRef<SGraphPin> InputWidget = (*pTargetWidget).ToSharedRef();
			EndWidgetGeometry = PinGeometries->Find(InputWidget);
		}
	}
}

void FConnectionDrawingPolicy_InworldGraph::DrawConnection(
	int32 LayerId, const FVector2D& Start, const FVector2D& End, const FConnectionParams& Params)
{
	const FVector2D& P0 = Start;
	const FVector2D& P1 = End;

	const FVector2D SplineTangent = ComputeSplineTangent(P0, P1);
	const FVector2D P0Tangent = (Params.StartDirection == EGPD_Output) ? SplineTangent : -SplineTangent;
	const FVector2D P1Tangent = (Params.EndDirection == EGPD_Input) ? SplineTangent : -SplineTangent;

	if (Settings->bTreatSplinesLikePins)
	{
		// Distance to consider as an overlap
		const float QueryDistanceTriggerThresholdSquared =
			FMath::Square(Settings->SplineHoverTolerance + Params.WireThickness * 0.5f);

		// Distance to pass the bounding box cull test. This is used for the bCloseToSpline output that can be used as a
		// dead zone to avoid mistakes caused by missing a double-click on a connection.
		const float QueryDistanceForCloseSquared =
			FMath::Square(FMath::Sqrt(QueryDistanceTriggerThresholdSquared) + Settings->SplineCloseTolerance);

		bool bCloseToSpline = false;
		{
			// The curve will include the endpoints but can extend out of a tight bounds because of the tangents
			// P0Tangent coefficient maximizes to 4/27 at a=1/3, and P1Tangent minimizes to -4/27 at a=2/3.
			const float MaximumTangentContribution = 4.0f / 27.0f;
			FBox2D Bounds(ForceInit);

			Bounds += FVector2D(P0);
			Bounds += FVector2D(P0 + MaximumTangentContribution * P0Tangent);
			Bounds += FVector2D(P1);
			Bounds += FVector2D(P1 - MaximumTangentContribution * P1Tangent);

			bCloseToSpline = Bounds.ComputeSquaredDistanceToPoint(LocalMousePosition) < QueryDistanceForCloseSquared;

			// Draw the bounding box for debugging
#if 0
#define DrawSpaceLine(Point1, Point2, DebugWireColor) \
	{                                                 \
		const FVector2D FakeTangent = (Point2 - Point1).GetSafeNormal();
FSlateDrawElement::MakeDrawSpaceSpline(DrawElementsList, LayerId, Point1, FakeTangent, Point2, FakeTangent, ClippingRect, 1.0f,
ESlateDrawEffect::None, DebugWireColor); }

			if (bCloseToSpline)
			{
				const FLinearColor BoundsWireColor = bCloseToSpline ? FLinearColor::Green : FLinearColor::White;

				FVector2D TL = Bounds.Min;
				FVector2D BR = Bounds.Max;
				FVector2D TR = FVector2D(Bounds.Max.X, Bounds.Min.Y);
				FVector2D BL = FVector2D(Bounds.Min.X, Bounds.Max.Y);

				DrawSpaceLine(TL, TR, BoundsWireColor);
				DrawSpaceLine(TR, BR, BoundsWireColor);
				DrawSpaceLine(BR, BL, BoundsWireColor);
				DrawSpaceLine(BL, TL, BoundsWireColor);
			}
#endif
		}

		if (bCloseToSpline)
		{
			// Find the closest approach to the spline
			FVector2D ClosestPoint(ForceInit);
			float ClosestDistanceSquared = FLT_MAX;

			const int32 NumStepsToTest = 16;
			const float StepInterval = 1.0f / (float) NumStepsToTest;
			FVector2D Point1 = FMath::CubicInterp(P0, P0Tangent, P1, P1Tangent, 0.0f);
			for (float TestAlpha = 0.0f; TestAlpha < 1.0f; TestAlpha += StepInterval)
			{
				const FVector2D Point2 = FMath::CubicInterp(P0, P0Tangent, P1, P1Tangent, TestAlpha + StepInterval);
				const FVector2D LocalMousePositionVector(LocalMousePosition.X, LocalMousePosition.Y);
				const FVector2D ClosestPointToSegment = FMath::ClosestPointOnSegment2D(LocalMousePositionVector, Point1, Point2);
				const float DistanceSquared = (LocalMousePosition - ClosestPointToSegment).SizeSquared();

				if (DistanceSquared < ClosestDistanceSquared)
				{
					ClosestDistanceSquared = DistanceSquared;
					ClosestPoint = ClosestPointToSegment;
				}

				Point1 = Point2;
			}

			// Record the overlap
			if (ClosestDistanceSquared < QueryDistanceTriggerThresholdSquared)
			{
				if (ClosestDistanceSquared < SplineOverlapResult.GetDistanceSquared())
				{
					const float SquaredDistToPin1 = (Params.AssociatedPin1 != nullptr) ? (P0 - ClosestPoint).SizeSquared() : FLT_MAX;
					const float SquaredDistToPin2 = (Params.AssociatedPin2 != nullptr) ? (P1 - ClosestPoint).SizeSquared() : FLT_MAX;

					SplineOverlapResult = FGraphSplineOverlapResult(Params.AssociatedPin1, Params.AssociatedPin2,
						ClosestDistanceSquared, SquaredDistToPin1, SquaredDistToPin2, true);
				}
			}
			else if (ClosestDistanceSquared < QueryDistanceForCloseSquared)
			{
				SplineOverlapResult.SetCloseToSpline(true);
			}
		}
	}

	// way to mark that line is Dashed is to use flag bUserFlag2
	if (!Params.bUserFlag2)
	{
		// Draw a dashed line

		// Parameters for the dashed line
		const float DashSize = 8.0f;	// dash size
		const float GapSize = 6.0f;		// gap size

		// total line length
		const float PathLength = FVector2D::Distance(Start, End);

		// how many lines and spaces should be drawn
		const int32 NumSegments = FMath::CeilToInt(PathLength / (DashSize + GapSize));

		if (NumSegments <= 0)
		{
			// Draw the spline itself
			FSlateDrawElement::MakeDrawSpaceSpline(DrawElementsList, LayerId, P0, P0Tangent, P1, P1Tangent, Params.WireThickness, ESlateDrawEffect::None, Params.WireColor);
			return;
		}

		// step along the path
		const float SegmentLength = PathLength / static_cast<float>(NumSegments);

		// draw each line separately
		for (int32 i = 0; i < NumSegments; ++i)
		{
			// start position of the current line
			const float StartT = i * SegmentLength / PathLength;
			// position of the end of the current dash(limited by the size of the dash)
			const float EndT = FMath::Min((i * SegmentLength + DashSize) / PathLength, 1.0f);

			// interpolate the positions for the start and end of the dash
			const FVector2D DashStart = FMath::Lerp(Start, End, StartT);
			const FVector2D DashEnd = FMath::Lerp(Start, End, EndT);

			// draw a dash
			FSlateDrawElement::MakeDrawSpaceSpline(DrawElementsList, LayerId, DashStart, P0Tangent * (EndT - StartT), DashEnd,
				P1Tangent * (EndT - StartT), Params.WireThickness, ESlateDrawEffect::None, Params.WireColor);
		}
	}
	else
	{
		// draw a regular line
		FSlateDrawElement::MakeDrawSpaceSpline(DrawElementsList, LayerId, P0, P0Tangent, P1, P1Tangent, Params.WireThickness, ESlateDrawEffect::None, Params.WireColor);
	}

	if (Params.bDrawBubbles || (MidpointImage != nullptr))
	{
		// This table maps distance along curve to alpha
		FInterpCurve<float> SplineReparamTable;
		const float SplineLength = MakeSplineReparamTable(P0, P0Tangent, P1, P1Tangent, SplineReparamTable);

		// Draw bubbles on the spline
		if (Params.bDrawBubbles)
		{
			const float BubbleSpacing = 64.f * ZoomFactor;
			const float BubbleSpeed = 192.f * ZoomFactor;
			const FVector2D BubbleSize = BubbleImage->ImageSize * ZoomFactor * 0.2f * Params.WireThickness;

			float Time = (FPlatformTime::Seconds() - GStartTime);
			const float BubbleOffset = FMath::Fmod(Time * BubbleSpeed, BubbleSpacing);
			const int32 NumBubbles = FMath::CeilToInt(SplineLength / BubbleSpacing);
			for (int32 i = 0; i < NumBubbles; ++i)
			{
				const float Distance = ((float) i * BubbleSpacing) + BubbleOffset;
				if (Distance < SplineLength)
				{
					const float Alpha = SplineReparamTable.Eval(Distance, 0.f);
					FVector2D BubblePos = FMath::CubicInterp(P0, P0Tangent, P1, P1Tangent, Alpha);
					BubblePos -= (BubbleSize * 0.5f);

					FSlateDrawElement::MakeBox(DrawElementsList, LayerId, FPaintGeometry(BubblePos, BubbleSize, ZoomFactor), BubbleImage, ESlateDrawEffect::None, Params.WireColor);
				}
			}
		}

		// Draw the midpoint image
		if (MidpointImage != nullptr)
		{
			// Determine the spline position for the midpoint
			const float MidpointAlpha = SplineReparamTable.Eval(SplineLength * 0.5f, 0.f);
			const FVector2D Midpoint = FMath::CubicInterp(P0, P0Tangent, P1, P1Tangent, MidpointAlpha);

			// Approximate the slope at the midpoint (to orient the midpoint image to the spline)
			const FVector2D MidpointPlusE = FMath::CubicInterp(P0, P0Tangent, P1, P1Tangent, MidpointAlpha + KINDA_SMALL_NUMBER);
			const FVector2D MidpointMinusE = FMath::CubicInterp(P0, P0Tangent, P1, P1Tangent, MidpointAlpha - KINDA_SMALL_NUMBER);
			const FVector2D SlopeUnnormalized = MidpointPlusE - MidpointMinusE;

			// Draw the arrow
			const FVector2D MidpointDrawPos = Midpoint - MidpointRadius;
			const float AngleInRadians =
				SlopeUnnormalized.IsNearlyZero() ? 0.0f : FMath::Atan2(SlopeUnnormalized.Y, SlopeUnnormalized.X);

			FSlateDrawElement::MakeRotatedBox(DrawElementsList, LayerId,
				FPaintGeometry(MidpointDrawPos, MidpointImage->ImageSize * ZoomFactor, ZoomFactor), MidpointImage,
				ESlateDrawEffect::None, AngleInRadians, TOptional<FVector2D>(), FSlateDrawElement::RelativeToElement,
				Params.WireColor);
		}
	}
}
