// Copyright (c) 2023 Codemongous LLC. All Rights Reserved.

#include "UI/SourcifyGraph/EdGraph_Sourcify.h"

#include "Async/Async.h"
#include "AssetRegistry/IAssetRegistry.h"
#include "AssetThumbnail.h"
#include "EdGraph/EdGraphPin.h"
#include "Engine/Blueprint.h"
#include "Misc/FilterCollection.h"
#include "Misc/IFilter.h"
#include "SSourcifyGraph.h"
#include "SourcifyCoreUEWrapperModule.h"
#include "UI/SourcifyGraph/EdGraphNode_Sourcify.h"

#include UE_INLINE_GENERATED_CPP_BY_NAME(EdGraph_Sourcify)

using namespace Sourcify;

FSourcifyNodeInfo::FSourcifyNodeInfo(Sourcify::DependencyNode::ReferenceInfo InInfo, bool InbReferencers)
	: ReferenceInfo(InInfo)
	, bReferencers(InbReferencers)
	, OverflowCount(0)
	, bExpandAllChildren(true)
	, ChildProvisionSize(1)
	, PassedFilters(true)
{
	Parents.Reserve(4);
}

bool FSourcifyNodeInfo::IsFirstParent(const FAssetIdentifier& InParentId) const
{
	return Parents.IsEmpty() || FAssetIdentifier{ Parents[0].ReferenceInfo.Node->AssetIdentifier.c_str() } == InParentId;
}

bool FSourcifyNodeInfo::IsADuplicate() const
{
	return Parents.Num() > 1;
}

int32 FSourcifyNodeInfo::ProvisionSize(const FAssetIdentifier& InParentId) const
{
	return IsFirstParent(InParentId) ? ChildProvisionSize : 1;
}

UEdGraph_Sourcify::UEdGraph_Sourcify()
{
	if (!IsTemplate())
	{
		AssetThumbnailPool = MakeShareable( new FAssetThumbnailPool(1024) );
	}

	//Settings = GetMutableDefault<UReferenceViewerSettings>();
}

void UEdGraph_Sourcify::BeginDestroy()
{
	AssetThumbnailPool.Reset();

	Super::BeginDestroy();
}

void UEdGraph_Sourcify::SetGraphRoot(std::shared_ptr<Sourcify::DependencyNode> GraphRootNode, const FIntPoint& GraphRootOrigin /*= FIntPoint(ForceInitToZero)*/)
{
	CurrentDependencyRootNode = GraphRootNode;
	CurrentGraphRootOrigin = GraphRootOrigin;
}

void UEdGraph_Sourcify::SetSourcifyGraph(TSharedPtr<SSourcifyGraph> InViewer)
{
	SourcifyGraph = InViewer;
}

UEdGraphNode_Sourcify* UEdGraph_Sourcify::RebuildGraph()
{
	RemoveAllNodes();
	UEdGraphNode_Sourcify* NewRootNode = ConstructNodes(CurrentDependencyRootNode, CurrentGraphRootOrigin);
	NotifyGraphChanged();
	
	return NewRootNode;
}

UEdGraphNode_Sourcify* UEdGraph_Sourcify::ConstructNodes(std::shared_ptr<Sourcify::DependencyNode> InGraphRoot, const FIntPoint& GraphRootOrigin)
{
	RootNodeInfo = FSourcifyNodeInfo(Sourcify::DependencyNode::ReferenceInfo{ FSourcifyModule::Get().GetIUE(), InGraphRoot, nullptr}, true);
	RecursivelyPopulateNodeInfos(true, RootNodeInfo.ReferenceInfo, RootNodeInfo, 0, 100/*Settings->GetSearchReferencerDepthLimit()*/);

	return BuildNodes();
}

UEdGraphNode_Sourcify* UEdGraph_Sourcify::BuildNodes()
{
	RemoveAllNodes();
	UEdGraphNode_Sourcify* RootNode = nullptr;

	bBreadthLimitReached = false;
	if (CurrentDependencyRootNode)
	{
		// Create the root node
		bool bRootIsDuplicated = false;

		RootNode = CreateReferenceNode();
		RootNode->SetupReferenceNode(RootNodeInfo.ReferenceInfo, CurrentGraphRootOrigin, { FAssetIdentifier{CurrentDependencyRootNode->AssetIdentifier.c_str()} }, RootNodeInfo.AssetData, /*bInAllowThumbnail = */ true/*!Settings->IsCompactMode()*/, /*bIsDuplicate*/ bRootIsDuplicated);
		RootNode->SetMakeCommentBubbleVisible(false/*Settings->IsShowPath()*/);
		RecursivelyCreateNodes(false, CurrentGraphRootOrigin, FAssetIdentifier{ CurrentDependencyRootNode->AssetIdentifier.c_str() }, RootNode, RootNodeInfo, 0, 100/*Settings->GetSearchReferencerDepthLimit()*/, /*bIsRoot*/ true);
	}

	NotifyGraphChanged();
	return RootNode;
}

