/**
 * Copyright 2022-2025 Theai, Inc. dba Inworld AI
 */
#include "Graph/Nodes/InworldNode_SubgraphNode.h"

#include "Dom/JsonObject.h"
#include "Graph/Assets/InworldGraphAsset.h"
#include "Graph/InworldGraph.h"
#include "Serialization/JsonReader.h"
#include "Serialization/JsonSerializer.h"

void UInworldNode_SubgraphNode::SetGraphAsset(UInworldGraphAsset* Asset)
{
	GraphAsset = Asset;
	GraphAssetObject = Asset;
	if (Asset != nullptr)
	{
		NodeName = GraphAsset->GetName();
	}
	else
	{
		NodeName = {};
	}
}

void UInworldNode_SubgraphNode::RefreshDetails()
{
	if (GetGraphAsset() != nullptr)
	{
		Description = GetGraphAsset()->Description;
	}
	else
	{
		Description = {};
	}
}

const FString& UInworldNode_SubgraphNode::GetDescription()
{
	if (GetGraphAsset() != nullptr)
	{
		return GetGraphAsset()->Description;
	}
	return Description;
}

const UInworldNode* UInworldNode_SubgraphNode::GetStartNode()
{
	UInworldGraphAsset* SubgraphAsset = GetGraphAsset();
	if (SubgraphAsset && SubgraphAsset->IsValidAsSubgraph())
	{
		return SubgraphAsset->RuntimeGraph->GetNodeByName(SubgraphAsset->RuntimeGraph->StartNodes[0]);
	}
	return nullptr;
}

const UInworldNode* UInworldNode_SubgraphNode::GetEndNode()
{
	UInworldGraphAsset* SubgraphAsset = GetGraphAsset();
	if (SubgraphAsset && SubgraphAsset->IsValidAsSubgraph())
	{
		return SubgraphAsset->RuntimeGraph->GetNodeByName(SubgraphAsset->RuntimeGraph->EndNodes[0]);
	}
	return nullptr;
}

UInworldGraphAsset* UInworldNode_SubgraphNode::GetGraphAsset()
{
	if (GraphAssetObject == nullptr)
	{
		GraphAssetObject = GraphAsset.LoadSynchronous();
	}
	return GraphAssetObject;
}

