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

#include "InworldRuntimeSubsystem.h"

#include "Engine/Engine.h"
#include "Engine/World.h"
#include "Graph/Assets/InworldGraphAsset.h"
#include "Graph/Edges/InworldEdge_WithCondition.h"
#include "Graph/InworldData.h"
#include "Graph/InworldEdge.h"
#include "Graph/InworldGraph.h"
#include "Graph/Nodes/InworldNode_Custom.h"
#include "InworldRuntime.h"
#include "InworldRuntimeSettings.h"
#include "InworldRuntimeUtils.h"
#include "JsonObjectConverter.h"
#include "Logging/MessageLog.h"
#include "Primitives/AEC/InworldAECTypes.h"
#include "Primitives/VAD/InworldVADTypes.h"
#include "Tools/InworldTextEmbedder.h"

THIRD_PARTY_INCLUDES_START
#include "InworldRuntimeLibrary/inworld.h"
THIRD_PARTY_INCLUDES_END

#define LOCTEXT_NAMESPACE "InworldRuntimeSubsystem"

TAtomic<UInworldRuntimeSubsystem*> UInworldRuntimeSubsystem::RuntimeSubsystemInstance(nullptr);

void UInworldRuntimeSubsystem::Initialize(FSubsystemCollectionBase& Collection)
{
	Super::Initialize(Collection);

	bInitialized = true;
	RuntimeSubsystemInstance.Store(this);

	RegisterCreationConfigType("UnrealObjectInstance");
#if WITH_EDITOR
	PIEStartedHandle = FWorldDelegates::OnPIEStarted.AddUObject(this, &UInworldRuntimeSubsystem::OnPIEStarted);
#endif
}

void UInworldRuntimeSubsystem::Deinitialize()
{
	bInitialized = false;
	RuntimeSubsystemInstance.Store(nullptr);
	ClearDataConverters();

#if WITH_EDITOR
	TextEmbedderInstances.Empty();
	FWorldDelegates::OnPIEStarted.Remove(PIEStartedHandle);
#endif

	Super::Deinitialize();
}

UInworldRuntimeSubsystem* UInworldRuntimeSubsystem::Get()
{
	if (IsEngineExitRequested())
	{
		return nullptr;
	}
	if (UInworldRuntimeSubsystem* Subsystem = RuntimeSubsystemInstance.Load())
	{
		return (IsValid(Subsystem) && Subsystem->bInitialized) ? Subsystem : nullptr;
	}
	if (UEngine* Engine = GEngine)
	{
		if (UInworldRuntimeSubsystem* LazySubsystem = Engine->GetEngineSubsystem<UInworldRuntimeSubsystem>())
		{
			if (IsValid(LazySubsystem) && LazySubsystem->bInitialized)
			{
				RuntimeSubsystemInstance.Store(LazySubsystem);
				return LazySubsystem;
			}
		}
	}
	return nullptr;
}

UInworldGraph* UInworldRuntimeSubsystem::RegisterGraph(UInworldGraph* TemplateGraph, UInworldGraphAsset* GraphAsset)
{
	if (!IsValid(TemplateGraph))
	{
		return nullptr;
	}

	{
		FReadScopeLock ReadLock(GraphRegistryLock);
		if (TWeakObjectPtr<UInworldGraph>* Found = GraphRegistry.Find(TemplateGraph->GraphId))
		{
			if (Found->IsValid())
			{
				UInworldGraph* Existing = Found->Get();
				UE_LOG(LogInworldGraph, Log, TEXT("UInworldRuntimeSubsystem::RegisterGraph - Existing Graph Found: %s"),
					*Existing->GraphId);
				return Existing;
			}
		}
	}

	UInworldGraph* RegisteredGraph{nullptr};
	if (GraphAsset)
	{
		RegisteredGraph = GraphAsset->CreateGraph(this);
	}
	else
	{
		RegisteredGraph = TemplateGraph;
	}
	if (!RegisteredGraph)
	{
		return nullptr;
	}

	{
		FWriteScopeLock WriteLock(GraphRegistryLock);

		// Clean any stale pointers in the graph before adding a new one.
		for (auto It = GraphRegistry.CreateIterator(); It; ++It)
		{
			// Check if the weak pointer is stale (object was destroyed/GC'd)
			if (It.Value().IsStale())
			{
				// Remove the current element from the map
				It.RemoveCurrent();
			}
		}

		GraphRegistry.Add(TemplateGraph->GraphId, RegisteredGraph);
		UE_LOG(LogInworldGraph, Log, TEXT("UInworldRuntimeSubsystem::RegisterGraph - Created new graph: %s"),
			*RegisteredGraph->GraphId);
	}

	return RegisteredGraph;
}