bool UEdGraph_Sourcify::AreAllNodesReadyToBeSourcified() const
{
	TMap<FSoftObjectPath, NodeErrorStatus> NodeMap;
	auto BuildReadyMap = [this, &NodeMap](auto&& BuildReadyMap, const FSourcifyNodeInfo& NodeInfo) -> void
		{
			if (NodeInfo.ReferenceInfo.Node)
			{
				NodeMap.Add({ NodeInfo.AssetData.GetSoftObjectPath(), NodeInfo.ReferenceInfo.Node->IsReadyForOutput() });
			}
			for (const auto& Child : NodeInfo.Children)
			{
				BuildReadyMap(BuildReadyMap, Child.Key);
			}
		};

	BuildReadyMap(BuildReadyMap, RootNodeInfo);

	// If the Sourcify Graph opens automatically at startup, the graph could be empty
	if (NodeMap.IsEmpty())
	{
		return false;
	}

	for (const auto& Node : NodeMap)
	{
		if (Node.Value == NodeErrorStatus::Fatal)
		{
			return false;
		}
	}

	return true;
}

void UEdGraph_Sourcify::RecursivelyPopulateNodeInfos(bool bReferencers, Sourcify::DependencyNode::ReferenceInfo ReferenceInfo, FSourcifyNodeInfo& NodeInfo, int32 CurrentDepth, int32 MaxDepth)
{
	if (ReferenceInfo.Node)
	{
		NodeInfo.ReferenceInfo = ReferenceInfo;

		for (int32 i = 0; i < ReferenceInfo.Node->Referencees.size(); i++)
		{
			auto& info = NodeInfo.Children.Add_GetRef({ FSourcifyNodeInfo{ReferenceInfo.Node->Referencees[i], true}, EDependencyPinCategory::LinkTypeHard });
			RecursivelyPopulateNodeInfos(true, ReferenceInfo.Node->Referencees[i], info.Key, CurrentDepth + 1, MaxDepth);
		}

		IAssetRegistry* AssetRegistry = IAssetRegistry::Get();
		if (!AssetRegistry)
		{
			return;
		}
		TArray<FAssetData> AssetData;
		AssetRegistry->GetAssetsByPackageName(ReferenceInfo.Node->AssetIdentifier.c_str(), AssetData);
		check(AssetData.Num() == 1);
		NodeInfo.AssetData = AssetData[0];
	}
}

void UEdGraph_Sourcify::GatherAssetData(TMap<FAssetIdentifier, FSourcifyNodeInfo>& InNodeInfos)
{
	// Grab the list of packages
	TSet<FName> PackageNames;
	for (TPair<FAssetIdentifier, FSourcifyNodeInfo>&  InfoPair : InNodeInfos)
	{
		FAssetIdentifier& AssetId = InfoPair.Key;
		if (!AssetId.IsValue() && !AssetId.PackageName.IsNone())
		{
			PackageNames.Add(AssetId.PackageName);
		}
	}

	// Retrieve the AssetData from the Registry
	TMap<FName, FAssetData> PackagesToAssetDataMap;
	UE::AssetRegistry::GetAssetForPackages(PackageNames.Array(), PackagesToAssetDataMap);

	// Populate the AssetData back into the NodeInfos
	for (TPair<FAssetIdentifier, FSourcifyNodeInfo>&  InfoPair : InNodeInfos)
	{
		InfoPair.Value.AssetData = PackagesToAssetDataMap.FindRef(InfoPair.Key.PackageName);
	}
}

