/**
 * Copyright 2022-2025 Theai, Inc. dba Inworld AI
 */

#include "Graph/InworldGraph.h"

#include "Async/Async.h"
#include "Engine/GameInstance.h"
#include "Engine/World.h"
#include "Graph/Assets/InworldGraphAsset.h"
#include "Graph/Edges/InworldEdge_WithCondition.h"
#include "Graph/InworldNode.h"
#include "Graph/Nodes/InworldNode_Custom.h"
#include "InworldRuntimeErrors.h"
#include "InworldRuntimeSettings.h"
#include "InworldRuntimeSubsystem.h"
#include "InworldRuntimeUtils.h"
#include "JsonObjectConverter.h"
#include "Misc/TVariant.h"
#include "Tasks/Task.h"

THIRD_PARTY_INCLUDES_START
#include "InworldRuntimeLibrary/inworld.h"
THIRD_PARTY_INCLUDES_END

DEFINE_LOG_CATEGORY(LogInworldGraph);

void UInworldGraph::BeginDestroy()
{
	if (CompileTask.IsValid() && !CompileTask.IsCompleted())
	{
		CompileTask.GetResult();
	}

	if (IsCompiled())
	{
		CancelAllExecutions();

		AsyncTask(ENamedThreads::AnyThread,
			[Executor = MoveTemp(executor_)]() mutable
			{
				if (Executor.IsValid())
				{
					Executor->Stop();
					Executor.Reset();
				}
			});

		for (const FString& Type : ComponentTypesRegistry)
		{
			if (UInworldRuntimeSubsystem* Subsystem = UInworldRuntimeSubsystem::Get())
			{
				Subsystem->UnregisterComponent(Type);
			}
		}
		for (UInworldEdge* Edge : CustomEdges)
		{
			Edge->DestroyEdge();
		}
		ComponentTypesRegistry.Empty();
		CustomEdges.Empty();
		UE_LOG(LogInworldGraph, Log, TEXT("UInworldGraph::BeginDestroy - %s"), *GraphId);
	}
	Super::BeginDestroy();
}

void UInworldGraph::GetGraphInstance(FOnGraphCompiled Callback)
{
	FOnGraphCompiledNative NativeCallback;
	if (Callback.IsBound())
	{
		NativeCallback.BindUFunction(Callback.GetUObject(), Callback.GetFunctionName());
	}
	GetGraphInstance(NativeCallback);
}

void UInworldGraph::GetGraphInstance(FOnGraphCompiledNative Callback)
{
	TArray<FString> Failed;
	if (!GenerateConfig(Failed))
	{
		UE_LOG(LogInworldGraph, Error, TEXT("UInworldGraph::GetGraphInstance - Graph JSON generating failed."));
		Callback.ExecuteIfBound(nullptr, false);
	}
	UInworldGraph* RegisteredGraph = UInworldRuntimeSubsystem::Get()->RegisterGraph(this);
	if (IsValid(RegisteredGraph))
	{
		RegisteredGraph->Compile(Callback);
	}
	else
	{
		Callback.ExecuteIfBound(nullptr, false);
	}
}