void UInworldRuntimeSubsystem::SendGraphFeedback(const FInworldGraphFeedback& GraphFeedback)
{
	FString StringRating = (GraphFeedback.Rating == 0)	? FString("neutral")
						   : (GraphFeedback.Rating > 0) ? FString("positive")
														: FString("negative");

	inworld::FeedbackRecorder::RecordFeedback(TCHAR_TO_UTF8(*GraphFeedback.GraphId), TCHAR_TO_UTF8(*GraphFeedback.ExecutionId),
		TCHAR_TO_UTF8(*StringRating), TCHAR_TO_UTF8(*GraphFeedback.Description));
}

inworld::AECFilterInterface& UInworldRuntimeSubsystem::GetAEC()
{
	static FCriticalSection AECCreationLock;
	static inworld::AECFilterInterface aec_filter = inworld::AECFilterInterface::null();
	FScopeLock InputScopedLock(&AECCreationLock);
	if (!aec_filter.is_valid())
	{
		inworld::LocalAECFilterConfig local_aec_config;
		inworld::Convert(FInworldAECCreationConfig{}, local_aec_config);
		if (const auto status_or_aec_filter = inworld::AECFilterFactory().CreateAECFilter(local_aec_config);
			status_or_aec_filter.ok())
		{
			aec_filter = status_or_aec_filter.value();
		}
		else
		{
			UE_LOG(LogInworldRuntime, Error, TEXT("Failed to create AEC: %s"),
				UTF8_TO_TCHAR(status_or_aec_filter.status().ToString().c_str()));
		}
	}
	return aec_filter;
}

inworld::VADInterface& UInworldRuntimeSubsystem::GetVAD()
{
	static FCriticalSection VADCreationLock;
	static inworld::VADInterface vad = inworld::VADInterface::null();
	FScopeLock InputScopedLock(&VADCreationLock);
	if (!vad.is_valid())
	{
		inworld::LocalVADConfig local_vad_config;
		inworld::Convert(FInworldVADCreationConfig{}, local_vad_config);
		if (const auto status_or_vad = inworld::VADFactory().CreateVAD(local_vad_config); status_or_vad.ok())
		{
			vad = status_or_vad.value();
		}
		else
		{
			UE_LOG(LogInworldRuntime, Error, TEXT("Failed to create VAD: %s"),
				UTF8_TO_TCHAR(status_or_vad.status().ToString().c_str()));
		}
	}
	return vad;
}

