/**
 * Copyright 2022-2025 Theai, Inc. dba Inworld AI
 */
#include "Primitives/TextEmbedder/InworldTextEmbedderTasks.h"

#include "Async/ParallelFor.h"
#include "Containers/Queue.h"
#include "InworldRuntimeTypes.h"
#include "InworldRuntimeUtils.h"
#include "Misc/Paths.h"
#include "Primitives/TextEmbedder/InworldTextEmbedding.h"
#include "Tasks/Task.h"

THIRD_PARTY_INCLUDES_START
#include "InworldRuntimeLibrary/inworld.h"
THIRD_PARTY_INCLUDES_END

static TVariant<TSharedPtr<inworld::TextEmbedderInterface>, FString> CreateTextEmbedder(
	const FInworldLocalTextEmbedderCreationConfig& LocalTextEmbedderCreationConfig)
{
#define RETURN_ERROR(Error) return TVariant<TSharedPtr<inworld::TextEmbedderInterface>, FString>{TInPlaceType<FString>(), Error};
#define RETURN_SUCCESS(Object)                                            \
	return TVariant<TSharedPtr<inworld::TextEmbedderInterface>, FString>{ \
		TInPlaceType<TSharedPtr<inworld::TextEmbedderInterface>>(), Object};

	if (const FString ModelFilePath{GetInworldExtrasPath("embedders", LocalTextEmbedderCreationConfig.Model)};
		!FPaths::FileExists(ModelFilePath))
	{
		RETURN_ERROR(FString::Format(TEXT("{0}: Model not found at {1}"), {*FString(__func__), *ModelFilePath}))
	}

	inworld::TextEmbedderCreationConfig text_embedder_config;
	inworld::LocalTextEmbedderConfig local_text_embedder_config;
	inworld::Convert(LocalTextEmbedderCreationConfig, local_text_embedder_config);
	text_embedder_config.set_local_config(local_text_embedder_config);
	if (const inworld::StatusOr_TextEmbedderInterface status_or_text_embedder{
			inworld::TextEmbedderFactory().CreateTextEmbedder(text_embedder_config)};
		status_or_text_embedder.ok())
	{
		RETURN_SUCCESS(MakeShared<inworld::TextEmbedderInterface>(status_or_text_embedder.value()));
	}
	else
	{
		RETURN_ERROR(FString::Format(
			TEXT("{0}: {1}"), {*FString(__func__), UTF8_TO_TCHAR(status_or_text_embedder.status().ToString().c_str())}))
	}
#undef RETURN_ERROR
#undef RETURN_SUCCESS
}

static TVariant<TSharedPtr<inworld::TextEmbedderInterface>, FString> CreateTextEmbedder(
	const FInworldRemoteTextEmbedderCreationConfig& RemoteTextEmbedderCreationConfig)
{
#define RETURN_ERROR(Error) return TVariant<TSharedPtr<inworld::TextEmbedderInterface>, FString>{TInPlaceType<FString>(), Error};
#define RETURN_SUCCESS(Object)                                            \
	return TVariant<TSharedPtr<inworld::TextEmbedderInterface>, FString>{ \
		TInPlaceType<TSharedPtr<inworld::TextEmbedderInterface>>(), Object};

	const FString& RuntimeApiKey{GetInworldRuntimeApiKey()};
	if (RuntimeApiKey.IsEmpty())
	{
		RETURN_ERROR(FString::Format(TEXT("{0}: Inworld Runtime API Key is not set. Check ProjectSettings."), {*FString(__func__)}))
	}

	inworld::TextEmbedderCreationConfig text_embedder_config;
	inworld::RemoteTextEmbedderConfig remote_text_embedder_config;
	inworld::Convert(RemoteTextEmbedderCreationConfig, remote_text_embedder_config);
	text_embedder_config.set_remote_config(remote_text_embedder_config);
	if (const inworld::StatusOr_TextEmbedderInterface status_or_text_embedder{
			inworld::TextEmbedderFactory().CreateTextEmbedder(text_embedder_config)};
		status_or_text_embedder.ok())
	{
		RETURN_SUCCESS(MakeShared<inworld::TextEmbedderInterface>(status_or_text_embedder.value()));
	}
	else
	{
		RETURN_ERROR(FString::Format(
			TEXT("{0}: {1}"), {*FString(__func__), UTF8_TO_TCHAR(status_or_text_embedder.status().ToString().c_str())}))
	}
#undef RETURN_ERROR
#undef RETURN_SUCCESS
}