TOptional<FString> UInworldGraph::CompileInternal()
{
	if (bIsCompiled)
	{
		UE_LOG(LogInworldGraph, Warning, TEXT("UInworldGraph::CompileGraph - Graph already compiled"));
		return TOptional<FString>{};
	}

	inworld::StatusOr_CompiledGraphInterfacePtr status_or_compiled_graph;

	// Temporary fix for Extras path
	GraphConfig.ReplaceInline(TEXT("INWORLD_EXTRAS_PATH"), *GetInworldExtrasPath());

	TSharedPtr<FJsonObject> GraphJsonObject;
	TSharedRef<TJsonReader<>> JsonReader = TJsonReaderFactory<>::Create(GraphConfig);
	if (!FJsonSerializer::Deserialize(JsonReader, GraphJsonObject) || !GraphJsonObject.IsValid())
	{
		const FString Error{FString::Format(TEXT("Unable to parse json=[{0}]"), {*GraphConfig})};
		UE_LOG(LogInworldGraph, Error, TEXT("UInworldGraph::CompileGraph - %s"), *Error);
		return Error;
	}

	const TSharedPtr<FJsonObject>* MainJsonObject;
	if (!GraphJsonObject->TryGetObjectField(TEXT("main"), MainJsonObject))
	{
		const FString Error{FString::Format(TEXT("Unable to parse main section json=[{0}]"), {*GraphConfig})};
		UE_LOG(LogInworldGraph, Error, TEXT("UInworldGraph::CompileGraph - %s"), *Error);
		return Error;
	}

	auto RegisterNodes = [](const TSharedPtr<FJsonObject>& GraphJsonObj) -> bool
	{
		const TArray<TSharedPtr<FJsonValue>>* NodesJsonArray;
		if (GraphJsonObj->TryGetArrayField(TEXT("nodes"), NodesJsonArray))
		{
			for (TSharedPtr<FJsonValue> NodeJson : *NodesJsonArray)
			{
				const TSharedPtr<FJsonObject>* NodeJsonObject;
				if (NodeJson->TryGetObject(NodeJsonObject))
				{
					FString Type;
					if ((*NodeJsonObject)->TryGetStringField(TEXT("type"), Type))
					{
						UInworldRuntimeSubsystem::Get()->RegisterNodeType(Type);
					}
				}
			}
		}
		return true;
	};
	auto RegisterEdges = [this](const TSharedPtr<FJsonObject>& GraphJsonObj,
							 TMap<FString, FUnrealObjectCreationConfig>& EdgesCreationConfigsMap) -> TOptional<FString>
	{
		const TArray<TSharedPtr<FJsonValue>>* EdgesJsonObject;
		if (GraphJsonObj->TryGetArrayField(TEXT("edges"), EdgesJsonObject))
		{
			for (TSharedPtr<FJsonValue> EdgeJson : *EdgesJsonObject)
			{
				const TSharedPtr<FJsonObject>* EdgeJsonObject;
				if (EdgeJson->TryGetObject(EdgeJsonObject))
				{
					FString ConditionId;
					(*EdgeJsonObject)->TryGetStringField(TEXT("condition_id"), ConditionId);
					if (ConditionId.IsEmpty())
					{
						continue;
					}

					FUnrealObjectCreationConfig* EdgesCreationConfig = EdgesCreationConfigsMap.Find(ConditionId);
					if (!EdgesCreationConfig)
					{
						const FString Error{
							FString::Format(TEXT("Failed to find custom condition edge creation config: {0}"), {*ConditionId})};
						UE_LOG(LogInworldGraph, Error, TEXT("UInworldGraph::CompileGraph - %s"), *Error);
						return Error;
					}
					int32 PrefixLen = EdgesCreationConfig->Type.Find(ConditionId) + ConditionId.Len();
					if (UClass * Class{LoadObject<UClass>(nullptr, *EdgesCreationConfig->Type.RightChop(PrefixLen))};
						Class != nullptr)
					{
						auto* CustomEdge = NewObject<UInworldEdge_WithCondition>(GetTransientPackage(), Class);
						CustomEdge->InitializeEdge(EdgesCreationConfig->InstanceData);
						if (CustomEdge)
						{
							if (UInworldRuntimeSubsystem::Get()->RegisterCustomEdgeConditionType(
									EdgesCreationConfig->Type, CustomEdge))
							{
								ComponentTypesRegistry.Add(EdgesCreationConfig->Type);
								CustomEdges.Add(CustomEdge);
							}
						}
					}
					else
					{
						const FString Error{
							FString::Format(TEXT("Failed to find custom condition edge class: {0}"), {*EdgesCreationConfig->Type})};
						UE_LOG(LogInworldGraph, Error, TEXT("UInworldGraph::CompileGraph - %s"), *Error);
						return Error;
					}
				}
			}
		}
		return TOptional<FString>{};
	};

	const FString GraphJsonId = (*MainJsonObject)->GetStringField(TEXT("id"));

	// edges creation configs map
	TMap<FString, FUnrealObjectCreationConfig> EdgesCreationConfigsMap;
	{
		const TArray<TSharedPtr<FJsonValue>>* ComponentsJson;
		if (GraphJsonObject->TryGetArrayField(TEXT("components"), ComponentsJson))
		{
			for (TSharedPtr<FJsonValue> ComponentJson : *ComponentsJson)
			{
				const TSharedPtr<FJsonObject>* ComponentJsonObject;
				if (ComponentJson->TryGetObject(ComponentJsonObject))
				{
					FString Type;
					(*ComponentJsonObject)->TryGetStringField(TEXT("type"), Type);
					if (Type.StartsWith("UCE/"))
					{
						FUnrealObjectCreationConfig EdgesCreationConfig;
						(*ComponentJsonObject)->TryGetStringField(TEXT("id"), EdgesCreationConfig.Id);
						EdgesCreationConfig.Type = Type;
						const TSharedPtr<FJsonObject>* CreationConfig;
						if ((*ComponentJsonObject)->TryGetObjectField(TEXT("creation_config"), CreationConfig))
						{
							const TSharedPtr<FJsonObject>* Properties;
							if ((*CreationConfig)->TryGetObjectField(TEXT("properties"), Properties))
							{
								FString InstanceData;
								if ((*Properties)->TryGetStringField(TEXT("InstanceData"), InstanceData))
								{
									EdgesCreationConfig.InstanceData = InstanceData;
									EdgesCreationConfigsMap.Add(EdgesCreationConfig.Id, EdgesCreationConfig);
								}
							}
						}
					}
				}
			}
		}
	}

	// subgraphs
	const TArray<TSharedPtr<FJsonValue>>* SubgraphsJsonArray;
	if (GraphJsonObject->TryGetArrayField(TEXT("subgraphs"), SubgraphsJsonArray))
	{
		for (TSharedPtr<FJsonValue> SubgraphJson : *SubgraphsJsonArray)
		{
			const TSharedPtr<FJsonObject>* SubgraphJsonObj;
			if (SubgraphJson->TryGetObject(SubgraphJsonObj))
			{
				RegisterNodes(*SubgraphJsonObj);
				RegisterEdges(*SubgraphJsonObj, EdgesCreationConfigsMap);
			}
		}
	}

	// main
	RegisterNodes(*MainJsonObject);
	RegisterEdges(*MainJsonObject, EdgesCreationConfigsMap);

	if (auto status_or_config_parser = inworld::ConfigParser::inworld_inworld_graphs_ConfigParser_CreateConfigParser__SWIG();
		status_or_config_parser.ok())
	{
		status_or_compiled_graph = status_or_config_parser.value()->ParseGraphFromJsonString(TCHAR_TO_UTF8(*GraphConfig));
	}

	if (status_or_compiled_graph.ok())
	{
		if (inworld::StatusOr_GraphExecutorFactoryPtr status_or_graph_executor_factory{
				inworld::GraphExecutorFactory::inworld_inworld_graphs_GraphExecutorFactory_CreateHelper__SWIG_0()};
			status_or_graph_executor_factory.ok())
		{
			TSharedPtr<inworld::GraphExecutorFactory> graph_executor_factory{status_or_graph_executor_factory.value()};
			graph_executor_factory->swig_owns_self_ = true;
			TSharedPtr<inworld::ExecutorInterface> executor = TSharedPtr<inworld::ExecutorInterface>(
				graph_executor_factory->CreateExecutorHelper(status_or_compiled_graph.value()).value());
			executor->swig_owns_self_ = true;
			executor->Start();
			AsyncTask(ENamedThreads::GameThread,
				[WeakThis = TWeakObjectPtr<UInworldGraph>(this), executor]()
				{
					if (WeakThis.IsValid())
					{
						WeakThis->executor_ = executor;
						WeakThis->bIsCompiled = true;
					}
				});
			UE_LOG(LogInworldGraph, Log, TEXT("UInworldGraph::CompileGraph - Successfully compiled graph"));
			return TOptional<FString>{};
		}
		else
		{
			const FString Error{FString::Format(TEXT("Error compiling graph: {0}"),
				{FString{UTF8_TO_TCHAR(status_or_graph_executor_factory.status().ToString().c_str())}})};
			UE_LOG(LogInworldGraph, Error, TEXT("UInworldGraph::CompileGraph - %s"), *Error);
			return Error;
		}
	}
	else
	{
		const FString Error{FString::Format(
			TEXT("Error compiling graph: {0}"), {FString{UTF8_TO_TCHAR(status_or_compiled_graph.status().ToString().c_str())}})};
		UE_LOG(LogInworldGraph, Error, TEXT("UInworldGraph::CompileGraph - %s"), *Error);
		return Error;
	}
}

