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

#pragma once

#include "CoreMinimal.h"
#include "Graph/InworldGraphRuntimeData.h"

#include "InworldGraphRuntimeData_Memory.generated.h"

USTRUCT(BlueprintType)
struct INWORLDCHARACTER_API FInworldFlashMemoryRecord
{
	GENERATED_BODY()

	/** The summarized text */
	UPROPERTY(BlueprintReadOnly, Category = "Inworld|Memory")
	FString Text;

	/** Topics associated with the text */
	UPROPERTY(BlueprintReadOnly, Category = "Inworld|Memory")
	TArray<FName> Topics;
};

USTRUCT(BlueprintType)
struct INWORLDCHARACTER_API FInworldFlashMemoryRecords
{
	GENERATED_BODY()

	/** List of summarized entries */
	UPROPERTY(BlueprintReadOnly, Category = "Inworld|Memory")
	TArray<FInworldFlashMemoryRecord> Records;
};

USTRUCT(BlueprintType)
struct INWORLDCHARACTER_API FInworldLongTermMemoryRecord
{
	GENERATED_BODY()

	/** Top level topic associated with the facts in this record */
	UPROPERTY(BlueprintReadOnly, Category = "Inworld|Memory")
	FName Topic;

	/** Embedding data associated with this record's topic */
	UPROPERTY(BlueprintReadOnly, Category = "Inworld|Memory")
	FInworldTextEmbedding TopicEmbedding;

	/** Summarization of facts */
	UPROPERTY(BlueprintReadOnly, Category = "Inworld|Memory")
	FString Text;
};

USTRUCT(BlueprintType)
struct INWORLDCHARACTER_API FInworldMemory
{
	GENERATED_BODY()

	/** Rolling list of summarized topic-fact entries. Has a configurable maximum size that when exceeded, we will merge this list
	 * with long-term memory */
	UPROPERTY(BlueprintReadOnly, Category = "Inworld|Memory")
	FInworldFlashMemoryRecords FlashMemory;

	/** Stable list of long-term memory entries. Items are only removed from this list when the client asks the character to
	 * "forget." This is built from summaries of flash memory items */
	UPROPERTY(BlueprintReadOnly, Category = "Inworld|Memory")
	TArray<FInworldLongTermMemoryRecord> LongTermMemory;
};

DECLARE_DYNAMIC_DELEGATE_OneParam(FOnInworldMemorySummaryGenerated, const FString&, Summary);

/**
 * @class UInworldGraphRuntimeData_Memory
 * @brief Runtime data component that is responsible for summarizing conversation history and then managing the storage of these
 * summaries in both short and long-term containers.
 */
UCLASS(BlueprintType, meta = (DisplayName = "Memory"))
class INWORLDCHARACTER_API UInworldGraphRuntimeData_Memory : public UInworldGraphRuntimeData
{
	GENERATED_BODY()

public:
	/**
	 * @brief Default constructor
	 */
	UInworldGraphRuntimeData_Memory();

	/**
	 * @brief Retrieves the memory wrapped with FInworldData_Struct for graph processing.
	 * @return Data handle containing current memory
	 */
	virtual FInworldDataHandle GetData() override;

	/**
	 * @brief Immediately triggers the asynchronous summary process. If a summary is already in process, it will queue another.
	 * @param EventHistory The dialogue to summarize
	 */
	UFUNCTION(BlueprintCallable, Category = "Inworld|Memory")
	void HandleEventHistoryAddition(class UInworldGraphRuntimeData_EventHistory* EventHistory);

	/**
	 * @brief Returns true if the memory system is enabled (default) and false otherwise.
	 */
	UFUNCTION(BlueprintPure, Category = "Inworld|Memory")
	bool IsEnabled() const
	{
		return bEnabled;
	}