void FInworldTextEmbedderUtils::TopKSimilarity(const TArray<float>& QueryEmbedding,
	const TArray<FInworldTextEmbedding>& FactEmbeddings, const int32& TopK, const float& Threshold, TArray<int32>& TopKIndices)
{
	// A thread-safe queue to collect scores from multiple threads.
	// Mpsc (Multiple Producer, Single Consumer) is perfect for this pattern.
	TQueue<FFactScore, EQueueMode::Mpsc> ScoresQueue;

	// --- Step 1: Calculate Scores in Parallel ---
	// ParallelFor will distribute the iterations of this loop across multiple worker threads.
	ParallelFor(FactEmbeddings.Num(),
		[&](int32 i)
		{
			const FInworldTextEmbedding& Embedding = FactEmbeddings[i];
			if (Embedding.Values.Num() != QueryEmbedding.Num())
			{
				// Note: Logging from multiple threads can jumble output. It's often better
				// to collect errors and log them once after the parallel section.
				return;	   // 'continue' is not used in lambdas, just return.
			}

			// This dot product calculation is the main work being parallelized.
			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 operation is thread-safe.
				ScoresQueue.Enqueue(FFactScore{i, Score});
			}
		});

	// --- Step 2: Reduce Results Serially ---
	// Now that the heavy lifting is done, we process the results on this thread.
	TArray<FFactScore> TopKHeap;
	TopKHeap.Reserve(TopK);
	FFactScore CurrentScore;

	// Dequeue all the scores and push them into a min-heap to find the Top K.
	while (ScoresQueue.Dequeue(CurrentScore))
	{
		if (TopKHeap.Num() < TopK)
		{
			// Fill the heap until it reaches the desired size.
			TopKHeap.HeapPush(CurrentScore);
		}
		else if (CurrentScore.Score > TopKHeap.HeapTop().Score)
		{
			// If the new score is better than the worst score in our heap, replace it.
			TopKHeap.HeapPopDiscard();
			TopKHeap.HeapPush(CurrentScore);
		}
	}

	// --- Step 3: Finalize Output ---
	// The heap is not sorted, so we sort it to get the results in order of score.
	TopKHeap.Sort(
		[](const FFactScore& A, const FFactScore& B)
		{
			return A.Score > B.Score;	 // Sort descending (best score first)
		});

	TopKIndices.Empty(TopKHeap.Num());
	for (const FFactScore& Item : TopKHeap)
	{
		TopKIndices.Add(Item.Index);
	}
}

void FInworldTextEmbedderUtils::NormalizeText(const FString& Text, FString& OutNormalizedText)
{
	OutNormalizedText = "";
	for (const TCHAR& Char : Text)
	{
		if (FChar::IsAlnum(Char) || Char == ' ')
		{
			OutNormalizedText.AppendChar(FChar::ToLower(Char));
		}
	}
}

void FInworldTextEmbedderUtils::NormalizeTextArray(const TArray<FString>& TextArray, TArray<FString>& OutNormalizedTextArray)
{
	OutNormalizedTextArray.Empty();
	OutNormalizedTextArray.SetNum(TextArray.Num());
	for (int i = 0; i < TextArray.Num(); ++i)
	{
		NormalizeText(TextArray[i], OutNormalizedTextArray[i]);
	}
}

bool FInworldTextEmbedderUtils::IsOpeningQuote(TCHAR Char)
{
	switch (Char)
	{
		case '\"':
		case '\'':
		case L'“':
		case L'‘':
			return true;
		default:
			return false;
	}
}

TCHAR FInworldTextEmbedderUtils::GetMatchingClosingQuote(TCHAR OpeningQuote)
{
	switch (OpeningQuote)
	{
		case '\"':
			return '\"';
		case '\'':
			return '\'';
		case L'“':
			return L'”';
		case L'‘':
			return L'’';
		default:
			return OpeningQuote;
	}
}

float FInworldTextEmbedderUtils::CosineSimilarity(const TArray<float>& TextEmbedding1, const TArray<float>& TextEmbedding2)
{
	if (TextEmbedding1.Num() != TextEmbedding2.Num())
	{
		UE_LOG(LogInworldTextEmbedder, Error, TEXT("CosineSimilarity - Dimension mismatch"));
		return 0;
	}
	if (TextEmbedding1.Num() == 0)
	{
		UE_LOG(LogInworldTextEmbedder, Error, TEXT("CosineSimilarity - Arrays can not be empty"));
		return 0;
	}

	float DotProduct = 0;
	float Norm1 = 0;
	float Norm2 = 0;

	for (int32 i = 0; i < TextEmbedding1.Num(); ++i)
	{
		DotProduct += TextEmbedding1[i] * TextEmbedding2[i];
		Norm1 += TextEmbedding1[i] * TextEmbedding1[i];
		Norm2 += TextEmbedding2[i] * TextEmbedding2[i];
	}

	if (Norm1 == 0 || Norm2 == 0)
	{
		return 0;
	}

	return DotProduct / (FMath::Sqrt(Norm1) * FMath::Sqrt(Norm2));
}