void UInworldGraph::Compile(FOnGraphCompiledNative Callback)
{
	if (bIsCompiled)
	{
		Callback.ExecuteIfBound(this, true);
		return;
	}

	{
		FScopeLock Lock(&CompileCallbacksLock);
		PendingCompileCallbacks.Add(Callback);
		if (bIsCompiling)
		{
			return;
		}
		bIsCompiling = true;
	}

	FOnGraphCompiledNative Coalesced;
	Coalesced.BindLambda(
		[this](UInworldGraph* Graph, bool bSuccess)
		{
			TArray<FOnGraphCompiledNative> Local;
			{
				FScopeLock Lock(&CompileCallbacksLock);
				Local = MoveTemp(PendingCompileCallbacks);
				PendingCompileCallbacks.Reset();
				bIsCompiling = false;
			}
			for (auto& Cb : Local)
			{
				Cb.ExecuteIfBound(Graph, bSuccess);
			}
		});

	CompileTask = UE::Tasks::Launch(UE_SOURCE_LOCATION,
		[WeakThis = TWeakObjectPtr<UInworldGraph>(this), Coalesced]() mutable
		{
			TOptional<FString> Error =
				WeakThis.IsValid() ? WeakThis->CompileInternal() : FString{TEXT("Inworld Graph Object no longer valid!")};

			AsyncTask(ENamedThreads::GameThread,
				[WeakThis, Coalesced, Error]()
				{
#if WITH_EDITOR
					if (Error.IsSet())
					{
						FMessageLog("PIE").Error(FText::FromString(Error.GetValue()));
					}
#endif
					if (WeakThis.IsValid())
					{
						Coalesced.ExecuteIfBound(WeakThis.Get(), !Error.IsSet());
					}
					else
					{
						Coalesced.ExecuteIfBound(nullptr, false);
					}
				});
		});
}

FString UInworldGraph::Execute(UObject* Instigator, const FInworldDataHandle& InputData,
	const TMap<FName, FInworldDataHandle>& RuntimeData, FOnGraphResult ResultCallback, FOnGraphStatus StatusCallback)
{
	FOnGraphResultNative NativeResultCallback;
	if (ResultCallback.IsBound())
	{
		NativeResultCallback.BindUFunction(ResultCallback.GetUObject(), ResultCallback.GetFunctionName());
	}

	FOnGraphStatusNative NativeStatusCallback;
	if (StatusCallback.IsBound())
	{
		NativeStatusCallback.BindUFunction(StatusCallback.GetUObject(), StatusCallback.GetFunctionName());
	}

	return Execute(Instigator, InputData, RuntimeData, NativeResultCallback, NativeStatusCallback);
}

FString UInworldGraph::Execute(UObject* Instigator, const FInworldDataHandle& InputData,
	const TMap<FName, FInworldDataHandle>& RuntimeData, FOnGraphResultNative ResultCallback, FOnGraphStatusNative StatusCallback)
{
	if (!bIsCompiled)
	{
		UE_LOG(LogInworldGraph, Error, TEXT("UInworldGraph::Execute - graph not compiled"));
		return {};
	}

	if (!InputData.IsValid())
	{
		UE_LOG(LogInworldGraph, Error, TEXT("UInworldGraph::Execute - input data is invalid"));
		return {};
	}

	if (!executor_.IsValid())
	{
		bIsCompiled = false;
		UE_LOG(LogInworldGraph, Error, TEXT("UInworldGraph::Execute - bad state."));
		return {};
	}

	const FString ExecutionId{FGuid::NewGuid().ToString(EGuidFormats::DigitsWithHyphensLower)};

#if !UE_BUILD_SHIPPING
	AddExecutionHistory(Instigator, ExecutionId, RuntimeData);
#endif

	{
		UInworldProcessContext* ProcessContext = UInworldProcessContext::Create(this, RuntimeData);
		FInworldGraphExecution NewExecution(ExecutionId, InputData, ProcessContext, ResultCallback, StatusCallback);
		FScopeLock Lock(&ExecutionQueueLock);
		if (ExecutionMode == EInworldGraphExecutionMode::LatestOnly)
		{
			HandleOneActiveMode(ExecutionId);
		}
		UInworldRuntimeSubsystem::Get()->RegisterProcessContext(ExecutionId, ProcessContext);
		ExecutionMap.Add(ExecutionId, NewExecution);
		if (CanStartNewExecution())
		{
			StartExecution(ExecutionId);
		}
		else
		{
			ExecutionQueue.Add(ExecutionId);
			UE_LOG(LogInworldGraph, Log, TEXT("UInworldGraph::Execute - Execution %s queued (position %d)"), *ExecutionId,
				ExecutionQueue.Num());
		}
	}

	return ExecutionId;
}

FString UInworldGraph::Execute(UObject* Instigator, const FInworldDataHandle& InputData, FOnGraphResultNative ResultCallback,
	FOnGraphStatusNative StatusCallback)
{
	return Execute(Instigator, InputData, {}, ResultCallback, StatusCallback);
}

void UInworldGraph::AddExecutionHistory(
	UObject* Instigator, const FString& ExecutionId, const TMap<FName, FInworldDataHandle>& RuntimeData)
{
	TObjectPtr<UObject> OwnerObject = IsValid(Instigator) ? Instigator : GetOuter();
	if (UActorComponent* ActorComponent = Cast<UActorComponent>(OwnerObject))
	{
		// if owned by a component, use the actor instead. (easier to read and understand for debugging)
		if (ActorComponent->GetOwner())
		{
			OwnerObject = ActorComponent->GetOwner();
		}
	}

	if (UInworldRuntimeSubsystem* RuntimeSubsystem = UInworldRuntimeSubsystem::Get())
	{
		FInworldGraphExecutionHistory& GraphExecutionHistory = RuntimeSubsystem->AddGraphExecutionHistory();
		GraphExecutionHistory.ExecutionId = ExecutionId;
		GraphExecutionHistory.GraphId = GraphId;
		GraphExecutionHistory.Owner = MakeWeakObjectPtr(OwnerObject);
		GraphExecutionHistory.OwnerName = IsValid(OwnerObject) ? OwnerObject->GetName() : FString(TEXT("none"));
	}
}

