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

#include "Graph/Nodes/InworldNode_Intent.h"

#include "Async/Async.h"
#include "Graph/InworldGraph.h"
#include "Graph/RuntimeData/InworldGraphRuntimeData_Intent.h"
#include "InworldBlueprintFunctionLibrary.h"
#include "InworldRuntimeSettings.h"
#include "InworldRuntimeUtils.h"
#include "Misc/TVariant.h"
#include "Primitives/Intent/InworldIntentTypes.h"
#include "Primitives/TextEmbedder/InworldTextEmbedderTasks.h"

THIRD_PARTY_INCLUDES_START
#include "InworldRuntimeLibrary/inworld.h"
THIRD_PARTY_INCLUDES_END

UInworldNode_Intent::UInworldNode_Intent()
{
	NodeName = "Intent";
	bExecuteProcessInGameThread = false;
}

UInworldNode_Intent* UInworldNode_Intent::CreateNative(UObject* Outer, const FString& NodeName,
	const FInworldIntentNodeCreationConfig& CreationConfig, const FInworldIntentNodeExecutionConfig& ExecutionConfig)
{
	// Create the node
	UInworldNode_Intent* NewNode = NewObject<UInworldNode_Intent>(Outer);
	NewNode->NodeName = NodeName;
	NewNode->CreationConfig = CreationConfig;
	NewNode->ExecutionConfig = ExecutionConfig;
	return NewNode;
}

FInworldData_MatchedIntents UInworldNode_Intent::Process_Default(
	const UInworldProcessContext* ProcessContext, const FInworldData_Text& Text, const FInworldIntentsData& Intents) const
{
	FInworldData_MatchedIntents MatchedIntents;
	if (ExecutionConfig.TopNIntents == 0 || Text.Text.IsEmpty())
	{
		return MatchedIntents;
	}

	FString NormalizedInput;
	FInworldTextEmbedderUtils::NormalizeText(Text.Text, NormalizedInput);

	FInworldIntentsData MutableIntentsData{Intents};

	// 1. Attempt to find strict (exact) matches.
	TArray<FInworldIntentMatch> ExactMatches;
	StrictMatch(NormalizedInput, MutableIntentsData, ExactMatches);

	if (ExactMatches.Num() >= ExecutionConfig.TopNIntents)
	{
		// If TopN strict matches are found, no need to search further. Add the top N and return.
		MatchedIntents.IntentMatches.Append(ExactMatches.GetData(), ExecutionConfig.TopNIntents);
		return MatchedIntents;
	}
	// Remove the current matches from the search.
	for (const FInworldIntentMatch& IntentMatch : ExactMatches)
	{
		MutableIntentsData.TextEmbeddings.Remove(IntentMatch.IntentId);
		MutableIntentsData.TextEmbeddingAssets.Remove(IntentMatch.IntentId);
	}

	// 2. If no strict matches, attempt to find matches using embeddings.
	TArray<FInworldIntentMatch> EmbeddingMatches;
	inworld::StatusOr_TextEmbedderInterface Status_Or_TextEmbedderInterface =
		ProcessContext->GetTextEmbedder(CreationConfig.EmbedderComponentId);
	if (Status_Or_TextEmbedderInterface.ok())
	{
		inworld::StatusOr_vector_float StatusOr_Vector_Float =
			Status_Or_TextEmbedderInterface.value().Embed(TCHAR_TO_UTF8(*NormalizedInput));
		if (StatusOr_Vector_Float.ok())
		{
			inworld::vector_float Vector_Float = MoveTemp(StatusOr_Vector_Float).value();
			TArray<float> InputEmbedding;
			const int32 VectorSize = Vector_Float.size();
			InputEmbedding.SetNumUninitialized(VectorSize);
			for (int32 i = 0; i < VectorSize; ++i)
			{
				InputEmbedding[i] = const_cast<inworld::vector_float&>(Vector_Float).get(i);
			}
			MatchWithEmbeddings(InputEmbedding, MutableIntentsData, EmbeddingMatches);
			int32 NumEmbeddingMatchesOverThreshold{0};
			for (const FInworldIntentMatch& EmbeddingMatch : EmbeddingMatches)
			{
				if (EmbeddingMatch.Score > ExecutionConfig.SimilarityThreshold)
				{
					NumEmbeddingMatchesOverThreshold++;
				}
				else
				{
					break;
				}
			}
			EmbeddingMatches.SetNum(NumEmbeddingMatchesOverThreshold);
		}
		else
		{
			UE_LOG(LogInworldGraph, Warning, TEXT("UInworldNode_Intent::Process - Failed to embed input text."));
			return MatchedIntents;
		}
	}
	else
	{
		UE_LOG(LogInworldGraph, Warning, TEXT("UInworldNode_Intent::Process - Failed to get text embedder."));
		return MatchedIntents;
	}

	// TODO: Add support for LLM intent matching

	// The helper function returns all potential matches, sorted by score. Add the top N.
	const int32 NumExactToAdd = FMath::Min(ExactMatches.Num(), ExecutionConfig.TopNIntents);
	if (NumExactToAdd > 0)
	{
		MatchedIntents.IntentMatches.Append(ExactMatches.GetData(), NumExactToAdd);
	}
	const int32 NumEmbeddingToAdd = FMath::Min(EmbeddingMatches.Num(), ExecutionConfig.TopNIntents - NumExactToAdd);
	if (NumEmbeddingToAdd > 0)
	{
		MatchedIntents.IntentMatches.Append(EmbeddingMatches.GetData(), NumEmbeddingToAdd);
	}
	return MatchedIntents;
}