namespace
{
static void DestroyCustomNode(void* opaque_this)
{
	if (UInworldNode_Custom* CustomNode = static_cast<UInworldNode_Custom*>(opaque_this))
	{
		CustomNode->DestroyNode();
	}
}

static void ExecuteCustomNode(void* opaque_this, int execution_id)
{
	UInworldNode_Custom* CustomNode = static_cast<UInworldNode_Custom*>(opaque_this);

	inworld::ProcessContext process_context = inworld::CustomNodeThreadedProcessExecutor::GetContext(execution_id);

	UInworldProcessContext* ProcessContext =
		UInworldRuntimeSubsystem::Get()->GetProcessContext({UTF8_TO_TCHAR(process_context.GetExecutionId().c_str())});
	UInworldProcessContext::SetProcessContext(&process_context);

	const inworld::vector_BaseData& inputs = inworld::CustomNodeThreadedProcessExecutor::GetInputs(execution_id);
	TArray<FInworldDataHandle> Inputs;
	for (int i = 0; i < inputs.size(); ++i)
	{
		Inputs.Add(UInworldRuntimeSubsystem::Get()->Convert(inputs.safe_get(i)));
	}

	inworld::CustomNodeThreadedProcessExecutor::SetOutput(
		execution_id, *inworld::ExecuteCustomNode(CustomNode, ProcessContext, Inputs)->CreateRuntimeData());
	UInworldProcessContext::SetProcessContext(nullptr);
}

static void CreateCustomNode(void* context, int execution_id)
{
	UClass* Class = static_cast<UClass*>(context);

	auto CustomNode = NewObject<UInworldNode_Custom>(GetTransientPackage(), Class);

	const std::string id{inworld::CustomNodeThreadedCreateExecutor::GetInputId(execution_id)};
	if (!CustomNode)
	{
		auto status = inworld::Status::inworld_absl_Status_createStatus__SWIG(
			inworld::StatusCode::kInvalidArgument, "Custom node id not registered");
		inworld::CustomNodeThreadedCreateExecutor::SetOutput(execution_id, status);
		return;
	}

	auto config_wrappers{inworld::CustomNodeThreadedCreateExecutor::GetInputConfigs(execution_id)};
	FUnrealObjectInstance* InstanceConfig = static_cast<FUnrealObjectInstance*>(config_wrappers.safe_get(0).value());
	CustomNode->InitializeNode(InstanceConfig->InstanceData);

	inworld::CustomNodeThreadedCreateExecutor::SetOutput(execution_id,
		inworld::CustomNodeWrapper{id, inworld::CustomNodeThreadedProcessExecutor((void*) CustomNode, (void*) ExecuteCustomNode),
			inworld::CustomNodeDestructor((void*) CustomNode, (void*) DestroyCustomNode)});
}
}	 // namespace

bool UInworldRuntimeSubsystem::RegisterNodeType(const FString& Type)
{
	if (Type.Equals("SubgraphNode"))
	{
		return true;
	}

	{
		FScopeLock Lock(&ComponentRegistryMutex);
		if (int32* Existing = ComponentRefCounts.Find(Type))
		{
			(*Existing)++;
			return true;
		}
		if (inworld::ComponentRegistry::GetInstance().IsComponentRegistered({TCHAR_TO_UTF8(*Type)}))
		{
			ComponentRefCounts.Add(Type, 1);
			return true;
		}
	}

	FString ClassPath = Type;
	if (ClassPath.StartsWith("UCN"))
	{
		ClassPath = ClassPath.RightChop(3);
	}
	UClass* LoadedClass = nullptr;
	auto LoadClassLambda = [&LoadedClass, ClassPath]()
	{
		LoadedClass = LoadObject<UClass>(nullptr, *ClassPath);
	};
	if (IsInGameThread())
	{
		LoadClassLambda();
	}
	else
	{
		FGraphEventRef LoadClassTask = FFunctionGraphTask::CreateAndDispatchWhenReady(
			MoveTemp(LoadClassLambda), TStatId(), nullptr, ENamedThreads::GameThread);
		FTaskGraphInterface::Get().WaitUntilTaskCompletes(LoadClassTask);
	}

	if (UClass* Class = LoadedClass; Class != nullptr)
	{
		FScopeLock Lock(&ComponentRegistryMutex);
		if (int32* Existing = ComponentRefCounts.Find(Type))
		{
			(*Existing)++;
			return true;
		}
		if (inworld::ComponentRegistry::GetInstance().IsComponentRegistered({TCHAR_TO_UTF8(*Type)}))
		{
			ComponentRefCounts.Add(Type, 1);
			return true;
		}
		const bool bRegistered = inworld::ComponentRegistry::GetInstance()
									 .RegisterCustomNode({TCHAR_TO_UTF8(*Type)},
										 inworld::CustomNodeThreadedCreateExecutor((void*) Class, (void*) CreateCustomNode))
									 .ok();
		if (bRegistered)
		{
			ComponentRefCounts.Add(Type, 1);
		}
		return bRegistered;
	}
	else
	{
		UE_LOG(LogInworldGraph, Error, TEXT("UInworldGraph::CompileGraph - Failed to find custom node class: %s"), *ClassPath);
	}
	return false;
}

