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

#pragma once

#include "CoreMinimal.h"
#include "Graph/InworldGraphRuntimeData.h"
#include "Primitives/Knowledge/InworldKnowledgeTypes.h"

#include "InworldGraphRuntimeData_Knowledge.generated.h"

class UInworldNode_Knowledge;

namespace inworld
{
class KnowledgeInterface;
}

/**
 * @class UInworldGraphRuntimeData_Knowledge
 * @brief Runtime data class for managing knowledge retrieval in Inworld graphs.
 *
 * This class manages knowledge sets used for information retrieval and question-answering
 * operations within Inworld graphs. It supports dynamic adding, removing, enabling, and
 * disabling of knowledge sets, with text embeddings used for semantic search and retrieval.
 */
UCLASS(BlueprintType, meta = (DisplayName = "Knowledge Data"))
class INWORLDRUNTIME_API UInworldGraphRuntimeData_Knowledge : public UInworldGraphRuntimeData
{
	GENERATED_BODY()

public:
	/**
	 * @brief Default constructor for the knowledge runtime data.
	 */
	UInworldGraphRuntimeData_Knowledge();

	/**
	 * @brief Retrieves the knowledge data wrapped in an InworldDataHandle.
	 *
	 * This override implementation returns the current knowledge base data,
	 * making it accessible to knowledge nodes within the graph system.
	 *
	 * @return Handle containing the knowledge base data
	 */
	virtual FInworldDataHandle GetData() override;

	/**
	 * @brief Initializes the knowledge runtime data.
	 *
	 * This override implementation processes initial knowledge assets, loads their
	 * text embeddings, and prepares the knowledge retrieval system for use during
	 * graph execution.
	 */
	virtual void Initialize_Implementation() override;

	/**
	 * @brief Sets the initial knowledge assets for the runtime data.
	 *
	 * This function configures the collection of knowledge assets that will be loaded
	 * and made available for knowledge retrieval during graph execution.
	 *
	 * @param KnowledgeAssets Map of knowledge set identifiers to their corresponding text embedding assets
	 */
	UFUNCTION(BlueprintCallable, Category = "Inworld|Knowledge")
	void Set(const TMap<FName, UInworldTextEmbeddingAsset*>& KnowledgeAssets);

	/**
	 * @brief Adds a new knowledge set to the runtime data, with an option to replace an existing set.
	 * @param Id The unique identifier of the knowledge set to be added.
	 * @param Knowledge The TextEmbeddings Asset to be added.
	 * @param Replace If true, replaces an existing knowledge set with the same ID. Defaults to false.
	 */
	UFUNCTION(BlueprintCallable, Category = "Inworld|Knowledge")
	void AddKnowledgeTextEmbeddingAsset(FName Id, const UInworldTextEmbeddingAsset* Knowledge, bool Replace = false);

	/**
	 * @brief Adds a new knowledge set to the runtime data, with an option to replace an existing set.
	 * @param Id The unique identifier of the knowledge set to be added.
	 * @param Knowledge The TextEmbeddings to be added.
	 * @param Replace If true, replaces an existing knowledge set with the same ID. Defaults to false.
	 */
	UFUNCTION(BlueprintCallable, Category = "Inworld|Knowledge")
	void AddKnowledgeTextEmbeddings(FName Id, const TArray<FInworldTextEmbedding>& Knowledge, bool Replace = false);

	/**
	 * @brief Removes a specific knowledge set identified by its unique ID.
	 * @param Id The unique identifier of the knowledge set to be removed.
	 */
	UFUNCTION(BlueprintCallable, Category = "Inworld|Knowledge")
	void RemoveKnowledge(FName Id);

	/**
	 * @brief Enables a specific knowledge set within the runtime data.
	 * @param Id The unique identifier for the knowledge set to be enabled.
	 */
	UFUNCTION(BlueprintCallable, Category = "Inworld|Knowledge")
	void EnableKnowledge(FName Id);

	/**
	 * @brief Disables a specific knowledge set within the runtime data.
	 * @param Id The unique identifier for the knowledge set to be disabled.
	 */
	UFUNCTION(BlueprintCallable, Category = "Inworld|Knowledge")
	void DisableKnowledge(FName Id);

	/**
	 * @brief Retrieves all knowledge set IDs currently stored in the runtime data.
	 */
	UFUNCTION(BlueprintPure, Category = "Inworld|Knowledge")
	TArray<FName> GetKnowledgeMapKeys() const;

protected:
	/** Initial knowledge assets that are automatically loaded during initialization of this runtime data */
	UPROPERTY(EditAnywhere, BlueprintReadOnly, Category = "Knowledge Runtime Data",
		meta = (ToolTip = "Initial knowledge assets that are automatically loaded during initialization of this runtime data"))
	TMap<FName, TObjectPtr<const UInworldTextEmbeddingAsset>> InitialKnowledge;

	/** Currently enabled knowledge embeddings available for retrieval */
	TMap<FName, TSharedPtr<const TArray<FInworldTextEmbedding>>> EnabledKnowledge;

	/** Currently disabled knowledge embeddings (temporarily unavailable) */
	TMap<FName, TSharedPtr<const TArray<FInworldTextEmbedding>>> DisabledKnowledge;

	/** References to enabled knowledge assets for management */
	TMap<FName, const UInworldTextEmbeddingAsset*> EnabledKnowledgeAssets;

	/** References to disabled knowledge assets for management */
	TMap<FName, const UInworldTextEmbeddingAsset*> DisabledKnowledgeAssets;

	friend UInworldNode_Knowledge;
};