void UInworldGraph::CancelExecution(const FString& ExecutionId)
{
	if (ExecutionId.IsEmpty())
	{
		return;
	}

	bool bRemoveNow = false;
	{
		FScopeLock Lock(&ExecutionQueueLock);
		if (FInworldGraphExecution* Execution = ExecutionMap.Find(ExecutionId))
		{
			if (Execution->Status == EInworldGraphExecutionStatus::Running)
			{
				if (ExecutionMode == EInworldGraphExecutionMode::Simultaneous)
				{
					if (auto* Subsystem = UInworldRuntimeSubsystem::Get())
					{
						Subsystem->DecrementExecutionCount();
					}
				}

				if (Execution->ExecutorHandle.IsValid())
				{
					Execution->ExecutorHandle->Cancel();
				}
			}
			else if (Execution->Status == EInworldGraphExecutionStatus::Pending)
			{
				ExecutionQueue.Remove(ExecutionId);
				bRemoveNow = true;
			}

			UpdateExecutionStatus(Execution, EInworldGraphExecutionStatus::Canceled);
			UE_LOG(LogInworldGraph, Log, TEXT("UInworldGraph::CancelExecution - Canceled execution: %s"), *ExecutionId);
		}
	}

	if (bRemoveNow)
	{
		RemoveExecutionFromPool(ExecutionId);
	}
	ProcessExecutionQueue();
}

void UInworldGraph::CancelAllExecutions()
{
	FScopeLock Lock(&ExecutionQueueLock);
	int32 CanceledRunningCount = 0;

	for (auto& ExecutionPair : ExecutionMap)
	{
		FInworldGraphExecution& Execution = ExecutionPair.Value;
		if (Execution.Status == EInworldGraphExecutionStatus::Running || Execution.Status == EInworldGraphExecutionStatus::Pending)
		{
			if (Execution.Status == EInworldGraphExecutionStatus::Running)
			{
				if (Execution.ExecutorHandle.IsValid())
				{
					Execution.ExecutorHandle->Cancel();
				}

				if (ExecutionMode == EInworldGraphExecutionMode::Simultaneous)
				{
					CanceledRunningCount++;
				}
			}
			Execution.Status = EInworldGraphExecutionStatus::Canceled;
		}
		if (auto* Subsystem = UInworldRuntimeSubsystem::Get())
		{
			Subsystem->UnregisterProcessContext(ExecutionPair.Key);
		}
	}

	if (ExecutionMode == EInworldGraphExecutionMode::Simultaneous && CanceledRunningCount > 0)
	{
		UInworldRuntimeSubsystem::Get()->DecrementExecutionCount(CanceledRunningCount);
	}

	ExecutionQueue.Empty();
	ExecutionMap.Empty();
	UE_LOG(LogInworldGraph, Log, TEXT("UInworldGraph::CancelAllExecutions - Canceled all executions"));
}

bool UInworldGraph::IsExecutionCanceled(const FString& ExecutionId)
{
	if (ExecutionId.IsEmpty())
	{
		return false;
	}

	FScopeLock Lock(&ExecutionQueueLock);
	if (const FInworldGraphExecution* Execution = ExecutionMap.Find(ExecutionId))
	{
		return Execution->Status == EInworldGraphExecutionStatus::Canceled;
	}

	return false;
}

EInworldGraphExecutionStatus UInworldGraph::GetExecutionStatus(const FString& ExecutionId)
{
	if (ExecutionId.IsEmpty())
	{
		return EInworldGraphExecutionStatus::Failed;
	}

	FScopeLock Lock(&ExecutionQueueLock);
	if (const FInworldGraphExecution* Execution = ExecutionMap.Find(ExecutionId))
	{
		return Execution->Status;
	}

	return EInworldGraphExecutionStatus::Failed;
}

FInworldGraphExecution UInworldGraph::GetExecutionInfo(const FString& ExecutionId)
{
	FScopeLock Lock(&ExecutionQueueLock);
	if (const FInworldGraphExecution* Execution = ExecutionMap.Find(ExecutionId))
	{
		return *Execution;
	}

	return FInworldGraphExecution();
}

TArray<FInworldGraphExecution> UInworldGraph::GetAllExecutions()
{
	FScopeLock Lock(&ExecutionQueueLock);
	TArray<FInworldGraphExecution> AllExecutions;

	for (const auto& ExecutionPair : ExecutionMap)
	{
		AllExecutions.Add(ExecutionPair.Value);
	}

	return AllExecutions;
}

int32 UInworldGraph::GetRunningExecutionsCount()
{
	FScopeLock Lock(&ExecutionQueueLock);
	int32 RunningCount = 0;

	for (const auto& ExecutionPair : ExecutionMap)
	{
		if (ExecutionPair.Value.Status == EInworldGraphExecutionStatus::Running)
		{
			RunningCount++;
		}
	}

	return RunningCount;
}

int32 UInworldGraph::GetPendingExecutionsCount()
{
	FScopeLock Lock(&ExecutionQueueLock);
	return ExecutionQueue.Num();
}

