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

#pragma once

#include "CoreMinimal.h"
#include "InworldTextEmbedding.h"
#include "Misc/Paths.h"
#include "Misc/TVariant.h"
#include "Primitives/TextEmbedder/InworldTextEmbedderTypes.h"
#include "Tasks/Task.h"

namespace inworld
{
class TextEmbedderInterface;
}	 // namespace inworld

/**
 * @struct FFactScore
 * @brief Structure representing a fact with its similarity score for ranking purposes.
 *
 * This structure pairs an index with a similarity score, used for ranking and sorting
 * facts or text embeddings by their relevance to a query.
 */
struct FFactScore
{
	/** Index of the fact or embedding in the source array */
	int32 Index;

	/** Similarity score indicating relevance (higher is more similar) */
	float Score;

	/** Comparison operator for sorting by score (ascending order) */
	bool operator<(const FFactScore& Other) const
	{
		return Score < Other.Score;
	}
};

namespace inworld
{
namespace Tasks
{
namespace TextEmbedder
{

/**
 * @brief Creates an asynchronous task for initializing a text embedder interface.
 *
 * This function creates a task that initializes the text embedding system with the provided
 * configuration. The task returns either a valid text embedder interface or an error message string.
 *
 * @param TextEmbedderCreationConfig Configuration parameters for text embedder creation (local or remote)
 * @return Task that resolves to either a shared pointer to TextEmbedderInterface or an error string
 */
UE::Tasks::TTask<TVariant<TSharedPtr<inworld::TextEmbedderInterface>, FString>> INWORLDRUNTIME_API CreateTextEmbedderTask(
	const FInworldTextEmbedderCreationConfig& TextEmbedderCreationConfig);

/**
 * @brief Creates an asynchronous task to find the top K most similar embeddings to a query.
 *
 * This function creates a task that compares a query embedding against a collection of fact embeddings
 * and returns the indices of the top K most similar ones that meet the similarity threshold.
 *
 * @param QueryEmbedding The embedding vector to compare against
 * @param FactEmbeddings Array of text embeddings to search through
 * @param K Maximum number of similar embeddings to return
 * @param Threshold Minimum similarity score required for inclusion
 * @return Task that resolves to an array of indices of the most similar embeddings
 */
UE::Tasks::TTask<TArray<int32>> INWORLDRUNTIME_API TopKSimilarityTask(
	const TArray<float>& QueryEmbedding, const TArray<FInworldTextEmbedding>& FactEmbeddings, int32 K, float Threshold);

/**
 * @brief Creates an asynchronous task to normalize an array of text strings.
 *
 * This function creates a task that applies text normalization to an array of strings,
 * cleaning and standardizing the text for better embedding and comparison results.
 *
 * @param TextArray Input array of text strings to normalize
 * @param OutNormalizedTextArray Output array to populate with normalized text strings
 * @return Task that completes when normalization is finished
 */
UE::Tasks::TTask<void> INWORLDRUNTIME_API NormalizeTextTask(
	const TArray<FString>& TextArray, TArray<FString>& OutNormalizedTextArray);

/**
 * @brief Creates an asynchronous task to calculate cosine similarity between two embeddings.
 *
 * This function creates a task that computes the cosine similarity between two embedding vectors,
 * returning a value between -1 and 1 where 1 indicates identical vectors and 0 indicates orthogonal vectors.
 *
 * @param TextEmbedding1 First embedding vector for comparison
 * @param TextEmbedding2 Second embedding vector for comparison
 * @return Task that resolves to the cosine similarity score as a float
 */
UE::Tasks::TTask<float> INWORLDRUNTIME_API CosineSimilarityTask(
	const TArray<float>& TextEmbedding1, const TArray<float>& TextEmbedding2);

/**
 * @brief Creates an asynchronous task to generate embeddings for an array of text strings.
 *
 * This function creates a task that processes an array of text strings through the text embedder
 * interface to generate their corresponding embedding vectors.
 *
 * @param TextEmbedderInterface The text embedder interface to use for generating embeddings
 * @param Text Array of text strings to convert to embeddings
 * @return Task that resolves to an optional array of text embeddings (empty if failed)
 */
UE::Tasks::TTask<TOptional<TArray<FInworldTextEmbedding>>> INWORLDRUNTIME_API EmbedTask(
	TSharedPtr<inworld::TextEmbedderInterface> TextEmbedderInterface, const TArray<FString>& Text);

/**
 * @brief Creates an asynchronous task to parse a single string into an array of text strings.
 *
 * This function creates a task that parses a formatted text string (potentially containing
 * multiple entries or structured data) into an array of individual text strings.
 *
 * @param Text The input string to parse into multiple text entries
 * @return Task that resolves to an array of parsed text strings
 */
UE::Tasks::TTask<TArray<FString>> INWORLDRUNTIME_API ParseTextArrayString(const FString& Text);

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

/**
 * @class FInworldTextEmbedderUtils
 * @brief Utility class providing static functions for text embedding operations.
 *
 * This class contains static utility functions for common text embedding tasks such as
 * similarity calculations, text normalization, and quote handling for text processing.
 */
class INWORLDRUNTIME_API FInworldTextEmbedderUtils
{
public:
	/**
	 * @brief Finds the top K most similar embeddings to a query embedding.
	 *
	 * This function compares a query embedding against a collection of fact embeddings
	 * and returns the indices of the most similar ones that meet the similarity threshold.
	 *
	 * @param QueryEmbedding The embedding vector to compare against
	 * @param FactEmbeddings Array of text embeddings to search through
	 * @param TopK Maximum number of similar embeddings to return
	 * @param Threshold Minimum similarity score required for inclusion
	 * @param TopKIndices Output array populated with indices of the most similar embeddings
	 */
	static void TopKSimilarity(const TArray<float>& QueryEmbedding, const TArray<FInworldTextEmbedding>& FactEmbeddings,
		const int32& TopK, const float& Threshold, TArray<int32>& TopKIndices);

