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

#pragma once

#include "CoreMinimal.h"
#include "Engine/MemberReference.h"
#include "Graph/InworldData.h"
#include "Graph/InworldEdge.h"
#include "Graph/InworldGraphRuntimeData.h"
#include "Graph/InworldNode.h"
#include "HAL/ThreadSafeBool.h"
#include "InworldProcessContext.h"

#include "InworldGraph.generated.h"

INWORLDRUNTIME_API DECLARE_LOG_CATEGORY_EXTERN(LogInworldGraph, Log, All);

struct FInworldLLMCreationConfig;

namespace inworld
{
class Graph;
class ExecutorInterface;
class ExecutionResult;
}	 // namespace inworld

class UInworldRuntimeSubsystem;

DECLARE_DYNAMIC_DELEGATE_TwoParams(FOnGraphCompiled, class UInworldGraph*, Graph, bool, bSuccess);
DECLARE_DELEGATE_TwoParams(FOnGraphCompiledNative, class UInworldGraph* /*Graph*/, bool /*bSuccess*/);

DECLARE_DYNAMIC_DELEGATE_ThreeParams(
	FOnGraphResult, const FString&, ExecutionId, const FString&, NodeId, const FInworldDataHandle&, Data);
DECLARE_DELEGATE_ThreeParams(
	FOnGraphResultNative, const FString& /*ExecutionId*/, const FString& /*NodeId*/, const FInworldDataHandle& /*Data*/);

/**
 * @enum EInworldGraphExecutionStatus
 * @brief Enumeration of possible execution states for graph instances.
 */
UENUM(BlueprintType)
enum class EInworldGraphExecutionStatus : uint8
{
	/** Execution is waiting to be started */
	Pending UMETA(DisplayName = "Pending"),
	/** Execution is currently running */
	Running UMETA(DisplayName = "Running"),
	/** Execution has completed successfully */
	Completed UMETA(DisplayName = "Completed"),
	/** Execution was canceled by user request */
	Canceled UMETA(DisplayName = "Canceled"),
	/** Execution failed due to an error */
	Failed UMETA(DisplayName = "Failed")
};

DECLARE_DYNAMIC_DELEGATE_TwoParams(FOnGraphStatus, const FString&, ExecutionId, EInworldGraphExecutionStatus, Status);
DECLARE_DELEGATE_TwoParams(FOnGraphStatusNative, const FString& /*ExecutionId*/, EInworldGraphExecutionStatus /*Status*/);

/**
 * @enum EInworldGraphExecutionMode
 * @brief Enumeration of execution modes that control how multiple graph executions are handled.
 */
UENUM(BlueprintType)
enum class EInworldGraphExecutionMode : uint8
{
	/** Allows multiple executions to run simultaneously across all graph instances.
	 *  New executions are queued when global limit is reached. */
	Simultaneous UMETA(DisplayName = "Simultaneous"),
	/** Cancels previous execution when starting new one, only keeping the latest execution active */
	LatestOnly UMETA(DisplayName = "Latest Only"),
	/** Sequential execution per graph instance: pending executions wait until current finishes (FIFO), no cancellations */
	Sequential UMETA(DisplayName = "Sequential")
};

/**
 * @struct FInworldGraphExecution
 * @brief Represents a single execution instance of an Inworld graph.
 *
 * This structure tracks all information related to a specific graph execution,
 * including timing data, status, input/output data, and callback functions.
 */
USTRUCT(BlueprintType)
struct INWORLDRUNTIME_API FInworldGraphExecution
{
	GENERATED_BODY()

	/** Unique identifier for this execution instance */
	UPROPERTY(
		BlueprintReadOnly, Category = "Inworld|GraphExecution", Meta = (ToolTip = "Unique identifier for this execution instance"))
	FString ExecutionId;

	/** Current status of the execution */
	UPROPERTY(BlueprintReadOnly, Category = "Inworld|GraphExecution", Meta = (ToolTip = "Current status of the execution"))
	EInworldGraphExecutionStatus Status = EInworldGraphExecutionStatus::Pending;

	/** Input data provided to the graph for this execution */
	UPROPERTY(BlueprintReadOnly, Category = "Inworld|GraphExecution",
		Meta = (ToolTip = "Input data provided to the graph for this execution"))
	FInworldDataHandle InputData;

	/** Process context containing runtime data accessible to nodes */
	UPROPERTY(Meta = (ToolTip = "Process context containing runtime data accessible to nodes"))
	TObjectPtr<UInworldProcessContext> ProcessContext;

	/** Callback function for receiving execution results */
	FOnGraphResultNative ResultCallback;

	/** Callback function for receiving status updates */
	FOnGraphStatusNative StatusCallback;