void UInworldGraph::SetExecutionMode(EInworldGraphExecutionMode Mode)
{
	if (ExecutionMode != Mode)
	{
		ExecutionMode = Mode;
		switch (ExecutionMode)
		{
			case EInworldGraphExecutionMode::Simultaneous:
				UE_LOG(LogInworldGraph, Log, TEXT("UInworldGraph::SetExecutionMode - Changed execution mode to: Simultaneous"));
				break;
			case EInworldGraphExecutionMode::LatestOnly:
				UE_LOG(LogInworldGraph, Log, TEXT("UInworldGraph::SetExecutionMode - Changed execution mode to: LatestOnly"));
				break;
			case EInworldGraphExecutionMode::Sequential:
				UE_LOG(LogInworldGraph, Log, TEXT("UInworldGraph::SetExecutionMode - Changed execution mode to: Sequential"));
				break;
			default:
				break;
		}

		if (Mode == EInworldGraphExecutionMode::LatestOnly)
		{
			FScopeLock Lock(&ExecutionQueueLock);
			TArray<FString> RunningExecutions;

			for (const auto& ExecutionPair : ExecutionMap)
			{
				if (ExecutionPair.Value.Status == EInworldGraphExecutionStatus::Running)
				{
					RunningExecutions.Add(ExecutionPair.Key);
				}
			}

			if (RunningExecutions.Num() > 1)
			{
				for (int32 i = 0; i < RunningExecutions.Num() - 1; ++i)
				{
					if (FInworldGraphExecution* Execution = ExecutionMap.Find(RunningExecutions[i]))
					{
						if (Execution->ExecutorHandle.IsValid() && Execution->Status == EInworldGraphExecutionStatus::Running)
						{
							Execution->ExecutorHandle->Cancel();
						}
						UpdateExecutionStatus(Execution, EInworldGraphExecutionStatus::Canceled);
					}
				}
			}
		}
	}
}

bool UInworldGraph::CanStartNewExecution() const
{
	if (ExecutionMode == EInworldGraphExecutionMode::LatestOnly)
	{
		for (const auto& ExecutionPair : ExecutionMap)
		{
			if (ExecutionPair.Value.Status == EInworldGraphExecutionStatus::Running)
			{
				return false;
			}
		}
		return true;
	}
	else if (ExecutionMode == EInworldGraphExecutionMode::Sequential)
	{
		for (const auto& ExecutionPair : ExecutionMap)
		{
			if (ExecutionPair.Value.Status == EInworldGraphExecutionStatus::Running)
			{
				return false;
			}
		}
		return true;
	}
	else
	{
		if (auto* Subsystem = UInworldRuntimeSubsystem::Get())
		{
			return Subsystem->CanStartExecution();
		}
		return false;
	}
}
void UInworldGraph::HandleOneActiveMode(const FString& NewExecutionId)
{
	for (auto& ExecutionPair : ExecutionMap)
	{
		FInworldGraphExecution& Execution = ExecutionPair.Value;
		if (Execution.ExecutionId != NewExecutionId)
		{
			if (Execution.ExecutorHandle.IsValid() && Execution.Status == EInworldGraphExecutionStatus::Running)
			{
				Execution.ExecutorHandle->Cancel();
			}
			Execution.Status = EInworldGraphExecutionStatus::Canceled;
		}
	}
	ExecutionQueue.RemoveAll(
		[&NewExecutionId](const FString& Id)
		{
			return Id != NewExecutionId;
		});
}

void UInworldGraph::ProcessExecutionQueue()
{
	FScopeLock Lock(&ExecutionQueueLock);
	while (ExecutionQueue.Num() > 0 && CanStartNewExecution())
	{
		FString ExecutionId = ExecutionQueue[0];
		ExecutionQueue.RemoveAt(0);

		if (FInworldGraphExecution* Execution = ExecutionMap.Find(ExecutionId))
		{
			if (Execution->Status == EInworldGraphExecutionStatus::Pending)
			{
				StartExecution(ExecutionId);
			}
		}
	}
}

void UInworldGraph::StartExecution(const FString& ExecutionId)
{
	FInworldGraphExecution* Execution = ExecutionMap.Find(ExecutionId);
	if (!Execution || Execution->Status != EInworldGraphExecutionStatus::Pending)
	{
		return;
	}
	UpdateExecutionStatus(Execution, EInworldGraphExecutionStatus::Running);
	Execution->StartTime = FDateTime::Now();

	if (ExecutionMode == EInworldGraphExecutionMode::Simultaneous)
	{
		if (auto* Subsystem = UInworldRuntimeSubsystem::Get())
		{
			Subsystem->IncrementExecutionCount();
		}
	}

	UE_LOG(LogInworldGraph, Log, TEXT("UInworldGraph::StartExecution - Starting execution: %s"), *ExecutionId);
	TSharedPtr<inworld::BaseData> inworld_data = MakeShareable<inworld::BaseData>(Execution->InputData->CreateRuntimeData());
	if (!inworld_data.IsValid())
	{
		UE_LOG(LogInworldGraph, Error, TEXT("UInworldGraph::StartExecution - Invalid input data for execution: %s"), *ExecutionId);
		CompleteExecution(ExecutionId, false);
		return;
	}

	UE::Tasks::Launch(UE_SOURCE_LOCATION,
		[executor = executor_, inworld_data_shared = inworld_data, ExecutionId, ResultCallback = Execution->ResultCallback,
			This = TWeakObjectPtr<UInworldGraph>(this)]()
		{
			if (!executor.IsValid() || !inworld_data_shared.IsValid())
			{
				AsyncTask(ENamedThreads::GameThread,
					[ExecutionId, ResultCallback, This]()
					{
						ResultCallback.ExecuteIfBound(ExecutionId, {}, {});
						if (This.IsValid())
						{
							This->CompleteExecution(ExecutionId, false);
						}
					});
				return;
			}

			inworld::ExecutionContext ctx;
			ctx.execution_id()->set_value(TCHAR_TO_UTF8(*ExecutionId));

			TSharedPtr<inworld::ExecutionResult> ExecutionResult =
				MakeShared<inworld::ExecutionResult>(executor->Execute(*inworld_data_shared, ctx));
			if (UInworldGraph* GraphPtr = This.Get())
			{
				FScopeLock Lock(&GraphPtr->ExecutionQueueLock);
				if (FInworldGraphExecution* Exec = GraphPtr->ExecutionMap.Find(ExecutionId))
				{
					Exec->ExecutorHandle = ExecutionResult;
				}
			}

			bool bSuccess = true;
			while (bSuccess && ExecutionResult->GetResultStream().HasNext())
			{
				if (!This.IsValid() || (This.IsValid() && This->IsExecutionCanceled(ExecutionId)))
				{
					bSuccess = false;
					break;
				}

				const inworld::StatusOr_ResultItem result{ExecutionResult->GetResultStream().Read()};
				if (!result.ok())
				{
					bSuccess = false;
					break;
				}
				UInworldRuntimeSubsystem* Subsystem = UInworldRuntimeSubsystem::Get();
				if (!IsValid(Subsystem))
				{
					bSuccess = false;
					break;
				}
				FInworldDataHandle OutputData{Subsystem->Convert(result.value().safe_value())};
				const FString NodeId{
					result.value().node_id()->has_value() ? UTF8_TO_TCHAR(result.value().node_id()->get_value()) : FString{}};

				if (FInworldData_Stream * Stream{OutputData.Unwrap<FInworldData_Stream>()}; Stream != nullptr)
				{
					int32 IndexCounter = 0;
					while (Stream->HasNext())
					{
						if (!This.IsValid() || (This.IsValid() && This->IsExecutionCanceled(ExecutionId)))
						{
							bSuccess = false;
							break;
						}
						FInworldDataHandle_Stream OutputChunk;
						OutputChunk.Data = Stream->GetNext().Data;
						OutputChunk.Index = IndexCounter++;
						if (!Stream->HasNext())
						{
							OutputChunk.IsFinal = true;
						}
						AsyncTask(ENamedThreads::GameThread,
							[ExecutionId, NodeId, OutputChunk, ResultCallback]()
							{
#if WITH_EDITOR
								if (OutputChunk.Data->IsA(FInworldData_Error::StaticStruct()))
								{
									FInworldData_Error* Error = (FInworldData_Error*) OutputChunk.Data.Get();
									FMessageLog("PIE").Error(FText::FromString(Error->Reason));
								}
#endif
								ResultCallback.ExecuteIfBound(ExecutionId, NodeId, OutputChunk);
							});
					}
				}
				else
				{
					if (This.IsValid() && !This->IsExecutionCanceled(ExecutionId))
					{
						AsyncTask(ENamedThreads::GameThread,
							[ExecutionId, NodeId, OutputData, ResultCallback]()
							{
#if WITH_EDITOR
								if (FInworldData_Error* Error = OutputData.Unwrap<FInworldData_Error>(); Error != nullptr)
								{
									FMessageLog("PIE").Error(FText::FromString(Error->Reason));
								}
#endif
								ResultCallback.ExecuteIfBound(ExecutionId, NodeId, OutputData);
							});
					}
				}
			}

			AsyncTask(ENamedThreads::GameThread,
				[ExecutionId, bSuccess, This]()
				{
					if (This.IsValid())
					{
						This->CompleteExecution(ExecutionId, bSuccess);
					}
				});
		});
}

