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

#include "Graph/Nodes/InworldNode_Knowledge.h"

#include "Async/Async.h"
#include "Async/ParallelFor.h"
#include "Graph/InworldGraph.h"
#include "InworldBlueprintFunctionLibrary.h"
#include "InworldRuntime/Public/Graph/RuntimeData/InworldGraphRuntimeData_Knowledge.h"
#include "InworldRuntimeSettings.h"
#include "InworldRuntimeUtils.h"
#include "Misc/TVariant.h"
#include "Primitives/TextEmbedder/InworldTextEmbedderTasks.h"

THIRD_PARTY_INCLUDES_START
#include "InworldRuntimeLibrary/inworld.h"
THIRD_PARTY_INCLUDES_END

UInworldNode_Knowledge::UInworldNode_Knowledge()
{
	NodeName = "Knowledge";
	bExecuteProcessInGameThread = false;
}

UInworldNode_Knowledge* UInworldNode_Knowledge::CreateNative(
	UObject* Outer, const FString& NodeName, const FString& EmbedderComponentId, const FInworldRetrievalConfig& RetrievalConfig)
{
	// Create the node
	UInworldNode_Knowledge* NewNode = NewObject<UInworldNode_Knowledge>(Outer);
	NewNode->NodeName = NodeName;
	NewNode->EmbedderComponentId = EmbedderComponentId;
	NewNode->RetrievalConfig = RetrievalConfig;
	return NewNode;
}

/**
 * Finds all facts within a single knowledge set that are above a similarity threshold.
 * This is the parallelized "gather" step for finding the global best results.
 * @param QueryEmbedding The vector we are searching with.
 * @param TextEmbeddings The knowledge set we are searching in.
 * @param Threshold The minimum similarity score to be considered a candidate.
 * @param OutCandidates The array to populate with all valid candidates found.
 */
void UInworldNode_Knowledge::FindCandidatesInSet(const TArray<float>& QueryEmbedding,
	const TArray<FInworldTextEmbedding>& TextEmbeddings, const float& Threshold, TArray<FKnowledgeCandidate>& OutCandidates) const
{
	// A thread-safe queue to collect results from all worker threads.
	TQueue<FKnowledgeCandidate, EQueueMode::Mpsc> CandidatesQueue;

	// Use ParallelFor to distribute the work across all available CPU cores.
	ParallelFor(TextEmbeddings.Num(),
		[&](int32 i)
		{
			const FInworldTextEmbedding& Embedding = TextEmbeddings[i];
			if (Embedding.Values.Num() != QueryEmbedding.Num())
			{
				return;	   // 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];
			}

			if (Score >= Threshold)
			{
				// Enqueue the result. This is a thread-safe operation.
				CandidatesQueue.Enqueue(FKnowledgeCandidate{Score, &TextEmbeddings, i});
			}
		});

	// After all parallel work is complete, drain the queue into the final output array.
	FKnowledgeCandidate Candidate;
	while (CandidatesQueue.Dequeue(Candidate))
	{
		OutCandidates.Add(Candidate);
	}
}

FInworldData_KnowledgeRecords UInworldNode_Knowledge::Process_Default(
	const UInworldProcessContext* ProcessContext, const FInworldData_Text& Text, const FInworldKnowledgeData& Knowledge) const
{
	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_Knowledge::Process - Failed to embed query."));
			}
		}
		else
		{
			UE_LOG(LogInworldGraph, Warning, TEXT("UInworldNode_Knowledge::Process - Failed to get text embedder."));
		}
	}

	if (QueryEmbedding.IsEmpty())
	{
		UE_LOG(LogInworldGraph, Warning, TEXT("UInworldNode_Knowledge::Process - Failed to embed query."));
		return FInworldData_KnowledgeRecords::Error(TEXT("UInworldNode_Knowledge::Process - Failed to embed query."));
	}

	TArray<FKnowledgeCandidate> AllCandidates;

	for (TSharedPtr<const TArray<FInworldTextEmbedding>> TextEmbeddings : Knowledge.TextEmbeddings)
	{
		if (TextEmbeddings.IsValid())
		{
			FindCandidatesInSet(QueryEmbedding, *TextEmbeddings.Get(), RetrievalConfig.Threshold, AllCandidates);
		}
	}

	for (const UInworldTextEmbeddingAsset* TextEmbeddingAsset : Knowledge.TextEmbeddingAssets)
	{
		if (IsValid(TextEmbeddingAsset))
		{
			FindCandidatesInSet(QueryEmbedding, TextEmbeddingAsset->TextEmbeddings, RetrievalConfig.Threshold, AllCandidates);
		}
	}

	// The KnowledgeSet.RetrievalConfig.TopK is superceded by MaxKnowledgeRecords. Since we want to gather across all knowledge
	// sets.

	// PASS 2: SELECT THE TOP K FROM ALL GATHERED CANDIDATES
	TArray<FKnowledgeCandidate> TopKHeap;
	TopKHeap.Reserve(RetrievalConfig.TopK);

	for (const FKnowledgeCandidate& Candidate : AllCandidates)
	{
		if (TopKHeap.Num() < RetrievalConfig.TopK)
		{
			TopKHeap.HeapPush(Candidate);
		}
		else if (Candidate.Score > TopKHeap.HeapTop().Score)
		{
			TopKHeap.HeapPopDiscard();
			TopKHeap.HeapPush(Candidate);
		}
	}

	// Now, build the final result from the true top candidates.
	FInworldData_KnowledgeRecords KnowledgeRecordData;
	KnowledgeRecordData.Records.Reserve(TopKHeap.Num());

	// Sort the heap for ordered results (best score first).
	TopKHeap.Sort(
		[](const FKnowledgeCandidate& A, const FKnowledgeCandidate& B)
		{
			return A.Score > B.Score;
		});

	for (const FKnowledgeCandidate& FinalCandidate : TopKHeap)
	{
		// Access the original text using the stored pointer and index.
		if (FinalCandidate.SourceEmbeddings && (*(FinalCandidate.SourceEmbeddings)).IsValidIndex(FinalCandidate.IndexInSource))
		{
			KnowledgeRecordData.Records.Add((*(FinalCandidate.SourceEmbeddings))[FinalCandidate.IndexInSource].Text);
		}
	}

	return KnowledgeRecordData;
}