namespace inworld
{
namespace Tasks
{
namespace TextEmbedder
{

UE::Tasks::TTask<TVariant<TSharedPtr<inworld::TextEmbedderInterface>, FString>> CreateTextEmbedderTask(
	const FInworldTextEmbedderCreationConfig& TextEmbedderCreationConfig)
{
#define RETURN_ERROR(Error) return TVariant<TSharedPtr<inworld::TextEmbedderInterface>, FString>{TInPlaceType<FString>(), Error};

	return UE::Tasks::Launch(UE_SOURCE_LOCATION,
		[TextEmbedderCreationConfig]()
		{
			switch (TextEmbedderCreationConfig.ComputeHost)
			{
				case EInworldComputeHost::LOCAL:
				{
					const FInworldLocalTextEmbedderCreationConfig& LocalTextEmbedderCreationConfig{
						TextEmbedderCreationConfig.LocalCreationConfig};
					return CreateTextEmbedder(LocalTextEmbedderCreationConfig);
				}
				break;
				case EInworldComputeHost::REMOTE:
				{
					const FInworldRemoteTextEmbedderCreationConfig& RemoteTextEmbedderCreationConfig{
						TextEmbedderCreationConfig.RemoteCreationConfig};
					return CreateTextEmbedder(RemoteTextEmbedderCreationConfig);
				}
				break;
				default:
					RETURN_ERROR(FString::Format(
						TEXT("{0}: Invalid Enum Value: TextEmbedderCreationConfig.ComputeHost"), {*FString(__func__)}));
			}
		});
#undef RETURN_ERROR
}

UE::Tasks::TTask<TArray<int32>> TopKSimilarityTask(
	const TArray<float>& QueryEmbedding, const TArray<FInworldTextEmbedding>& FactEmbeddings, int32 K, float Threshold)
{
	return UE::Tasks::Launch(UE_SOURCE_LOCATION,
		[QueryEmbedding, FactEmbeddings, K, Threshold]()
		{
			TArray<int32> TopKSimilarityResult;
			FInworldTextEmbedderUtils::TopKSimilarity(QueryEmbedding, FactEmbeddings, K, Threshold, TopKSimilarityResult);
			return TopKSimilarityResult;
		});
}

UE::Tasks::TTask<void> NormalizeTextTask(const TArray<FString>& TextArray, TArray<FString>& OutNormalizedTextArray)
{
	return UE::Tasks::Launch(UE_SOURCE_LOCATION,
		[TextArray, &OutNormalizedTextArray]()
		{
			FInworldTextEmbedderUtils::NormalizeTextArray(TextArray, OutNormalizedTextArray);
		});
}

UE::Tasks::TTask<float> CosineSimilarityTask(const TArray<float>& TextEmbedding1, const TArray<float>& TextEmbedding2)
{
	return UE::Tasks::Launch(UE_SOURCE_LOCATION,
		[TextEmbedding1, TextEmbedding2]()
		{
			return FInworldTextEmbedderUtils::CosineSimilarity(TextEmbedding1, TextEmbedding2);
		});
}

UE::Tasks::TTask<TOptional<TArray<FInworldTextEmbedding>>> EmbedTask(
	TSharedPtr<inworld::TextEmbedderInterface> TextEmbedderInterface, const TArray<FString>& Text)
{
	if (!TextEmbedderInterface.IsValid())
	{
		UE_LOG(LogInworldTextEmbedder, Error, TEXT("Failed to embed text: Text Embedder Interface is not valid."));
		return UE::Tasks::MakeCompletedTask<TOptional<TArray<FInworldTextEmbedding>>>();
	}

	TSharedPtr<TArray<FString>> NormalizedTextPtr = MakeShared<TArray<FString>>();
	UE::Tasks::TTask<void> NormalizeTextTask = inworld::Tasks::TextEmbedder::NormalizeTextTask(Text, *NormalizedTextPtr);

	return UE::Tasks::Launch(
		UE_SOURCE_LOCATION,
		// FIX: Use an "init-capture" to create an explicitly mutable copy of Text.
		[TextEmbedderInterface, TextCopy = Text, NormalizedTextPtr]() mutable -> TOptional<TArray<FInworldTextEmbedding>>
		{
			inworld::vector_string texts;
			texts.reserve(NormalizedTextPtr->Num());
			for (const FString& String : *NormalizedTextPtr)
			{
				texts.push_back(std::string(TCHAR_TO_UTF8(*String)));
			}

			inworld::StatusOr_vector_vector_float statusOr_Vector_Vector_Float = TextEmbedderInterface->EmbedBatch(texts);
			if (!statusOr_Vector_Vector_Float.ok())
			{
				UE_LOG(LogInworldTextEmbedder, Error, TEXT("Failed to embed text: %s"),
					UTF8_TO_TCHAR(statusOr_Vector_Vector_Float.status().ToString().c_str()));
				return TOptional<TArray<FInworldTextEmbedding>>();
			}

			inworld::vector_vector_float vector_Vector_Float = statusOr_Vector_Vector_Float.value();
			const int32 NumEmbeddings = vector_Vector_Float.size();

			TArray<FInworldTextEmbedding> TextEmbeddings;
			TextEmbeddings.SetNum(NumEmbeddings);

			ParallelFor(NumEmbeddings,
				[&](int32 i)
				{
					TextEmbeddings[i].Text = MoveTemp(TextCopy[i]);
					TextEmbeddings[i].NormalizedText = MoveTemp((*NormalizedTextPtr)[i]);
					inworld::vector_float vector_Float = vector_Vector_Float.get(i);

					const int32 NumValues = vector_Float.size();
					TextEmbeddings[i].Values.Reserve(NumValues);
					for (int j = 0; j < NumValues; j++)
					{
						TextEmbeddings[i].Values.Add(vector_Float.get(j));
					}
				});

			return TOptional<TArray<FInworldTextEmbedding>>(MoveTemp(TextEmbeddings));
		},
		NormalizeTextTask);
}

UE::Tasks::TTask<TArray<FString>> ParseTextArrayString(const FString& Text)
{
	return UE::Tasks::Launch(UE_SOURCE_LOCATION,
		[Text]()
		{
			TArray<FString> TextEntries;
			FString TextBuffer = Text.TrimStartAndEnd();
			int32 NextCommaIndex = -1;

			if (TextBuffer.Len() >= 2 && TextBuffer[0] == '(' && TextBuffer[TextBuffer.Len() - 1] == ')')
			{
				TextBuffer = TextBuffer.Mid(1, TextBuffer.Len() - 2).TrimStart().TrimEnd();
			}

			while (TextBuffer.Len() > 0)
			{
				const TCHAR FirstChar = TextBuffer[0];
				if (FInworldTextEmbedderUtils::IsOpeningQuote(FirstChar))
				{
					const TCHAR ClosingQuote = FInworldTextEmbedderUtils::GetMatchingClosingQuote(FirstChar);
					TextBuffer = TextBuffer.Mid(1);

					int32 NextQuoteIndex = -1;
					if (TextBuffer.FindChar(ClosingQuote, NextQuoteIndex))
					{
						TextEntries.Add(TextBuffer.Mid(0, NextQuoteIndex));
						TextBuffer = TextBuffer.Mid(NextQuoteIndex + 1).TrimStart();

						if (TextBuffer.Len() == 0)
						{
							break;
						}
						else if (TextBuffer.FindChar(',', NextCommaIndex))
						{
							TextBuffer = TextBuffer.Mid(NextCommaIndex + 1).TrimStart();
						}
						else
						{
							UE_LOG(LogInworldTextEmbedder, Error,
								TEXT("Failed to parse text array: invalid format after a quoted entry. Expected a comma. "
									 "Buffer: %s"),
								*TextBuffer);
							return TArray<FString>{};
						}
					}
					else
					{
						UE_LOG(LogInworldTextEmbedder, Error, TEXT("Failed to embed text: invalid format"));
						return TArray<FString>{};
					}
				}
				else if (TextBuffer.FindChar(',', NextCommaIndex))
				{
					TextEntries.Add(TextBuffer.Mid(0, NextCommaIndex).TrimEnd());
					TextBuffer = TextBuffer.Mid(NextCommaIndex + 1).TrimStart();
				}
				else
				{
					TextEntries.Add(TextBuffer.TrimEnd());
					break;
				}
			}
			return TextEntries;
		});
}

}	 // namespace TextEmbedder
}	 // namespace Tasks
}	 // namespace inworld
