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

#include "Graph/RuntimeData/InworldGraphRuntimeData_Memory.h"

#include "Async/Async.h"
#include "Dom/JsonObject.h"
#include "Graph/Assets/InworldGraphAsset.h"
#include "InworldRuntime/Public/Graph/RuntimeData/InworldGraphRuntimeData_EventHistory.h"
#include "InworldRuntimeTypes.h"
#include "JsonObjectConverter.h"
#include "Primitives/TextEmbedder/InworldTextEmbedderTasks.h"
#include "Tasks/Task.h"
#include "Tools/InworldTextEmbedder.h"
#include "UObject/ConstructorHelpers.h"

UInworldGraphRuntimeData_Memory::UInworldGraphRuntimeData_Memory()
{
	Key = "Memory";

	const ConstructorHelpers::FObjectFinder<UInworldGraphAsset> FlashMemoryGraphAssetFinder(
		TEXT("/InworldCharacter/Blueprint/Graph/IG_Memory.IG_Memory"));
	if (FlashMemoryGraphAssetFinder.Succeeded())
	{
		FlashMemoryGraphAsset = FlashMemoryGraphAssetFinder.Object;
	}
}

FInworldDataHandle UInworldGraphRuntimeData_Memory::GetData()
{
	// Async access protection
	FScopeLock MemoryCopyScopedLock(&MemoryCopyCriticalSection);
	MemoryCopy = Memory;
	return FInworldData_Struct(MemoryCopy);
}

void UInworldGraphRuntimeData_Memory::HandleEventHistoryAddition(class UInworldGraphRuntimeData_EventHistory* EventHistory)
{
	if (!bEnabled)
	{
		return;
	}

	if (!IsValid(EventHistory))
	{
		UE_LOG(LogInworldGraph, Warning,
			TEXT("UInworldGraphRuntimeData_Memory::Initialize_Implementation -"
				 " Null event history."));
		return;
	}

	EventHistoryCounter++;
	if (EventHistoryCounter < EventHistoryLimitForSummary)
	{
		return;
	}
	EventHistoryCounter = 0;

	const FInworldHistoryEventData EventData = EventHistory->GetHistoryEventData();

	// Ignore empty
	if (EventData.SpeechEvents.Num() <= 0)
	{
		UE_LOG(LogInworldGraph, Log, TEXT("UInworldGraphRuntimeData_Memory::HandleEventHistoryAddition - No speech events"));
		return;
	}

	FInworldData_Text InputData;

	// Append queued if we have any
	if (!QueuedEvents.IsEmpty())
	{
		InputData.Text.Append(QueuedEvents);
	}

	// Compile new events into a single string
	const int32 Max = FMath::Min(EventData.SpeechEvents.Num(), EventHistoryLimitForSummary);
	for (int32 i = 0; i < Max; ++i)
	{
		InputData.Text +=
			FString::Printf(TEXT("%s: \"%s\"\n"), *EventData.SpeechEvents[i].AgentName, *EventData.SpeechEvents[i].Utterance);
	}

	// Bail on empty input
	if (InputData.Text.IsEmpty())
	{
		return;
	}

	// Grab graph instance if we haven't done so already
	if (!IsValid(FlashMemoryGraphInstance) && FlashMemoryGraphAsset)
	{
		QueuedEvents.Append(InputData.Text);
		FlashMemoryGraphAsset->GetGraphInstance(
			FOnGraphCompiledNative::CreateUObject(this, &UInworldGraphRuntimeData_Memory::OnGraphCompiled));
		return;
	}

	// Only execute one summary at a time so that we don't accumulate summaries out of order
	if (!ActiveExecutionId.IsEmpty() || !IsValid(FlashMemoryGraphInstance))
	{
		QueuedEvents.Append(InputData.Text);
		return;
	}

	// Clear the queued queue since we are doing the execution
	QueuedEvents.Empty();

	// Trigger graph execution
	ActiveExecutionId = FlashMemoryGraphInstance->Execute(this, InputData.Wrap(), TMap<FName, FInworldDataHandle>(),
		FOnGraphResultNative::CreateUObject(this, &UInworldGraphRuntimeData_Memory::OnGraphResult));
}