void UInworldGraph::CompleteExecution(const FString& ExecutionId, bool bSuccess)
{
	{
		FScopeLock Lock(&ExecutionQueueLock);
		if (FInworldGraphExecution* Execution = ExecutionMap.Find(ExecutionId))
		{
			if (Execution->Status == EInworldGraphExecutionStatus::Running)
			{
				if (ExecutionMode == EInworldGraphExecutionMode::Simultaneous)
				{
					if (auto* Subsystem = UInworldRuntimeSubsystem::Get())
					{
						Subsystem->DecrementExecutionCount();
					}
				}
				Execution->FinishTime = FDateTime::Now();
				UpdateExecutionStatus(
					Execution, bSuccess ? EInworldGraphExecutionStatus::Completed : EInworldGraphExecutionStatus::Failed);
				UE_LOG(LogInworldGraph, Log, TEXT("UInworldGraph::CompleteExecution - Execution %s completed with status: %s"),
					*ExecutionId, bSuccess ? TEXT("Success") : TEXT("Failed"));
			}
		}
	}
	RemoveExecutionFromPool(ExecutionId);
	ProcessExecutionQueue();
}

void UInworldGraph::RemoveExecutionFromPool(const FString& ExecutionId)
{
	if (ExecutionId.IsEmpty())
	{
		return;
	}

	FScopeLock Lock(&ExecutionQueueLock);
	ExecutionMap.Remove(ExecutionId);
	if (auto* Subsystem = UInworldRuntimeSubsystem::Get())
	{
		Subsystem->UnregisterProcessContext(ExecutionId);
	}
}

void UInworldGraph::UpdateExecutionStatus(FInworldGraphExecution* Execution, EInworldGraphExecutionStatus Status)
{
	Execution->Status = Status;
	Execution->StatusCallback.ExecuteIfBound(Execution->ExecutionId, Status);
}

TArray<FString> UInworldGraph::GetNodeNames() const
{
	TArray<FString> NodeNames;
	for (UInworldNode* Node : Nodes)
	{
		if (Node != nullptr)
		{
			NodeNames.Add(Node->NodeName);
		}
	}
	return NodeNames;
}

void UInworldGraph::SetGraphId(const FString& InGraphId)
{
	GraphId = InGraphId;
}

UInworldNode* UInworldGraph::GetNodeByName(const FString& NodeName) const
{
	for (UInworldNode* Node : Nodes)
	{
		if (Node != nullptr && Node->NodeName.Equals(NodeName, ESearchCase::IgnoreCase))
		{
			return Node;
		}
	}
	return nullptr;
}