namespace
{
static void MeetsCondition(void* context, int index)
{
	if (UInworldEdge_WithCondition* CustomEdge = static_cast<UInworldEdge_WithCondition*>(context))
	{
		inworld::EdgeContext edge_context = inworld::EdgeConditionThreadedExecutor::GetEdgeContext(index);
		UInworldProcessContext* ProcessContext =
			UInworldRuntimeSubsystem::Get()->GetProcessContext({UTF8_TO_TCHAR(edge_context.GetExecutionId().c_str())});
		FInworldDataHandle Input{UInworldRuntimeSubsystem::Get()->Convert(inworld::EdgeConditionThreadedExecutor::GetInput(index))};
		inworld::EdgeConditionThreadedExecutor::SetOutput(index, inworld::ExecuteMeetsCondition(CustomEdge, ProcessContext, Input));
	}
}
}	 // namespace

bool UInworldRuntimeSubsystem::RegisterCustomEdgeConditionType(const FString& Type, UInworldEdge_WithCondition* Edge)
{
	if (Edge == nullptr)
	{
		return false;
	}
	FScopeLock Lock(&ComponentRegistryMutex);
	if (inworld::ComponentRegistry::GetInstance().IsComponentRegistered({TCHAR_TO_UTF8(*Type)}))
	{
		return true;
	}
	return inworld::ComponentRegistry::GetInstance()
		.RegisterCustomEdgeCondition(
			{TCHAR_TO_UTF8(*Type)}, inworld::EdgeConditionThreadedExecutor((void*) Edge, (void*) MeetsCondition))
		.ok();
}

namespace
{
static void DestroyCustomConfig(void* opaque_this)
{
	FMemory::Free(opaque_this);
}

static void DeserializeCustomConfig(void* context, int execution_id)
{
	UScriptStruct* ScriptStruct = static_cast<UScriptStruct*>(context);
	auto input{inworld::CustomConfigThreadedDeserializeExecutor::GetInput(execution_id)};
	uint8* Struct = (uint8*) FMemory::Malloc(ScriptStruct->GetStructureSize());
	ScriptStruct->InitializeDefaultValue(Struct);

	TSharedPtr<FJsonObject> JsonObject;
	TSharedRef<TJsonReader<>> JsonReader = TJsonReaderFactory<>::Create(UTF8_TO_TCHAR(input.c_str()));
	FJsonSerializer::Deserialize(JsonReader, JsonObject);
	FJsonObjectConverter::JsonObjectToUStruct(JsonObject.ToSharedRef(), ScriptStruct, Struct);

	inworld::CustomConfigThreadedDeserializeExecutor::SetOutput(
		execution_id, inworld::CustomConfigWrapper(TCHAR_TO_UTF8(*ScriptStruct->GetPathName()), (void*) Struct,
						  inworld::ConfigWrapperDestructor((void*) Struct, (void*) DestroyCustomConfig)));
}

static void SerializeCustomConfig(void* context, int execution_id)
{
	UScriptStruct* ScriptStruct = static_cast<UScriptStruct*>(context);
	auto input = inworld::CustomConfigThreadedSerializeExecutor::GetInput(execution_id);
	FString JsonString;
	FJsonObjectConverter::UStructToJsonObjectString(ScriptStruct, input.value(), JsonString);
	inworld::CustomConfigThreadedSerializeExecutor::SetOutput(execution_id, TCHAR_TO_UTF8(*JsonString));
}

}	 // namespace

bool UInworldRuntimeSubsystem::RegisterCreationConfigType(const FString& Type)
{
	UStruct* Struct;
	if (Type.Equals("UnrealObjectInstance"))
	{
		Struct = TBaseStructure<FUnrealObjectInstance>::Get();
	}
	else
	{
		Struct = LoadObject<UStruct>(nullptr, *Type);
	}

	if (Struct != nullptr)
	{
		return inworld::SerializerRegistry::GetInstance()
			.RegisterCustomConfig(TCHAR_TO_UTF8(*Type),
				inworld::CustomConfigThreadedDeserializeExecutor((void*) Struct, (void*) DeserializeCustomConfig),
				inworld::CustomConfigThreadedSerializeExecutor((void*) Struct, (void*) SerializeCustomConfig))
			.ok();
	}
	return false;
}