void UInworldGraphRuntimeData_Memory::AppendFlashMemoryToLongTermMemory(const FName& Topic, const FString& Text)
{
	// Apply to existing categories
	bool bFound = false;
	for (FInworldLongTermMemoryRecord& LongTermRecord : Memory.LongTermMemory)
	{
		if (Topic == LongTermRecord.Topic)
		{
			bFound = true;
			if (LongTermRecord.Text.IsEmpty())
			{
				LongTermRecord.Text = Text;
			}
			else
			{
				LongTermRecord.Text += FString::Printf(TEXT(", %s"), *Text);
			}
		}
	}

	// If we didn't find an existing category, then we need to create a new long-term record, including an embedding
	if (!bFound)
	{
		FInworldLongTermMemoryRecord& NewRecord = Memory.LongTermMemory.AddDefaulted_GetRef();
		NewRecord.Text = Text;
		NewRecord.Topic = Topic;

		if (!PendingTopicsToEmbed.Contains(Topic) && !InProgressTopicEmbeddings.Contains(Topic))
		{
			PendingTopicsToEmbed.Add(Topic);
		}
	}
}

void UInworldGraphRuntimeData_Memory::AppendFlashMemoryToLongTermMemory()
{
	if (Memory.FlashMemory.Records.Num() <= 0)
	{
		return;
	}

	// Reserve access to memory structure from async operations (ie from the memory retrieval node)
	FScopeLock MemoryScopedLock(&MemoryCopyCriticalSection);

	// Apply the oldest X number of flash memory records to long-term memory storage
	const int32 Size = FMath::Min(Memory.FlashMemory.Records.Num(), FlashMemorySize);
	for (int32 i = Size - 1; i >= 0; --i)
	{
		const FInworldFlashMemoryRecord& FlashRecord = Memory.FlashMemory.Records[i];

		// Add categoryless entries to the "general" category
		if (FlashRecord.Topics.Num() <= 0)
		{
			AppendFlashMemoryToLongTermMemory(FName(TEXT("General")), FlashRecord.Text);
		}

		for (const FName& Topic : FlashRecord.Topics)
		{
			AppendFlashMemoryToLongTermMemory(Topic, FlashRecord.Text);
		}

		// Remove from flash memory
		Memory.FlashMemory.Records.RemoveAt(i);
	}

	for (const FInworldLongTermMemoryRecord& Record : Memory.LongTermMemory)
	{
		UE_LOG(LogInworldGraph, Log, TEXT("%s: '%s'"), *Record.Topic.ToString(), *Record.Text);
	}

	// Now, if we have long-term memory entries without an embedding, apply them now
	EmbedCategories();
}

void UInworldGraphRuntimeData_Memory::EmbedCategories()
{
	if (!IsValid(TextEmbedder))
	{
		TextEmbedder = UInworldTextEmbedder::CreateInworldTextEmbedder(EmbedderComponentId);
	}

	if (IsValid(TextEmbedder) && PendingTopicsToEmbed.Num() > 0 && InProgressTopicEmbeddings.Num() <= 0)
	{
		InProgressTopicEmbeddings.Append(PendingTopicsToEmbed);
		PendingTopicsToEmbed.Empty();

		TArray<FString> TopicsToEmbed;
		for (const FName& InProgressTopicEmbedding : InProgressTopicEmbeddings)
		{
			TopicsToEmbed.Add(InProgressTopicEmbedding.ToString());
		}
		TextEmbedder->EmbedAsync(
			TopicsToEmbed, FOnTextEmbeddedNative::CreateUObject(this, &UInworldGraphRuntimeData_Memory::OnTextEmbedderResults));
	}
}

void UInworldGraphRuntimeData_Memory::DebugDumpFlashMemory()
{
	if (bDebugDumpFlashMemoryAfterUpdates)
	{
		UE_LOG(LogInworldGraph, Log, TEXT("UInworldGraphRuntimeData_Memory::OnGraphResult - Updated flash memory: "));
		for (const FInworldFlashMemoryRecord& Record : Memory.FlashMemory.Records)
		{
			FString TopicString;
			for (const FName& Topic : Record.Topics)
			{
				TopicString += FString::Printf(TEXT("%s, "), *Topic.ToString());
			}
			UE_LOG(LogInworldGraph, Log, TEXT("%s { %s }"), *Record.Text, *TopicString);
		}
	}
}

