#include "Graph/Nodes/InworldNode_Memory.h"

#include "Graph/InworldGraph.h"
#include "Graph/RuntimeData/InworldGraphRuntimeData_Memory.h"
#include "InworldRuntimeSettings.h"
#include "InworldRuntimeUtils.h"
#include "Primitives/TextEmbedder/InworldTextEmbedderTasks.h"

THIRD_PARTY_INCLUDES_START
#include "InworldRuntimeLibrary/inworld.h"
THIRD_PARTY_INCLUDES_END

UInworldNode_MemoryRetrieval::UInworldNode_MemoryRetrieval()
{
	NodeName = "MemoryRetrieval";
	bExecuteProcessInGameThread = false;
}

FInworldData_Text UInworldNode_MemoryRetrieval::Process_Default(
	const UInworldProcessContext* ProcessContext, const FInworldData_Text& Text, const FInworldMemory& Memory) const
{
	if (Memory.FlashMemory.Records.Num() <= 0 && Memory.LongTermMemory.Num() <= 0)
	{
		return FInworldData_Text(TEXT("None"));
	}

	TArray<float> QueryEmbedding;
	if (!Text.Text.IsEmpty())
	{
		FString NormalizedText;

		FInworldTextEmbedderUtils::NormalizeText(Text.Text, NormalizedText);

		// Embed the normalized text to get a vector for similarity search.
		inworld::StatusOr_TextEmbedderInterface Status_Or_TextEmbedderInterface =
			ProcessContext->GetTextEmbedder(EmbedderComponentId);
		if (Status_Or_TextEmbedderInterface.ok())
		{
			inworld::StatusOr_vector_float StatusOr_Vector_Float =
				Status_Or_TextEmbedderInterface.value().Embed(TCHAR_TO_UTF8(*NormalizedText));
			if (StatusOr_Vector_Float.ok())
			{
				inworld::vector_float Vector_Float = MoveTemp(StatusOr_Vector_Float).value();

				const int32 VectorSize = Vector_Float.size();
				QueryEmbedding.SetNumUninitialized(VectorSize);
				for (int32 i = 0; i < VectorSize; ++i)
				{
					QueryEmbedding[i] = Vector_Float.get(i);
				}
			}
			else
			{
				UE_LOG(LogInworldGraph, Warning, TEXT("UInworldNode_MemoryRetrieval::Process_Default - Failed to embed query."));
			}
		}
		else
		{
			UE_LOG(LogInworldGraph, Warning, TEXT("UInworldNode_MemoryRetrieval::Process_Default - Failed to get text embedder."));
		}
	}

	if (QueryEmbedding.Num() <= 0)
	{
		return FInworldData_Text(TEXT("None"));
	}

	FInworldData_Text Summary;
	const int32 MaxFlash = FMath::Min(Memory.FlashMemory.Records.Num(), FlashMemorySize);
	for (int32 i = 0; i < MaxFlash; ++i)
	{
		Summary.Text += FString::Printf(TEXT("%i). %s\n"), i + 1, *Memory.FlashMemory.Records[i].Text);
	}

	int32 Count = 0;
	for (const FInworldLongTermMemoryRecord& Entry : Memory.LongTermMemory)
	{
		const FInworldTextEmbedding& Embedding = Entry.TopicEmbedding;
		if (Embedding.Values.Num() != QueryEmbedding.Num())
		{
			continue;	 // Skip if embedding dimensions don't match.
		}

		// Calculate the dot product (similarity score).
		float Score = 0;
		for (int32 j = 0; j < Embedding.Values.Num(); ++j)
		{
			Score += QueryEmbedding[j] * Embedding.Values[j];
		}

		// Only include it if the category is close to input
		if (Score >= RetrievalConfig.Threshold)
		{
			Summary.Text += FString::Printf(TEXT("%i). %s\n"), Count + MaxFlash + 1, *Entry.Text);
			Count++;
		}

		if (Count >= RetrievalConfig.TopK)
		{
			break;
		}
	}
	return Summary;
}