bool UInworldRuntimeSubsystem::UnregisterComponent(const FString& Type)
{
	FScopeLock Lock(&ComponentRegistryMutex);
	if (int32* Existing = ComponentRefCounts.Find(Type))
	{
		if (*Existing > 1)
		{
			(*Existing) -= 1;
			return true;
		}
		const bool bUnregistered = inworld::ComponentRegistry::GetInstance().UnregisterComponent({TCHAR_TO_UTF8(*Type)}).ok();
		if (bUnregistered)
		{
			ComponentRefCounts.Remove(Type);
		}
		return bUnregistered;
	}
	return inworld::ComponentRegistry::GetInstance().UnregisterComponent({TCHAR_TO_UTF8(*Type)}).ok();
}

void UInworldRuntimeSubsystem::RegisterProcessContext(
	const FString& ExecutionId, TWeakObjectPtr<class UInworldProcessContext> ProcessContext)
{
	FWriteScopeLock Lock(ProcessContextRegistryLock);
	ProcessContextRegistry.Add(ExecutionId, ProcessContext);
}

void UInworldRuntimeSubsystem::UnregisterProcessContext(const FString& ExecutionId)
{
	FWriteScopeLock Lock(ProcessContextRegistryLock);
	ProcessContextRegistry.Remove(ExecutionId);
}

UInworldProcessContext* UInworldRuntimeSubsystem::GetProcessContext(const FString& ExecutionId) const
{
	FReadScopeLock Lock(ProcessContextRegistryLock);
	if (const TWeakObjectPtr<UInworldProcessContext>* Item{ProcessContextRegistry.Find(ExecutionId)}; Item != nullptr)
	{
		return Item->Get();
	}
	return nullptr;
}

void UInworldRuntimeSubsystem::RegisterNode(UInworldNode* Node)
{
	if (Node)
	{
		FWriteScopeLock Lock(NodeRegistryLock);
		NodeRegistry.Add(TWeakObjectPtr<UInworldNode>(Node));
	}
}

void UInworldRuntimeSubsystem::UnregisterNode(UInworldNode* Node)
{
	if (Node)
	{
		FWriteScopeLock Lock(NodeRegistryLock);
		NodeRegistry.Remove(TWeakObjectPtr<UInworldNode>(Node));
	}
}

bool UInworldRuntimeSubsystem::IsNodeValid(UInworldNode* Node) const
{
	FReadScopeLock Lock(NodeRegistryLock);
	if (Node == nullptr)
	{
		return false;
	}
	for (const TWeakObjectPtr<UInworldNode>& WeakNode : NodeRegistry)
	{
		if (WeakNode.Get() == Node)
		{
			return WeakNode.IsValid();
		}
	}
	return false;
}

bool UInworldRuntimeSubsystem::CanStartExecution() const
{
	FScopeLock Lock(&ExecutionsLock);
	// 0 means unlimited executions
	return MaxSimultaneousExecutions == 0 || CurrentExecutions < MaxSimultaneousExecutions;
}

void UInworldRuntimeSubsystem::IncrementExecutionCount()
{
	FScopeLock Lock(&ExecutionsLock);
	CurrentExecutions++;
}

void UInworldRuntimeSubsystem::DecrementExecutionCount(int32 EndCount)
{
	FScopeLock Lock(&ExecutionsLock);
	if (CurrentExecutions > 0)
	{
		CurrentExecutions = CurrentExecutions - EndCount;
		if (CurrentExecutions < 0)
		{
			CurrentExecutions = 0;
		}
	}
}

int32 UInworldRuntimeSubsystem::GetCurrentExecutions() const
{
	FScopeLock Lock(&ExecutionsLock);
	return CurrentExecutions;
}

int32 UInworldRuntimeSubsystem::GetMaxSimultaneousExecutions() const
{
	FScopeLock Lock(&ExecutionsLock);
	return MaxSimultaneousExecutions;
}

void UInworldRuntimeSubsystem::SetMaxSimultaneousExecutions(int32 MaxCount)
{
	FScopeLock Lock(&ExecutionsLock);
	MaxSimultaneousExecutions = FMath::Max(0, MaxCount);
}