UEdGraphNode_Sourcify* UEdGraph_Sourcify::RecursivelyCreateNodes(bool bInReferencers, const FIntPoint& InNodeLoc, const FAssetIdentifier& InParentId, UEdGraphNode_Sourcify* InParentNode, const FSourcifyNodeInfo& InNodeInfo, int32 InCurrentDepth, int32 InMaxDepth, bool bIsRoot)
{
	const FSourcifyNodeInfo& NodeInfo = InNodeInfo;
	int32 NodeProvSize = 1;

	UEdGraphNode_Sourcify* NewNode = nullptr;
	if (bIsRoot)
	{
		NewNode = InParentNode;
		NodeProvSize = NodeInfo.ProvisionSize(FAssetIdentifier(NAME_None));
	}
	else
	{
		NewNode = CreateReferenceNode();
		NewNode->SetupReferenceNode(NodeInfo.ReferenceInfo, InNodeLoc, { FAssetIdentifier{ NodeInfo.ReferenceInfo.Node->AssetIdentifier.c_str() } }, NodeInfo.AssetData, /*bInAllowThumbnail*/ /*!Settings->IsCompactMode() &&*/ NodeInfo.PassedFilters, /*bIsADuplicate*/ NodeInfo.Parents.Num() > 1);
		NewNode->SetIsFiltered(!NodeInfo.PassedFilters);
		NodeProvSize = NodeInfo.ProvisionSize(InParentId);
	}

	// Add compilation 
	if (InNodeInfo.ReferenceInfo.Node->Type == Sourcify::DependencyNode::NodeType::Blueprint)
	{
		if (UBlueprint* BP = Cast<UBlueprint>(FSourcifyCoreUEWrapper::Unwrap(InNodeInfo.ReferenceInfo.Node->Object)))
		{
			if (!BP->OnCompiled().IsBoundToObject(this))
			{
				BP->OnCompiled().AddUObject(this, &UEdGraph_Sourcify::OnBPCompiled);	
			}
		}
	}

	const bool bIsFirstOccurance = bIsRoot || NodeInfo.IsFirstParent(InParentId);
	FIntPoint ChildLoc = InNodeLoc;
	if (bIsFirstOccurance) // Only expand the first parent
	{
		// position the children nodes
		const int32 ColumnWidth = /*Settings->IsCompactMode() ? 500 :*/ 500;
		ChildLoc.X += bInReferencers ? -ColumnWidth : ColumnWidth;

		int32 NodeSizeY = /*Settings->IsCompactMode() ? 100 :*/ 100;
		//NodeSizeY += /*Settings->IsShowPath() ? 40 :*/ 0;
		
		// Space out the nodes vertically by doubling their reported size and some buffer
		NodeSizeY += /*Settings->IsShowPath() ? 40 :*/ 120;

		ChildLoc.Y -= (NodeProvSize - 1) * NodeSizeY * 0.5 ;

		int32 Breadth = 0;
		int32 ChildIdx = 0;
		for (; ChildIdx < InNodeInfo.Children.Num(); ChildIdx++)
		{
			const TPair<FSourcifyNodeInfo, EDependencyPinCategory>& Pair = InNodeInfo.Children[ChildIdx];
			if (!InNodeInfo.bExpandAllChildren)
			{
				break;
			}

			const FSourcifyNodeInfo& ChildId = Pair.Key;
			int32 ChildProvSize = 0;
			FAssetIdentifier AssetIdentifier = FAssetIdentifier{ NodeInfo.ReferenceInfo.Node->AssetIdentifier.c_str() };
			if (ChildId.IsFirstParent(AssetIdentifier))
			{
				ChildProvSize = ChildId.ProvisionSize(AssetIdentifier);
			}

			// The provision size will always be at least 1 if it should be shown, factoring in filters, breadth, duplicates, etc.
			if (ChildProvSize > 0)
			{
				ChildLoc.Y += (ChildProvSize - 1) * NodeSizeY * 0.5;

				UEdGraphNode_Sourcify* ChildNode = RecursivelyCreateNodes(bInReferencers, ChildLoc, AssetIdentifier, NewNode, ChildId, InCurrentDepth + 1, InMaxDepth);

				if (bInReferencers)
				{
					ChildNode->GetDependencyPin()->PinType.PinCategory = ::GetName(Pair.Value);
					NewNode->AddReferencer(ChildNode);
				}
				else
				{
					ChildNode->GetReferencerPin()->PinType.PinCategory = ::GetName(Pair.Value);
					ChildNode->AddReferencer(NewNode);
				}

				ChildLoc.Y += NodeSizeY * (ChildProvSize + 1) * 0.5;
				Breadth++;
			}
		}
	}

	return NewNode;
}

const TSharedPtr<FAssetThumbnailPool>& UEdGraph_Sourcify::GetAssetThumbnailPool() const
{
	return AssetThumbnailPool;
}

UEdGraphNode_Sourcify* UEdGraph_Sourcify::CreateReferenceNode()
{
	const bool bSelectNewNode = false;
	return Cast<UEdGraphNode_Sourcify>(CreateNode(UEdGraphNode_Sourcify::StaticClass(), bSelectNewNode));
}

void UEdGraph_Sourcify::RemoveAllNodes()
{
	TArray<UEdGraphNode*> NodesToRemove = Nodes;
	for (int32 NodeIndex = 0; NodeIndex < NodesToRemove.Num(); ++NodeIndex)
	{
		if (UEdGraphNode_Sourcify* SourcifyNode = Cast<UEdGraphNode_Sourcify>(NodesToRemove[NodeIndex]))
		{
			SourcifyNode->UnRegisterDelegates();
		}
		RemoveNode(NodesToRemove[NodeIndex]);
	}
}

void UEdGraph_Sourcify::OnBPCompiled(UBlueprint* BP) const
{
	if (bIsGraphRefreshPending)
	{
		return;
	}

	// The BP triggers this callback JUST BEFORE it's done compiling, not truly afterwards.
	// So if we don't do the full check below, ensures will trigger in the Kismet compiler
	// That's gross
	bIsGraphRefreshPending = true;
	TWeakObjectPtr<const UEdGraph_Sourcify> WeakThis = this;
	Async(EAsyncExecution::ThreadPool, [BP, WeakThis]
		{
			while (BP->bBeingCompiled)
			{
				FPlatformProcess::Sleep(0.1);
			}

			Async(EAsyncExecution::TaskGraphMainThread, [WeakThis]
				{
					FModuleManager::LoadModuleChecked<FSourcifyModule>("Sourcify").RefreshCurrentGraph();
					if (WeakThis.IsValid())
					{
						WeakThis->bIsGraphRefreshPending = false;
					}
				});
		});
}