/**
 * Finds exact matches for the input text within the configured intent phrases.
 * Optimized to avoid copying data during iteration.
 *
 * @param NormalizedTextInput The normalized user input string.
 * @param IntentData The runtime data containing all intents.
 * @param OutMatches A reference to an array that will be populated with any matches found.
 */
void UInworldNode_Intent::StrictMatch(
	const FString& NormalizedTextInput, const FInworldIntentsData& IntentData, TArray<FInworldIntentMatch>& OutMatches) const
{
	OutMatches.Empty();

	for (const TTuple<FName, TSharedPtr<const TArray<FInworldTextEmbedding>>>& Intent : IntentData.TextEmbeddings)
	{
		if (!Intent.Value.IsValid())
		{
			continue;
		}
		const FName IntentName = Intent.Key;
		const TArray<FInworldTextEmbedding> TextEmbeddings = *Intent.Value.Get();

		for (const FInworldTextEmbedding& TextEmbedding : TextEmbeddings)
		{
			if (NormalizedTextInput == TextEmbedding.NormalizedText)
			{
				OutMatches.Add(FInworldIntentMatch{IntentName, 1.0f});
				// Found a match for this intent set, no need to check its other phrases.
				break;
			}
		}
	}

	for (const TTuple<FName, const UInworldTextEmbeddingAsset*>& Intent : IntentData.TextEmbeddingAssets)
	{
		const FName IntentName = Intent.Key;
		const TArray<FInworldTextEmbedding>& TextEmbeddings = Intent.Value->TextEmbeddings;

		for (const FInworldTextEmbedding& TextEmbedding : TextEmbeddings)
		{
			if (NormalizedTextInput == TextEmbedding.NormalizedText)
			{
				OutMatches.Add(FInworldIntentMatch{IntentName, 1.0f});
				// Found a match for this intent set, no need to check its other phrases.
				break;
			}
		}
	}
}

/**
 * Finds the best-matching intents using text embeddings and cosine similarity.
 * Optimized to avoid copying data during iteration and to use MoveTemp for embedding results.
 *
 * @param NormalizedTextInput The normalized user input string.
 * @param IntentData The runtime data containing all intents.
 * @param OutMatches A reference to an array that will be populated with matches, sorted by similarity score.
 */
void UInworldNode_Intent::MatchWithEmbeddings(
	const TArray<float>& InputEmbedding, const FInworldIntentsData& IntentData, TArray<FInworldIntentMatch>& OutMatches) const
{
	OutMatches.Empty();

	for (const TTuple<FName, TSharedPtr<const TArray<FInworldTextEmbedding>>>& Intent : IntentData.TextEmbeddings)
	{
		const FName IntentName = Intent.Key;
		if (!Intent.Value.IsValid())
		{
			continue;
		}
		const TArray<FInworldTextEmbedding> TextEmbeddings = *Intent.Value.Get();

		float BestSimilarity = 0.f;
		for (const FInworldTextEmbedding& TextEmbedding : TextEmbeddings)
		{
			const float Similarity = FInworldTextEmbedderUtils::CosineSimilarity(InputEmbedding, TextEmbedding.Values);
			BestSimilarity = FMath::Max(BestSimilarity, Similarity);
		}

		OutMatches.Push({IntentName, BestSimilarity});
	}

	for (const TTuple<FName, const UInworldTextEmbeddingAsset*>& Intent : IntentData.TextEmbeddingAssets)
	{
		const FName IntentName = Intent.Key;
		const TArray<FInworldTextEmbedding> TextEmbeddings = Intent.Value->TextEmbeddings;

		float BestSimilarity = 0.f;
		for (const FInworldTextEmbedding& TextEmbedding : TextEmbeddings)
		{
			const float Similarity = FInworldTextEmbedderUtils::CosineSimilarity(InputEmbedding, TextEmbedding.Values);
			BestSimilarity = FMath::Max(BestSimilarity, Similarity);
		}

		OutMatches.Push({IntentName, BestSimilarity});
	}

	// Sort by similarity score in descending order to easily get the top N results.
	OutMatches.Sort(
		[](const FInworldIntentMatch& A, const FInworldIntentMatch& B)
		{
			return A.Score > B.Score;
		});
}

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

	// text embedder component
	{
		TArray<TSharedPtr<FJsonValue>> Components = GraphJson->GetArrayField(TEXT("components"));
		const FString ComponentId = CreationConfig.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;
}