void UInworldRuntimeSubsystem::SetCustomDataConverter(const FString& TypeId, const TFunction<FInworldDataHandle(void*)>& Converter)
{
	FWriteScopeLock WriteLock(ConverterRegistryLock);
	if (!CustomDataAsFunctionsMap.Contains(TypeId))
	{
		CustomDataAsFunctionsMap.Add(TypeId, Converter);
	}
}

#if WITH_EDITOR
UInworldTextEmbedder* UInworldRuntimeSubsystem::GetOrCreateTextEmbedder(const FString& ConfigId)
{
	if (TextEmbedderInstances.Contains(ConfigId))
	{
		if (auto* ExistingEmbedder = TextEmbedderInstances[ConfigId].Get())
		{
			return ExistingEmbedder;
		}
		else
		{
			TextEmbedderInstances.Remove(ConfigId);
		}
	}

	const UInworldRuntimeSettings* Settings = GetDefault<UInworldRuntimeSettings>();
	if (!Settings)
	{
		UE_LOG(LogInworldRuntime, Error, TEXT("GetOrCreateTextEmbedder: InworldRuntimeSettings not found."));
		return nullptr;
	}

	if (Settings->TextEmbedderCreationConfig.Find(ConfigId))
	{
		UInworldTextEmbedder* NewEmbedder = UInworldTextEmbedder::CreateInworldTextEmbedder(ConfigId);
		if (NewEmbedder)
		{
			TextEmbedderInstances.Add(ConfigId, NewEmbedder);
			return NewEmbedder;
		}
	}

	UE_LOG(LogInworldRuntime, Error, TEXT("Failed to find or create TextEmbedder for ConfigId: %s"), *ConfigId);
	return nullptr;
}
#endif

