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

#pragma once

#include "CoreMinimal.h"
#include "Graph/RuntimeData/InworldGraphRuntimeData_Knowledge.h"
#include "InworldNode_Custom.h"

namespace inworld
{
class TextEmbedderInterface;
namespace graph
{
struct CompileContext;
}	 // namespace graph
}	 // namespace inworld

#include "InworldNode_Knowledge.generated.h"

/**
 * @class UInworldNode_Knowledge
 * @brief A node that retrieves knowledge records based on text input queries.
 *
 * This node takes text input queries and retrieves relevant knowledge records
 * from specified knowledge bases. It inherits from TypedNode with
 * KnowledgeRecords as output and Text as input.
 *
 * @input FInworldData_Text (Query)
 * @output FInworldData_KnowledgeRecords
 */
UCLASS()
class INWORLDRUNTIME_API UInworldNode_Knowledge : public UInworldNode_Custom
{
	GENERATED_BODY()

public:
	/**
	 * @brief Default constructor for the knowledge retrieval node.
	 *
	 * Initializes the node with default settings for knowledge base querying and retrieval operations.
	 */
	UInworldNode_Knowledge();

	/**
	 * @brief Native utility function. Creates a new Knowledge node instance with the specified configuration
	 * @param Outer The UObject that will own this node
	 * @param NodeName The name to assign to the node
	 * @param EmbedderComponentId The text embedder component ID to use for semantic search
	 * @param RetrievalConfig Configuration settings for knowledge retrieval operations
	 * @return Newly created Knowledge node instance
	 */
	static UInworldNode_Knowledge* CreateNative(UObject* Outer, const FString& NodeName, const FString& EmbedderComponentId,
		const FInworldRetrievalConfig& RetrievalConfig);

private:
	/**
	 * @brief Processes text queries to retrieve relevant knowledge records using semantic search.
	 *
	 * This function performs semantic search on the available knowledge base using text embeddings
	 * to find the most relevant knowledge records for the input query. It combines embedding-based
	 * similarity search with configurable ranking and filtering options.
	 *
	 * @param ProcessContext Runtime context containing execution state and runtime data
	 * @param Text The query text to search knowledge records for
	 * @param Knowledge The knowledge data containing text embeddings and source records
	 * @return Data structure containing the most relevant knowledge records with similarity scores
	 */
	UFUNCTION()
	virtual UPARAM(meta = (DisplayName = "Knowledge Records")) FInworldData_KnowledgeRecords Process_Default(
		const UInworldProcessContext* ProcessContext, const FInworldData_Text& Text, const FInworldKnowledgeData& Knowledge) const;

	/**
	 * @brief Generates JSON configuration for the knowledge retrieval node.
	 *
	 * This override generates the JSON configuration needed for the knowledge node's
	 * processing operations, including embedder component settings, retrieval parameters,
	 * and similarity thresholds.
	 *
	 * @param GraphJson The JSON object to populate with the node's configuration data
	 * @return Error code indicating success or failure of configuration generation
	 */
	virtual EInworldRuntimeErrors GetJsonConfig(const TSharedRef<FJsonObject>& GraphJson) override;

	// void CompareQueryToFacts(const TArray<float>& QueryEmbedding, const TArray<FInworldTextEmbedding>& TextEmbeddings, const
	// int32& TopK,	const float& Threshold, FInworldData_KnowledgeRecords& KnowledgeRecord) const;

	/**
	 * @struct FKnowledgeCandidate
	 * @brief Helper structure for tracking knowledge retrieval candidates during search operations.
	 *
	 * This structure is used internally to manage potential knowledge matches during the
	 * similarity search process. It tracks the similarity score, source embedding array,
	 * and index position for efficient candidate ranking and selection.
	 */
	struct FKnowledgeCandidate
	{
		/** The similarity score for this knowledge candidate (higher scores indicate better matches) */
		float Score;

		/** Non-owning pointer to the source embedding array containing this candidate */
		const TArray<FInworldTextEmbedding>* SourceEmbeddings;

		/** Index position of this candidate within the source embedding array */
		int32 IndexInSource;

		/**
		 * @brief Comparison operator for min-heap ordering during candidate selection.
		 * @param Other The other candidate to compare against
		 * @return True if this candidate's score is less than the other's score
		 */
		bool operator<(const FKnowledgeCandidate& Other) const
		{
			return Score < Other.Score;
		}
	};

	/**
	 * @brief Finds knowledge candidates in a specific embedding set that meet the similarity threshold.
	 *
	 * This function performs similarity calculations between the query embedding and each
	 * embedding in the provided set, collecting candidates that exceed the similarity threshold.
	 * It's used as part of the larger knowledge retrieval process.
	 *
	 * @param QueryEmbedding The embedding vector of the input query text
	 * @param TextEmbeddings The set of text embeddings to search against
	 * @param Threshold The minimum similarity threshold for including candidates
	 * @param OutCandidates Output array to store qualifying knowledge candidates
	 */
	void FindCandidatesInSet(const TArray<float>& QueryEmbedding, const TArray<FInworldTextEmbedding>& TextEmbeddings,
		const float& Threshold, TArray<FKnowledgeCandidate>& OutCandidates) const;

public:
	/** The text embedder component ID to use for generating query embeddings and performing semantic search */
	UPROPERTY(EditAnywhere, BlueprintReadWrite, Category = "Inworld|Knowledge",
		meta = (GetOptions = "InworldRuntimeSettings.GetTextEmbedderIds",
			ToolTip = "The text embedder component ID to use for generating query embeddings and performing semantic search"))
	FString EmbedderComponentId{"Default"};

	/** Configuration settings for knowledge retrieval including top-K results, similarity thresholds, and source filtering */
	UPROPERTY(EditAnywhere, BlueprintReadWrite, Category = "Inworld|Knowledge",
		Meta = (ToolTip =
					"Configuration settings for knowledge retrieval including top-K results, similarity thresholds, and source filtering"))
	FInworldRetrievalConfig RetrievalConfig;
};