	/** Time when the execution was requested */
	UPROPERTY(BlueprintReadOnly, Category = "Inworld|GraphExecution", Meta = (ToolTip = "Time when the execution was requested"))
	FDateTime RequestTime;

	/** Time when the execution actually started */
	UPROPERTY(BlueprintReadOnly, Category = "Inworld|GraphExecution", Meta = (ToolTip = "Time when the execution actually started"))
	FDateTime StartTime;

	/** Time when the execution finished (completed, canceled, or failed) */
	UPROPERTY(BlueprintReadOnly, Category = "Inworld|GraphExecution",
		Meta = (ToolTip = "Time when the execution finished (completed, canceled, or failed)"))
	FDateTime FinishTime;

	TSharedPtr<inworld::ExecutionResult> ExecutorHandle = nullptr;

	FInworldGraphExecution()
	{
		RequestTime = FDateTime::Now();
	}

	FInworldGraphExecution(const FString& InExecutionId, const FInworldDataHandle& InInputData,
		UInworldProcessContext* InProcessContext, FOnGraphResultNative InResultCallback, FOnGraphStatusNative InStatusCallback)
		: ExecutionId(InExecutionId)
		, InputData(InInputData)
		, ProcessContext(InProcessContext)
		, ResultCallback(InResultCallback)
		, StatusCallback(InStatusCallback)
	{
		RequestTime = FDateTime::Now();
	}
};

/**
 * @brief Represents a blueprintable and programmable runtime graph.
 * Provides functionality for creating, compiling, and executing graphs asynchronously using callbacks.
 * The graph can manage nodes, edges, and execution states for real-time execution purposes.
 */
UCLASS(Blueprintable, BlueprintType)
class INWORLDRUNTIME_API UInworldGraph : public UObject
{
	GENERATED_BODY()

public:
	/**
	 * @brief Called before the object is destroyed
	 * Overrides the UObject's BeginDestroy function to handle cleanup
	 */
	virtual void BeginDestroy() override;

	/**
	 * @brief Creates and compile an instance of the given InworldGraph asset.
	 * @param InGraphScope Use managed singleton graph or local
	 * @param Callback Called when the graph initialization is finished
	 * @return An instance of the UInworldGraph archetype
	 */
	UFUNCTION(BlueprintCallable, Category = "Inworld|Graph")
	void GetGraphInstance(FOnGraphCompiled Callback);
	void GetGraphInstance(FOnGraphCompiledNative Callback);

	/**
	 * @brief Executes the graph with the given input data. This is an asynchronous function, the
	 * results will be returned in the given callback function. The result is returned as a data handle. The
	 * type of the data depends on the end node of the graph. InworldData.h contains the supported data types.
	 * See InworldBlueprintFunctionLibrary for the conversion functions.
	 *
	 * Input Example:
	 *
	 * FInworldData_Text TextData;
	 * TextData.Text = TEXT("Here is an example.");
	 * FInworldDataHandle Handle = UInworldBlueprintFunctionLibrary::WrapInworldData(TextData);
	 *
	 * @param Instigator The object/actor/component that started this execution
	 * @param InputData The input data that flows through the graph
	 * @param RuntimeData The data accessible globally by all nodes and edges in the graph
	 * @param ResultCallback Blueprint callback function to receive execution results
	 * @param StatusCallback Blueprint callback function to receive status updates
	 * @return A string identifier for the execution instance (used by graph result callback and CancelExecution)
	 */
	UFUNCTION(
		BlueprintCallable, Category = "Inworld|Graph", meta = (DefaultToSelf = "Instigator", AutoCreateRefTerm = "RuntimeData"))
	FString Execute(UObject* Instigator, const FInworldDataHandle& InputData, const TMap<FName, FInworldDataHandle>& RuntimeData,
		FOnGraphResult ResultCallback, FOnGraphStatus StatusCallback);
	FString Execute(UObject* Instigator, const FInworldDataHandle& InputData, const TMap<FName, FInworldDataHandle>& RuntimeData,
		FOnGraphResultNative ResultCallback = nullptr, FOnGraphStatusNative StatusCallback = nullptr);
	FString Execute(UObject* Instigator, const FInworldDataHandle& InputData, FOnGraphResultNative ResultCallback = nullptr,
		FOnGraphStatusNative StatusCallback = nullptr);

	/**
	 * @brief Cancels an ongoing graph execution
	 * @param ExecutionId The identifier of the execution to cancel
	 */
	UFUNCTION(BlueprintCallable, Category = "Inworld|Graph")
	void CancelExecution(const FString& ExecutionId);

	/**
	 * @brief Cancels all ongoing and pending executions
	 */
	UFUNCTION(BlueprintCallable, Category = "Inworld|Graph")
	void CancelAllExecutions();