FInworldDataHandle UInworldRuntimeSubsystem::Convert(const inworld::BaseData& Data)
{
	const FString TypeId{inworld::type_id(&Data).c_str()};
	{
		FReadScopeLock ReadLock(ConverterRegistryLock);
		if (const TFunction<FInworldDataHandle(const inworld::BaseData&)>* Found = IdentifiedDataAsFunctionsMap.Find(TypeId))
		{
			return (*Found)(Data);
		}
	}

	if (inworld::CustomDataWrapper AsCustom = inworld::BaseDataAs_CustomDataWrapper(Data); AsCustom.is_valid())
	{
		const FString CustomTypeId{FString(AsCustom.type_id())};
		FReadScopeLock ReadLock(ConverterRegistryLock);
		if (const TFunction<FInworldDataHandle(void*)>* Found = CustomDataAsFunctionsMap.Find(CustomTypeId))
		{
			return (*Found)(AsCustom.value());
		}
		else
		{
			UE_LOG(LogInworldRuntime, Warning, TEXT("No converter for custom type %s"), *CustomTypeId)
			return {};
		}
	}

#define TRY_RESOLVE(Type)                                                                   \
	if (inworld::BaseDataAs_##Type(Data).is_valid())                                        \
	{                                                                                       \
		TFunction<FInworldDataHandle(const inworld::BaseData&)> Convertor =                 \
			[](const inworld::BaseData& Data) -> FInworldDataHandle                         \
		{                                                                                   \
			if (inworld::Type AsType = inworld::BaseDataAs_##Type(Data); AsType.is_valid()) \
			{                                                                               \
				return FInworldData_##Type(AsType);                                         \
			}                                                                               \
			return {};                                                                      \
		};                                                                                  \
		{                                                                                   \
			FWriteScopeLock W(ConverterRegistryLock);                                       \
			if (!IdentifiedDataAsFunctionsMap.Contains(TypeId))                             \
			{                                                                               \
				IdentifiedDataAsFunctionsMap.Add(TypeId, Convertor);                        \
			}                                                                               \
		}                                                                                   \
		return Convertor(Data);                                                             \
	}

	TRY_RESOLVE(Error);
	TRY_RESOLVE(Text);
	TRY_RESOLVE(Audio);
	TRY_RESOLVE(SafetyResult);
	TRY_RESOLVE(KnowledgeRecords);
	TRY_RESOLVE(MatchedIntents);
	TRY_RESOLVE(EventHistory);
	TRY_RESOLVE(Json);
	TRY_RESOLVE(LLMChatRequest);
	TRY_RESOLVE(LLMChatResponse);
	TRY_RESOLVE(DataStream_String);
	TRY_RESOLVE(DataStream_TTSOutput);
	TRY_RESOLVE(DataStream_AudioChunk);
	TRY_RESOLVE(TTSRequest);
	TRY_RESOLVE(MatchedKeywords);
	TRY_RESOLVE(ClassificationResult);
#undef TRY_RESOLVE

	const FString DataType{inworld::type_id(&Data).c_str()};
	UE_LOG(LogInworldRuntime, Warning, TEXT("No converter for type %s"), *DataType)
	return {};
}

void UInworldRuntimeSubsystem::ClearDataConverters()
{
	FWriteScopeLock WriteLock(ConverterRegistryLock);
	IdentifiedDataAsFunctionsMap.Empty();
}

FInworldGraphExecutionHistory& UInworldRuntimeSubsystem::AddGraphExecutionHistory()
{
	return GraphHistoryList.AddDefaulted_GetRef();
}

// TICK FOR DEBUG DATA.
// TODO: We probably want to move this to its own dedicated subsystem?
void UInworldRuntimeSubsystem::Tick(float DeltaTime)
{
#if !UE_BUILD_SHIPPING
	const UInworldRuntimeSettings* Settings{GetDefault<UInworldRuntimeSettings>()};
	if ((Settings != nullptr) && Settings->bEnableGraphExecutionHistory)
	{
		UpdateGraphHistory();
	}
#endif
}

void UInworldRuntimeSubsystem::UpdateGraphHistory()
{
	constexpr int SPANCOUNT = 100;
	const std::string execIdKey(TCHAR_TO_UTF8(TEXT("execution_id")));
	const std::string dataStoreKey(TCHAR_TO_UTF8(TEXT("data_store")));
	const std::string outputKey(TCHAR_TO_UTF8(TEXT("output")));

	inworld::vector_SpanData SpanDataList = inworld::DequeueSpanTraces(SPANCOUNT);
	if (!SpanDataList.empty())
	{
		for (int i = 0; i < SpanDataList.size(); ++i)
		{
			inworld::SpanData SpanData = SpanDataList.get(i);

			FString NodeName(UTF8_TO_TCHAR(SpanData.GetName().c_str()));
			FString NodeSpanId(UTF8_TO_TCHAR(SpanData.GetSpanId().c_str()));
			FString NodeParentSpanId(UTF8_TO_TCHAR(SpanData.GetParentSpanId().c_str()));

			FString ExecutionId;
			inworld::map_string_string attributes = SpanData.GetAttributes();
			if (attributes.has_key(execIdKey))
			{
				ExecutionId = UTF8_TO_TCHAR(attributes.get(execIdKey).c_str());
			}

			FString dataStoreJsonString;
			if (attributes.has_key(dataStoreKey))
			{
				dataStoreJsonString = FString(UTF8_TO_TCHAR(attributes.get(dataStoreKey).c_str()));
			}

			FString outputJsonString;
			if (attributes.has_key(outputKey))
			{
				outputJsonString = FString(UTF8_TO_TCHAR(attributes.get(outputKey).c_str()));
			}

			FInworldGraphExecutionHistory* GraphExecutionHistory =
				!ExecutionId.IsEmpty() ? GraphHistoryList.FindByKey(ExecutionId) : nullptr;

			if (GraphExecutionHistory)
			{
				FInworldNodeHistory* NodeHistory = GraphExecutionHistory->NodeHistory.FindByKey(NodeName);
				if (NodeHistory == nullptr)
				{
					NodeHistory = &GraphExecutionHistory->NodeHistory.AddDefaulted_GetRef();
					NodeHistory->Name = NodeName;
					NodeHistory->SpanId = NodeSpanId;
					NodeHistory->ParentSpanId = NodeParentSpanId;
				}

				NodeHistory->StartTime = SpanData.GetStartTimeInNanos();
				NodeHistory->Duration = SpanData.GetDurationInNanos();

				// Milliseconds is easier to understand.
				NodeHistory->DurationInMilliSeconds = NodeHistory->Duration / 1000000LL;

				NodeHistory->DataStoreAsJson = dataStoreJsonString;
				NodeHistory->OutputAsJson = outputJsonString;

				// sort nodes based on 'StartTime'
				GraphExecutionHistory->NodeHistory.Sort();
			}
		}

		const UInworldRuntimeSettings* Settings{GetDefault<UInworldRuntimeSettings>()};
		if (Settings != nullptr)
		{
			if (GraphHistoryList.Num() > Settings->GraphExecutionHistoryMaxCount)
			{
				constexpr int32 MinCount = 10;
				int32 MinRemoveCount = FMath::Max(MinCount, GraphHistoryList.Num() - Settings->GraphExecutionHistoryMaxCount);
#if ENGINE_MAJOR_VERSION == 5 && ENGINE_MINOR_VERSION < 4
				GraphHistoryList.RemoveAt(0, MinRemoveCount, false);
#else
				GraphHistoryList.RemoveAt(0, MinRemoveCount, EAllowShrinking::No);
#endif
			}
		}
	}
}

FString UInworldRuntimeSubsystem::GetNodeHistoryOuput(const FInworldNodeHistory& Node)
{
	FString ExtractedText;
	if (!Node.OutputAsJson.IsEmpty())
	{
		TSharedPtr<FJsonObject> JsonObject;
		TSharedRef<TJsonReader<>> JsonReader = TJsonReaderFactory<>::Create(Node.OutputAsJson);
		if (FJsonSerializer::Deserialize(JsonReader, JsonObject) && JsonObject.IsValid())
		{
			const TSharedPtr<FJsonObject>* DataObject;
			if (JsonObject->TryGetObjectField(TEXT("data"), DataObject))
			{
				if ((*DataObject)->TryGetStringField(TEXT("text"), ExtractedText))
				{
					// Replace all newlines with spaces
					ExtractedText = ExtractedText.Replace(TEXT("\r"), TEXT(" ")).Replace(TEXT("\n"), TEXT(" "));
				}
			}
		}
	}
	return ExtractedText;
}

void UInworldRuntimeSubsystem::AddExecutionHistoryError(const FString& ExecutionId, const FString& ErrorMessage)
{
	FInworldGraphExecutionHistory* GraphExecutionHistory =
		!ExecutionId.IsEmpty() ? GraphHistoryList.FindByKey(ExecutionId) : nullptr;
	if (GraphExecutionHistory)
	{
		GraphExecutionHistory->AddError(ErrorMessage);
	}
}

void UInworldRuntimeSubsystem::GetAllExecutionHistory(
	UObject* HistoryOwner, TArray<const FInworldGraphExecutionHistory*>& OwnerHistoryList) const
{
	if (HistoryOwner)
	{
		for (const FInworldGraphExecutionHistory& GraphHistory : GraphHistoryList)
		{
			if (GraphHistory.Owner.Get() == HistoryOwner)
			{
				OwnerHistoryList.Add(&GraphHistory);
			}
		}
	}
}

void UInworldRuntimeSubsystem::GetAllUniqueOwner(TArray<UObject*>& OwnerList) const
{
	OwnerList.Empty();
	for (const FInworldGraphExecutionHistory& GraphHistory : GraphHistoryList)
	{
		if (GraphHistory.Owner.IsValid())
		{
			OwnerList.AddUnique(GraphHistory.Owner.Get());
		}
	}
}

TStatId UInworldRuntimeSubsystem::GetStatId() const
{
	RETURN_QUICK_DECLARE_CYCLE_STAT(UInworldRuntimeSubsystem, STATGROUP_Tickables);
}

void UInworldRuntimeSubsystem::OnPIEStarted(UGameInstance* InGameInstance)
{
	// For each PIE session, we want to clear history and start again
	ClearGraphExecutionHistory();
}

void UInworldRuntimeSubsystem::ClearGraphExecutionHistory()
{
	GraphHistoryList.Empty();
}

#undef LOCTEXT_NAMESPACE