EInworldRuntimeErrors UInworldNode_Knowledge::GetJsonConfig(const TSharedRef<FJsonObject>& GraphJson)
{
	Super::GetJsonConfig(GraphJson);

	// text embedder component
	{
		TArray<TSharedPtr<FJsonValue>> Components = GraphJson->GetArrayField(TEXT("components"));
		const FString ComponentId = EmbedderComponentId;
		bool bExists = false;
		for (const TSharedPtr<FJsonValue>& Component : Components)
		{
			const TSharedPtr<FJsonObject> ComponentObj = Component->AsObject();
			if (ComponentObj.IsValid() && ComponentObj->GetStringField(TEXT("id")) == ComponentId)
			{
				bExists = true;
				break;
			}
		}
		if (!bExists)
		{
			TSharedPtr<FJsonObject> Comp = MakeShared<FJsonObject>();
			Comp->SetStringField(TEXT("id"), ComponentId);
			Comp->SetStringField(TEXT("type"), TEXT("TextEmbedderInterface"));

			if (const UInworldRuntimeSettings* Settings = GetDefault<UInworldRuntimeSettings>())
			{
				if (const FInworldTextEmbedderCreationConfig* TextEmbedderCfg =
						Settings->TextEmbedderCreationConfig.Find(ComponentId))
				{
					TSharedPtr<FJsonObject> CreationCfg = MakeShared<FJsonObject>();
					TSharedPtr<FJsonObject> CreationProps = MakeShared<FJsonObject>();

					if (TextEmbedderCfg->ComputeHost == EInworldComputeHost::REMOTE)
					{
						CreationCfg->SetStringField(TEXT("type"), TEXT("RemoteTextEmbedderConfig"));
						CreationProps->SetStringField(TEXT("provider"), TextEmbedderCfg->RemoteCreationConfig.Provider);
						CreationProps->SetStringField(TEXT("model_name"), TextEmbedderCfg->RemoteCreationConfig.Model);
						CreationProps->SetStringField(TEXT("api_key"), Settings->RuntimeApiKey);
					}
					else
					{
						CreationCfg->SetStringField(TEXT("type"), TEXT("LocalTextEmbedderConfig"));
						CreationProps->SetStringField(
							TEXT("model_path"), GetInworldExtrasPath("embedders", TextEmbedderCfg->LocalCreationConfig.Model));

						TSharedPtr<FJsonObject> DeviceObj = MakeShared<FJsonObject>();
						inworld::Device device{FInworldRuntimeUtils::GetDevice(
							(TextEmbedderCfg->LocalCreationConfig.LocalComputeHost == EInworldLocalComputeHost::CPU)
								? inworld::DeviceType::kCPU
								: inworld::DeviceType::kCUDA)};
						DeviceObj->SetStringField(
							TEXT("type"), device.type() == inworld::DeviceType::kCPU ? TEXT("CPU") : TEXT("CUDA"));
						DeviceObj->SetNumberField(TEXT("index"), device.index());

						TSharedPtr<FJsonObject> DeviceInfoObj = MakeShared<FJsonObject>();
						DeviceInfoObj->SetStringField(TEXT("name"), UTF8_TO_TCHAR(device.info().name().c_str()));
						DeviceInfoObj->SetNumberField(TEXT("free_memory_bytes"), device.info().free_memory_bytes());
						DeviceInfoObj->SetNumberField(TEXT("total_memory_bytes"), device.info().total_memory_bytes());

						DeviceObj->SetObjectField(TEXT("info"), DeviceInfoObj);
						CreationProps->SetObjectField(TEXT("device"), DeviceObj);
					}

					CreationCfg->SetObjectField(TEXT("properties"), CreationProps);
					Comp->SetObjectField(TEXT("creation_config"), CreationCfg);
				}
			}
			Components.Add(MakeShared<FJsonValueObject>(Comp));
			GraphJson->SetArrayField(TEXT("components"), Components);
		}
	}
	return EInworldRuntimeErrors::SUCCESS;
}