	/**
	 * @brief Set if the memory system is enabled or not.
	 * @param bNewEnabled New enable status of the memory system.
	 */
	UFUNCTION(BlueprintCallable, Category = "Inworld|Memory")
	void SetEnabled(const bool bNewEnabled)
	{
		bEnabled = bNewEnabled;
	}

protected:
	void AppendFlashMemoryToLongTermMemory(const FName& Topic, const FString& Text);
	void AppendFlashMemoryToLongTermMemory();
	void EmbedCategories();

	UFUNCTION(BlueprintCallable, CallInEditor, Category = "Inworld|Memory")
	void DebugDumpFlashMemory();

	UFUNCTION()
	void OnGraphCompiled(class UInworldGraph* Graph, bool bSuccess);

	UFUNCTION()
	void OnGraphResult(const FString& ExecutionId, const FString& NodeId, const FInworldDataHandle& Data);

	UFUNCTION()
	void OnTextEmbedderResults(bool bSuccess, const TArray<FInworldTextEmbedding>& TextEmbeddings);

	UPROPERTY(EditAnywhere, Category = "Inworld|Memory")
	bool bEnabled{true};

	/** The graph that summarizes event history into a list of fact-topic pairs */
	UPROPERTY(EditAnywhere, Category = "Inworld|Memory")
	class UInworldGraphAsset* FlashMemoryGraphAsset{nullptr};

	/** When flash memory is 2x this value, it will trigger a long-term memory update. */
	UPROPERTY(EditAnywhere, Category = "Inworld|Memory")
	int32 FlashMemorySize{4};

	/** An instance of the FlashMemoryGraphAsset */
	UPROPERTY(Transient)
	class UInworldGraph* FlashMemoryGraphInstance{nullptr};

	UPROPERTY(EditAnywhere, Category = "Inworld|Memory")
	bool bDebugDumpFlashMemoryAfterUpdates{false};

	/** The text embedder component ID to use for generating long-term memory category embeddings */
	UPROPERTY(EditAnywhere, BlueprintReadWrite, Category = "Inworld|Memory",
		Meta = (GetOptions = "InworldRuntimeSettings.GetTextEmbedderIds"))
	FString EmbedderComponentId{"Default"};

	/** Text embedder instance, used to embed long-term memory categories */
	UPROPERTY(Transient)
	class UInworldTextEmbedder* TextEmbedder{nullptr};

	/** Rolling list of summarized topic-fact entries. Has a configurable maximum size that when exceeded, we will merge this list
	 * with long-term memory */
	UPROPERTY(Transient)
	FInworldMemory Memory;

	/** Memory copy for async access from the memory retrieval node */
	UPROPERTY(Transient)
	FInworldMemory MemoryCopy;

	/** Critical section for synchronizing copy of memory during async operations */
	FCriticalSection MemoryCopyCriticalSection;

	/** List of topics that need embedding and will process next */
	UPROPERTY(Transient)
	TArray<FName> PendingTopicsToEmbed;

	/** List of topics that are being processed by the text embedder */
	UPROPERTY(Transient)
	TArray<FName> InProgressTopicEmbeddings;

	/** If there is an active execution of the memory graph, then this variable will contain its ID. Empty otherwise */
	UPROPERTY(Transient, VisibleInstanceOnly, Category = "Inworld|Memory")
	FString ActiveExecutionId;

	/** We only execute one (asynchronous) summary at a time, queue other summaries until the current is finished */
	UPROPERTY(Transient, VisibleInstanceOnly, Category = "Inworld|Memory")
	FString QueuedEvents;

	/** Limit before we summarize event history and then clear it */
	UPROPERTY(VisibleInstanceOnly, Transient, Category = "Inworld|Memory")
	int32 EventHistoryLimitForSummary{6};

	/** Counter of event history additions. */
	UPROPERTY(
		VisibleInstanceOnly, Transient, Category = "Inworld|Memory", Meta = (ToolTip = "Number of turns before memory triggers"))
	int32 EventHistoryCounter{0};
};