EInworldRuntimeErrors UInworldNode_SubgraphNode::GetJsonConfig(const TSharedRef<FJsonObject>& GraphJson)
{
	if (GetGraphAsset() == nullptr || !GetGraphAsset()->RuntimeGraph)
	{
		return EInworldRuntimeErrors::SUBGRAPH_ASSET_NOTVALID;
	}

	if (!GetGraphAsset()->IsValidAsSubgraph())
	{
		return EInworldRuntimeErrors::SUBGRAPH_ONE_STARTEND;
	}

	FString SubgraphId = GraphAssetObject.GetPath();
	TSharedPtr<FJsonObject> Main = GraphJson->GetObjectField(TEXT("main"));

	// Parse subgraph config  from the asset
	TSharedPtr<FJsonObject> SubgraphJson;
	{
		TSharedRef<TJsonReader<>> JsonReader = TJsonReaderFactory<>::Create(GetGraphAsset()->GetGraphConfig());
		if (!FJsonSerializer::Deserialize(JsonReader, SubgraphJson) || !SubgraphJson.IsValid())
		{
			return EInworldRuntimeErrors::CANT_PARSE_JSON;
		}
	}

	auto BuildIdSet = [](const TArray<TSharedPtr<FJsonValue>>& ArrayValues)
	{
		TSet<FString> Result;
		for (const TSharedPtr<FJsonValue>& Value : ArrayValues)
		{
			const TSharedPtr<FJsonObject>* ObjectPtr = nullptr;
			if (Value.IsValid() && Value->TryGetObject(ObjectPtr) && ObjectPtr && (*ObjectPtr).IsValid())
			{
				FString Id;
				if ((*ObjectPtr)->TryGetStringField(TEXT("id"), Id))
				{
					Result.Add(MoveTemp(Id));
				}
			}
		}
		return Result;
	};

	// Merge components by unique id
	{
		const TArray<TSharedPtr<FJsonValue>>& GraphComponents = GraphJson->GetArrayField(TEXT("components"));
		TSet<FString> GraphComponentIds = BuildIdSet(GraphComponents);
		TArray<TSharedPtr<FJsonValue>> MergedComponents = GraphComponents;

		const TArray<TSharedPtr<FJsonValue>>& SubgraphComponents = SubgraphJson->GetArrayField(TEXT("components"));
		for (const TSharedPtr<FJsonValue>& SubgraphComponent : SubgraphComponents)
		{
			TSharedPtr<FJsonObject> SubgraphComponentObject = SubgraphComponent->AsObject();
			if (SubgraphComponentObject.IsValid())
			{
				FString Id = SubgraphComponentObject->GetStringField(TEXT("id"));
				if (!GraphComponentIds.Contains(Id))
				{
					MergedComponents.Add(SubgraphComponent);
					GraphComponentIds.Add(MoveTemp(Id));
				}
			}
		}

		if (MergedComponents.Num() > 0)
		{
			GraphJson->SetArrayField(TEXT("components"), MergedComponents);
		}
	}

	// Merge nested subgraphs by unique id
	{
		const TArray<TSharedPtr<FJsonValue>>& GraphSubgraphs = GraphJson->GetArrayField(TEXT("subgraphs"));
		TSet<FString> GraphSubgraphIds = BuildIdSet(GraphSubgraphs);
		TArray<TSharedPtr<FJsonValue>> MergedSubgraphs = GraphSubgraphs;

		const TArray<TSharedPtr<FJsonValue>>& SubgraphSubgraphs = SubgraphJson->GetArrayField(TEXT("subgraphs"));
		for (const TSharedPtr<FJsonValue>& SubgraphSubgraph : SubgraphSubgraphs)
		{
			TSharedPtr<FJsonObject> SubgraphSubgraphObject = SubgraphSubgraph->AsObject();
			if (SubgraphSubgraphObject.IsValid())
			{
				FString Id = SubgraphSubgraphObject->GetStringField(TEXT("id"));
				if (!GraphSubgraphIds.Contains(Id))
				{
					MergedSubgraphs.Add(SubgraphSubgraph);
					GraphSubgraphIds.Add(MoveTemp(Id));
				}
			}
		}

		// Convert subgraph main into a subgraph
		TSharedPtr<FJsonObject> SubgraphMain = SubgraphJson->GetObjectField(TEXT("main"));
		if (!GraphSubgraphIds.Contains(SubgraphId))
		{
			TSharedPtr<FJsonObject> SubgraphJsonConverted = MakeShared<FJsonObject>();
			SubgraphJsonConverted->SetStringField(TEXT("id"), SubgraphId);

			const TArray<TSharedPtr<FJsonValue>>& Nodes = SubgraphMain->GetArrayField(TEXT("nodes"));
			SubgraphJsonConverted->SetArrayField(TEXT("nodes"), Nodes);

			const TArray<TSharedPtr<FJsonValue>>& Edges = SubgraphMain->GetArrayField(TEXT("edges"));
			SubgraphJsonConverted->SetArrayField(TEXT("edges"), Edges);

			const TArray<TSharedPtr<FJsonValue>>& StartNodes = SubgraphMain->GetArrayField(TEXT("start_nodes"));
			SubgraphJsonConverted->SetArrayField(TEXT("start_nodes"), StartNodes);

			const TArray<TSharedPtr<FJsonValue>>& EndNodes = SubgraphMain->GetArrayField(TEXT("end_nodes"));
			SubgraphJsonConverted->SetArrayField(TEXT("end_nodes"), EndNodes);

			MergedSubgraphs.Add(MakeShared<FJsonValueObject>(SubgraphJsonConverted));
		}

		if (MergedSubgraphs.Num() > 0)
		{
			GraphJson->SetArrayField(TEXT("subgraphs"), MergedSubgraphs);
		}
	}

	TSharedPtr<FJsonObject> NodeObj = MakeShared<FJsonObject>();
	NodeObj->SetStringField(TEXT("id"), GetNodeName());
	NodeObj->SetStringField(TEXT("type"), TEXT("SubgraphNode"));
	NodeObj->SetStringField(TEXT("subgraph_id"), SubgraphId);

	TArray<TSharedPtr<FJsonValue>> Nodes = Main->GetArrayField(TEXT("nodes"));
	Nodes.Add(MakeShared<FJsonValueObject>(NodeObj));
	Main->SetArrayField(TEXT("nodes"), Nodes);

	return EInworldRuntimeErrors::SUCCESS;
}