bool UInworldGraph::GenerateConfig(TArray<FString>& Failed)
{
	UE_LOG(LogInworldGraph, Log, TEXT("UInworldGraph::GenerateConfig - %s"), *GraphId);
	// generate json
	TSharedRef<FJsonObject> GraphJson = MakeShared<FJsonObject>();
	GraphJson->SetStringField(TEXT("schema_version"), TEXT("1.2.0"));
	GraphJson->SetArrayField(TEXT("components"), TArray<TSharedPtr<FJsonValue>>{});
	GraphJson->SetArrayField(TEXT("subgraphs"), TArray<TSharedPtr<FJsonValue>>{});

	TSharedPtr<FJsonObject> Main = MakeShared<FJsonObject>();
	Main->SetStringField(TEXT("id"), GraphId);
	Main->SetArrayField(TEXT("nodes"), TArray<TSharedPtr<FJsonValue>>{});
	Main->SetArrayField(TEXT("edges"), TArray<TSharedPtr<FJsonValue>>{});
	Main->SetArrayField(TEXT("start_nodes"), TArray<TSharedPtr<FJsonValue>>{});
	Main->SetArrayField(TEXT("end_nodes"), TArray<TSharedPtr<FJsonValue>>{});
	GraphJson->SetObjectField(TEXT("main"), Main);

	// nodes
	for (UInworldNode* Node : Nodes)
	{
		if (IsValid(Node))
		{
			EInworldRuntimeErrors ErrorCode = Node->GetJsonConfig(GraphJson);
			if (ErrorCode != EInworldRuntimeErrors::SUCCESS)
			{
				const FString ErrorMsg =
					FString::Printf(TEXT("%s : %s"), *Node->GetNodeName(), *FInworldRuntimeErrors::GetErrorMessage(ErrorCode));
				Failed.Add(ErrorMsg);
				UE_LOG(LogInworldGraph, Error, TEXT("UInworldGraph::GetGraphConfig - %s"), *ErrorMsg);
			}
		}
		else
		{
			const FString ErrorMsg = TEXT("Contains not valid runtime node.");
			Failed.Add(ErrorMsg);
			UE_LOG(LogInworldGraph, Error, TEXT("UInworldGraph::GetGraphConfig - %s"), *ErrorMsg);
		}
	}

	// edges
	for (UInworldEdge* Edge : Edges)
	{
		if (IsValid(Edge))
		{
			EInworldRuntimeErrors ErrorCode = Edge->GetJsonConfig(GraphJson);
			if (ErrorCode != EInworldRuntimeErrors::SUCCESS)
			{
				const FString ErrorMsg = FString::Printf(
					TEXT("%s : %s"), *Edge->GetClass()->GetName(), *FInworldRuntimeErrors::GetErrorMessage(ErrorCode));
				Failed.Add(ErrorMsg);
				UE_LOG(LogInworldGraph, Error, TEXT("UInworldGraph::GetGraphConfig - %s"), *ErrorMsg);
			}
		}
		else
		{
			const FString ErrorMsg = TEXT("Contains not valid runtime edge.");
			Failed.Add(ErrorMsg);
			UE_LOG(LogInworldGraph, Error, TEXT("UInworldGraph::GetGraphConfig - %s"), *ErrorMsg);
		}
	}

	// start nodes
	{
		TArray<TSharedPtr<FJsonValue>> JsonStartNodes = Main->GetArrayField(TEXT("start_nodes"));
		for (const FString& StartNodeId : StartNodes)
		{
			JsonStartNodes.Add(MakeShared<FJsonValueString>(StartNodeId));
		}
		Main->SetArrayField(TEXT("start_nodes"), JsonStartNodes);
	}

	// end nodes
	{
		TArray<TSharedPtr<FJsonValue>> JsonEndNodes = Main->GetArrayField(TEXT("end_nodes"));
		for (const FString& EndNodeId : EndNodes)
		{
			JsonEndNodes.Add(MakeShared<FJsonValueString>(EndNodeId));
		}
		Main->SetArrayField(TEXT("end_nodes"), JsonEndNodes);
	}

	if (Failed.Num() == 0)
	{
		const TSharedRef<TJsonWriter<TCHAR, TPrettyJsonPrintPolicy<TCHAR>>> Writer =
			TJsonWriterFactory<TCHAR, TPrettyJsonPrintPolicy<TCHAR>>::Create(&GraphConfig);
		FJsonSerializer::Serialize(GraphJson, Writer);
		return true;
	}
	else
	{
		GraphConfig.Empty();
		return false;
	}
}

bool UInworldGraph::CheckNodeByName(const FString& NodeName) const
{
	for (UInworldNode* Node : Nodes)
	{
		if (Node != nullptr && Node->NodeName.Equals(NodeName, ESearchCase::IgnoreCase))
		{
			return true;
		}
	}
	return false;
}

inworld::graph::CompileContext::CompileContext(const UObject* InWorldContextObject, inworld::Graph& InGraph)
	: WorldContextObject{InWorldContextObject}, graph{InGraph}
{
}