	/**
	 * @brief Checks if an execution is canceled
	 * @param ExecutionId The identifier of the execution to check
	 * @return True if the execution is canceled, false otherwise
	 */
	UFUNCTION(BlueprintCallable, Category = "Inworld|Graph")
	bool IsExecutionCanceled(const FString& ExecutionId);

	/**
	 * @brief Gets the status of a specific execution
	 * @param ExecutionId The identifier of the execution
	 * @return The current status of the execution
	 */
	UFUNCTION(BlueprintPure, Category = "Inworld|Graph")
	EInworldGraphExecutionStatus GetExecutionStatus(const FString& ExecutionId);

	/**
	 * @brief Gets information about a specific execution
	 * @param ExecutionId The identifier of the execution
	 * @return The execution information, or empty struct if not found
	 */
	UFUNCTION(BlueprintPure, Category = "Inworld|Graph")
	FInworldGraphExecution GetExecutionInfo(const FString& ExecutionId);

	/**
	 * @brief Gets all current executions (running and pending)
	 * @return Array of all current executions
	 */
	UFUNCTION(BlueprintPure, Category = "Inworld|Graph")
	TArray<FInworldGraphExecution> GetAllExecutions();

	/**
	 * @brief Gets the number of currently running executions
	 * @return Number of running executions
	 */
	UFUNCTION(BlueprintPure, Category = "Inworld|Graph")
	int32 GetRunningExecutionsCount();

	/**
	 * @brief Gets the number of pending executions
	 * @return Number of pending executions
	 */
	UFUNCTION(BlueprintPure, Category = "Inworld|Graph")
	int32 GetPendingExecutionsCount();

	/**
	 * @brief Sets the execution mode for the graph
	 * @param Mode The execution mode to set
	 */
	UFUNCTION(BlueprintCallable, Category = "Inworld|Graph")
	void SetExecutionMode(EInworldGraphExecutionMode Mode);

	/**
	 * @brief Gets the current execution mode
	 * @return The current execution mode
	 */
	UFUNCTION(BlueprintPure, Category = "Inworld|Graph")
	EInworldGraphExecutionMode GetExecutionMode() const
	{
		return ExecutionMode;
	}

	/**
	 * @brief Retrieves the process context associated with a given execution ID
	 * @param ExecutionId The string identifier of the execution to retrieve the process context for
	 * @return Pointer to a UInworldProcessContext if found, nullptr if not
	 */
	UInworldProcessContext* GetProcessContext(const FString& ExecutionId) const
	{
		FScopeLock Lock(&ExecutionQueueLock);
		if (ExecutionMap.Contains(ExecutionId))
		{
			return ExecutionMap.Find(ExecutionId)->ProcessContext;
		}
		return nullptr;
	}

	/**
	 * @brief Retrieves the names of all nodes in the graph. This is mainly used as a UPROPERTY utility function.
	 * @return Array of node names as strings
	 */
	UFUNCTION()
	TArray<FString> GetNodeNames() const;

	/**
	 * @brief Checks if the graph has been compiled
	 * @return True if the graph is compiled, false otherwise
	 */
	UFUNCTION(BlueprintPure, Category = "Inworld|Graph")
	bool IsCompiled() const
	{
		return bIsCompiled;
	}

protected:
	/**
	 * @brief Updates the execution status of a given Inworld graph execution instance.
	 * Modifies the current execution's status to the specified status value, and invokes the status update event.
	 *
	 * @param Execution A pointer to the execution instance whose status is being updated.
	 * @param Status The new status to assign to the execution instance.
	 */
	static void UpdateExecutionStatus(FInworldGraphExecution* Execution, EInworldGraphExecutionStatus Status);

private:
	friend class UInworldGraphAsset;
	friend class UInworldRuntimeSubsystem;
#if WITH_EDITOR
	friend class UEdGraph_InworldGraph;
#endif

	void Compile(FOnGraphCompiledNative Callback);
	TOptional<FString> CompileInternal();
	bool GenerateConfig(TArray<FString>& Failed);

public:
	/** SetGraphId - set the name of the graph
	 * @brief Utility function for changing the graph's string identifier
	 */
	void SetGraphId(const FString& InGraphId);

	/** GetNodeByName - search for node with name
	 * @brief Utility function for grabbing a member node by its unique id
	 * @param NodeName unique id of the node to look for
	 * @param NodeClass The class to convert the node to on return
	 * @return Pointer to the node with the given name, or nullptr if not found
	 */
	UFUNCTION(BlueprintPure, Category = "Inworld|Graph", Meta = (DeterminesOutputType = "NodeClass"))
	UInworldNode* GetNodeByName(const FString& NodeName, TSubclassOf<UInworldNode> NodeClass) const
	{
		return GetNodeByName(NodeName);
	}
	UInworldNode* GetNodeByName(const FString& NodeName) const;