void UInworldGraphRuntimeData_Memory::OnGraphCompiled(class UInworldGraph* Graph, bool bSuccess)
{
	if (bSuccess)
	{
		FlashMemoryGraphInstance = Graph;

		if (!QueuedEvents.IsEmpty())
		{
			FInworldData_Text InputData;
			InputData.Text = QueuedEvents;
			QueuedEvents.Empty();

			// Trigger graph execution
			ActiveExecutionId = FlashMemoryGraphInstance->Execute(this, InputData.Wrap(), TMap<FName, FInworldDataHandle>(),
				FOnGraphResultNative::CreateUObject(this, &UInworldGraphRuntimeData_Memory::OnGraphResult));
		}
	}
}

void UInworldGraphRuntimeData_Memory::OnGraphResult(
	const FString& ExecutionId, const FString& NodeId, const FInworldDataHandle& Data)
{
	if (ExecutionId == ActiveExecutionId)
	{
		ActiveExecutionId.Empty();

		FString JsonString;
		if (const FInworldData_Text* TextData = Data.Unwrap<FInworldData_Text>())
		{
			JsonString = TextData->Text;
		}

		if (JsonString.IsEmpty())
		{
			UE_LOG(LogInworldGraph, Warning, TEXT("UInworldGraphRuntimeData_Memory::OnGraphResult - JsonString is empty"));
			return;
		}

		// Create a JSON object from the string
		TSharedPtr<FJsonObject> JsonObject;
		TSharedRef<TJsonReader<>> Reader = TJsonReaderFactory<>::Create(JsonString);

		if (!FJsonSerializer::Deserialize(Reader, JsonObject) || !JsonObject.IsValid())
		{
			UE_LOG(LogInworldGraph, Warning,
				TEXT("UInworldGraphRuntimeData_Memory::OnGraphResult - Failed to deserialize JsonObject."));
			return;
		}

		// Convert the JSON object to our struct
		FInworldFlashMemoryRecords Records;
		if (FJsonObjectConverter::JsonObjectToUStruct(JsonObject.ToSharedRef(), &Records))
		{
			// Reserve access to memory structure from async operations (ie from the memory retrieval node)
			FScopeLock MemoryScopedLock(&MemoryCopyCriticalSection);

			// Add to flash memory. Items go from newest to oldest
			Memory.FlashMemory.Records.InsertDefaulted(0, Records.Records.Num());
			for (int32 i = 0; i < Records.Records.Num(); ++i)
			{
				Memory.FlashMemory.Records[i] = Records.Records[i];
			}

			DebugDumpFlashMemory();

			// If we've reached flash memory limit then apply our current flash memory to long-term memory
			if (Memory.FlashMemory.Records.Num() >= (2 * FlashMemorySize))
			{
				AppendFlashMemoryToLongTermMemory();
			}
		}
		else
		{
			UE_LOG(LogInworldGraph, Warning, TEXT("UInworldGraphRuntimeData_Memory::OnGraphResult - JsonObjectToUStruct failed"));
		}
	}
	else
	{
		UE_LOG(LogInworldGraph, Warning, TEXT("UInworldGraphRuntimeData_Memory::OnGraphResult - Invalid execution ID"));
	}
}

void UInworldGraphRuntimeData_Memory::OnTextEmbedderResults(bool bSuccess, const TArray<FInworldTextEmbedding>& TextEmbeddings)
{
	if (!bSuccess)
	{
		PendingTopicsToEmbed.Append(InProgressTopicEmbeddings);
		InProgressTopicEmbeddings.Empty();
		return;
	}

	// Reserve access to memory structure from async operations (ie from the memory retrieval node)
	FScopeLock MemoryScopedLock(&MemoryCopyCriticalSection);

	// Apply embeddings to our long-term memory entries
	for (const FInworldTextEmbedding& TextEmbedding : TextEmbeddings)
	{
		const FName Topic = FName(TextEmbedding.NormalizedText);
		for (FInworldLongTermMemoryRecord& LongTermRecord : Memory.LongTermMemory)
		{
			if (Topic == LongTermRecord.Topic)
			{
				LongTermRecord.TopicEmbedding = TextEmbedding;
			}
		}
	}
	InProgressTopicEmbeddings.Empty();

	// Do the next round of embeddings if any categories are pending
	EmbedCategories();
}