namespace inworld
{
namespace graph
{

template <typename TInterface>
struct TInterfaceAttributes;

template <typename TConfig>
struct TBaseConfigAttributes;

template <typename TConfig>
struct TConfigAttributes;

#ifndef INWORLD_COMPONENT_TEMPLATE
#define INWORLD_COMPONENT_TEMPLATE(Type)                                                                                          \
	template <>                                                                                                                   \
	struct TInterfaceAttributes<inworld::Type##Interface>                                                                         \
	{                                                                                                                             \
		using LibStatusOrInterfaceType = inworld::StatusOr_##Type##Interface;                                                     \
                                                                                                                                  \
		static inline void AddLibInterface(const FString& Id, inworld::Graph& graph, const inworld::Type##Interface& interface)   \
		{                                                                                                                         \
			graph.AddComponent_##Type##Interface({TCHAR_TO_UTF8(*Id)}, interface);                                                \
		}                                                                                                                         \
                                                                                                                                  \
		static inline LibStatusOrInterfaceType GetLibInterface(                                                                   \
			const FString& Id, const inworld::CreationContext& creation_context)                                                  \
		{                                                                                                                         \
			return const_cast<inworld::CreationContext&>(creation_context).GetComponent_##Type##Interface({TCHAR_TO_UTF8(*Id)});  \
		}                                                                                                                         \
		static inline FString ToTypedId(const FString& Id)                                                                        \
		{                                                                                                                         \
			return FString(#Type) + FString("_") + Id;                                                                            \
		}                                                                                                                         \
	};                                                                                                                            \
                                                                                                                                  \
	template <>                                                                                                                   \
	struct TBaseConfigAttributes<FInworld##Type##CreationConfig>                                                                  \
	{                                                                                                                             \
		using LibInterfaceType = inworld::Type##Interface;                                                                        \
                                                                                                                                  \
		template <typename LibConfigType>                                                                                         \
		static inline TInterfaceAttributes<LibInterfaceType>::LibStatusOrInterfaceType CreateLibInterface(                        \
			const FString& Id, const inworld::CreationContext& creation_context, const LibConfigType& lib_config)                 \
		{                                                                                                                         \
			return inworld::ComponentRegistry::GetInstance().CreateComponent_##Type##Interface(                                   \
				{TCHAR_TO_UTF8(*Id)}, creation_context, lib_config);                                                              \
		}                                                                                                                         \
                                                                                                                                  \
		static inline const FInworld##Type##CreationConfig* GetConfig(const FString& Id)                                          \
		{                                                                                                                         \
			return GetDefault<UInworldRuntimeSettings>()->Type##CreationConfig.Find(Id);                                          \
		}                                                                                                                         \
	};                                                                                                                            \
                                                                                                                                  \
	template <>                                                                                                                   \
	struct TConfigAttributes<FInworldLocal##Type##CreationConfig> : public TBaseConfigAttributes<FInworld##Type##CreationConfig>  \
	{                                                                                                                             \
		using LibConfigType = inworld::Local##Type##Config;                                                                       \
	};                                                                                                                            \
	template <>                                                                                                                   \
	struct TConfigAttributes<FInworldRemote##Type##CreationConfig> : public TBaseConfigAttributes<FInworld##Type##CreationConfig> \
	{                                                                                                                             \
		using LibConfigType = inworld::Remote##Type##Config;                                                                      \
	};                                                                                                                            \
                                                                                                                                  \
	template void CompileContext::Load<FInworld##Type##CreationConfig>(const FString& Id);                                        \
	template void CompileContext::Load<FInworld##Type##CreationConfig>(                                                           \
		const FString& Id, const FInworld##Type##CreationConfig& Config);                                                         \
	template inworld::Type##Interface CompileContext::Get<inworld::Type##Interface>(const FString& Id);

INWORLD_COMPONENT_TEMPLATE(LLM)
INWORLD_COMPONENT_TEMPLATE(StreamingSTT)
INWORLD_COMPONENT_TEMPLATE(STT)
INWORLD_COMPONENT_TEMPLATE(TextEmbedder)
INWORLD_COMPONENT_TEMPLATE(TTS)
#undef INWORLD_COMPONENT_TEMPLATE
#endif

template <typename TConfig>
inline void CompileContext::Load(const FString& Id)
{
	if (const TConfig * DefaultConfig{TBaseConfigAttributes<TConfig>::GetConfig(Id)}; DefaultConfig != nullptr)
	{
		Load(Id, *DefaultConfig);
	}
	else
	{
		UE_LOG(LogInworldGraph, Error, TEXT("No default config for Id: %s."), *Id);
	}
}

template <typename TConfig>
inline void CompileContext::Load(const FString& Id, const TConfig& Config)
{
	using TLibInterfaceType = typename TBaseConfigAttributes<TConfig>::LibInterfaceType;
	using TLibStatusOrInterfaceType = typename TInterfaceAttributes<TLibInterfaceType>::LibStatusOrInterfaceType;
	const FString ComponentId{TInterfaceAttributes<TLibInterfaceType>::ToTypedId(Id)};

	ComponentMutexMapMutex.Lock();
	if (!ComponentMutexMap.Contains(ComponentId))
	{
		ComponentMutexMap.Emplace(ComponentId, MakeShared<FCriticalSection>());
	}

	TSharedPtr<FCriticalSection> ComponentMutex{ComponentMutexMap[ComponentId]};
	ComponentMutex->Lock();
	ComponentMutexMapMutex.Unlock();

	if (Get<TLibInterfaceType>(Id).is_valid())
	{
		UE_LOG(LogInworldGraph, Warning, TEXT("CompileContext::Load: Component %s already loaded with another configuration."),
			*ComponentId);
	}
	else
	{
		switch (Config.ComputeHost)
		{
			case EInworldComputeHost::LOCAL:
				Create(Id, Config.LocalCreationConfig);
				break;
			case EInworldComputeHost::REMOTE:
				Create(Id, Config.RemoteCreationConfig);
				break;
		}
	}

	ComponentMutex->Unlock();
}

template <typename TConfig>
void CompileContext::Create(const FString& Id, const TConfig& Config)
{
	using TLibConfigType = typename TConfigAttributes<TConfig>::LibConfigType;
	using TLibInterfaceType = typename TConfigAttributes<TConfig>::LibInterfaceType;
	using TLibStatusOrInterfaceType = typename TInterfaceAttributes<TLibInterfaceType>::LibStatusOrInterfaceType;
	const FString ComponentId{TInterfaceAttributes<TLibInterfaceType>::ToTypedId(Id)};

	if constexpr (std::is_same_v<TConfig, FInworldLocalTTSCreationConfig>)
	{
		UE_LOG(LogInworldGraph, Error, TEXT("TTS Component does not support the LOCAL compute host."));
		return;
	}
	else
	{
		[&]()
		{
			TLibConfigType lib_config;
			inworld::Convert(Config, lib_config);

			if (TLibStatusOrInterfaceType status_or_interface{TConfigAttributes<TConfig>::CreateLibInterface(
					ComponentId, graph.GetCreationContextHelper().value(), lib_config)};
				status_or_interface.ok())
			{
				TInterfaceAttributes<TLibInterfaceType>::AddLibInterface(ComponentId, graph, status_or_interface.value());
			}
			else
			{
				UE_LOG(LogInworldGraph, Error, TEXT("CompileContext::Load: Component %s failed to load with configuration."),
					UTF8_TO_TCHAR(status_or_interface.status().ToString().c_str()));
			}
		}();
	}
}

template <typename TInterface>
inline TInterface CompileContext::Get(const FString& Id)
{
	using TLibStatusOrInterfaceType = typename TInterfaceAttributes<TInterface>::LibStatusOrInterfaceType;
	const FString ComponentId{TInterfaceAttributes<TInterface>::ToTypedId(Id)};

	if (inworld::StatusOr_CreationContext status_or_creation_context{graph.GetCreationContextHelper()};
		status_or_creation_context.ok())
	{
		if (TLibStatusOrInterfaceType status_or_interface{
				TInterfaceAttributes<TInterface>::GetLibInterface(ComponentId, status_or_creation_context.value())};
			status_or_interface.ok())
		{
			return status_or_interface.value();
		}
	}
	return TInterface::null();
}

}	 // namespace graph

}	 // namespace inworld