	/** CheckNodeByName - determine if we have a node with a given name
	 * @brief Utility function for determining if a node exists in the graph
	 * @param NodeName unique id of the node to look for
	 * @return true if a node with the name exists, false otherwise
	 */
	bool CheckNodeByName(const FString& NodeName) const;

	UPROPERTY(EditDefaultsOnly, Category = "Inworld|Graph", meta = (ToolTip = "Unique identifier for the graph instance"))
	FString GraphId{};

	UPROPERTY(EditDefaultsOnly, Instanced, Category = "Inworld|Graph",
		meta = (TitleProperty = "NodeName", ToolTip = "Collection of all nodes in the graph"))
	TArray<TObjectPtr<UInworldNode>> Nodes;

	UPROPERTY(EditDefaultsOnly, Instanced, Category = "Inworld|Graph",
		meta = (TitleProperty = "{Source} > {Destination}", ToolTip = "Collection of all edges connecting nodes in the graph"))
	TArray<TObjectPtr<UInworldEdge>> Edges;

	UPROPERTY(EditDefaultsOnly, Category = "Inworld|Graph",
		meta = (GetOptions = "GetNodeNames", ToolTip = "Name of the nodes where graph execution begins"))
	TArray<FString> StartNodes;

	UPROPERTY(EditDefaultsOnly, Category = "Inworld|Graph",
		meta = (GetOptions = "GetNodeNames", ToolTip = "Names of nodes where graph execution can terminate"))
	TArray<FString> EndNodes;

	/** Execution mode for this graph instance */
	UPROPERTY(EditAnywhere, BlueprintReadWrite, Category = "Inworld|Graph|Execution",
		meta = (ToolTip = "Defines how executions are handled: Simultaneous, LatestOnly, Sequential"))
	EInworldGraphExecutionMode ExecutionMode = EInworldGraphExecutionMode::Simultaneous;

	/*
	 * Temporary storage for all data output during graph execution
	 */
	UPROPERTY(Transient, Meta = (ToolTip = "Temporary storage for all data output during graph execution"))
	TArray<FInworldDataHandle> AllOutputData;

	UPROPERTY(Transient, Meta = (ToolTip = "JSON configuration string for the graph"))
	FString GraphConfig{};

private:
	UE::Tasks::TTask<void> CompileTask;
	TSharedPtr<inworld::ExecutorInterface> executor_;

	bool bIsCompiled{false};
	FThreadSafeBool bIsCompiling{false};
	mutable FCriticalSection CompileCallbacksLock;
	TArray<FOnGraphCompiledNative> PendingCompileCallbacks;

	/** Instance-level execution tracking */
	UPROPERTY(Transient)
	TMap<FString, FInworldGraphExecution> ExecutionMap;
	TArray<FString> ExecutionQueue;
	mutable FCriticalSection ExecutionQueueLock;

	/** Instance-level components registry **/
	TSet<FString> ComponentTypesRegistry;
	UPROPERTY(Transient)
	TArray<TObjectPtr<UInworldEdge>> CustomEdges;

	/** Internal functions for execution management */
	void RemoveExecutionFromPool(const FString& ExecutionId);
	void ProcessExecutionQueue();
	void StartExecution(const FString& ExecutionId);
	void CompleteExecution(const FString& ExecutionId, bool bSuccess);
	bool CanStartNewExecution() const;
	void HandleOneActiveMode(const FString& NewExecutionId);

	// for debugging purposes. keep track of 'who' is lauching this execution.
	void AddExecutionHistory(UObject* Instigator, const FString& ExecutionId, const TMap<FName, FInworldDataHandle>& RuntimeData);
};

struct FInworldKnowledgeCreationConfig;
struct FInworldLLMCreationConfig;
struct FInworldSTTCreationConfig;
struct FInworldTextEmbedderCreationConfig;
struct FInworldTTSCreationConfig;

namespace inworld
{
class CreationContext;
namespace graph
{
struct CompileContext
{
public:
	CompileContext(const UObject* InWorldContextObject, inworld::Graph& InGraph);

	const UObject* GetWorldContextObject() const
	{
		return WorldContextObject;
	}

	inworld::Graph& GetGraph() const
	{
		return graph;
	}

	template <typename TConfig>
	void Load(const FString& Id);
	template <typename TConfig>
	void Load(const FString& Id, const TConfig& Config);

	template <typename TInterface>
	TInterface Get(const FString& Id);

private:
	const UObject* WorldContextObject;
	inworld::Graph& graph;

	FCriticalSection ComponentMutexMapMutex;
	TMap<FString, TSharedPtr<FCriticalSection>> ComponentMutexMap;

	template <typename TConfig>
	void Create(const FString& Id, const TConfig& Config);
};

}	 // namespace graph
}	 // namespace inworld