	/**
	 * @brief Calculates the cosine similarity between two embedding vectors.
	 *
	 * Computes the cosine similarity between two embedding vectors, returning a value
	 * between -1 and 1 where 1 indicates identical vectors and 0 indicates orthogonal vectors.
	 *
	 * @param TextEmbedding1 First embedding vector for comparison
	 * @param TextEmbedding2 Second embedding vector for comparison
	 * @return The cosine similarity score as a float
	 */
	static float CosineSimilarity(const TArray<float>& TextEmbedding1, const TArray<float>& TextEmbedding2);

	/**
	 * @brief Normalizes a single text string for better embedding consistency.
	 *
	 * Applies text normalization to clean and standardize the input text,
	 * improving the quality of text embeddings and comparisons.
	 *
	 * @param Text Input text string to normalize
	 * @param OutNormalizedText Output string populated with the normalized text
	 */
	static void NormalizeText(const FString& Text, FString& OutNormalizedText);

	/**
	 * @brief Normalizes an array of text strings for consistent embedding processing.
	 *
	 * Applies text normalization to an entire array of text strings, cleaning and
	 * standardizing each entry for better embedding and comparison results.
	 *
	 * @param TextArray Input array of text strings to normalize
	 * @param OutNormalizedTextArray Output array populated with normalized text strings
	 */
	static void NormalizeTextArray(const TArray<FString>& TextArray, TArray<FString>& OutNormalizedTextArray);

	/**
	 * @brief Checks if a character is an opening quote character.
	 *
	 * Determines whether the given character represents an opening quotation mark,
	 * supporting various quote types for text parsing operations.
	 *
	 * @param Char The character to check
	 * @return True if the character is an opening quote, false otherwise
	 */
	static bool IsOpeningQuote(TCHAR Char);

	/**
	 * @brief Gets the matching closing quote character for a given opening quote.
	 *
	 * Returns the appropriate closing quotation mark that matches the provided
	 * opening quote character, supporting proper quote pairing in text parsing.
	 *
	 * @param OpeningQuote The opening quote character
	 * @return The corresponding closing quote character
	 */
	static TCHAR GetMatchingClosingQuote(TCHAR OpeningQuote);
};
