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

#pragma once
#include "Graph/InworldData.h"
#include "Graph/InworldGraphRuntimeData.h"

#include "InworldProcessContext.generated.h"

namespace inworld
{
class ProcessContext;
class StatusOr_TextEmbedderInterface;
}	 // namespace inworld

/**
 * @class UInworldProcessContext
 * @brief Context object for managing runtime data during graph execution.
 *
 * This class provides access to runtime data that can be shared across nodes
 * during graph execution. It allows nodes to store and retrieve data using
 * key-value pairs, supporting various data types including basic types,
 * objects, and structures.
 */
UCLASS(BlueprintType)
class INWORLDRUNTIME_API UInworldProcessContext : public UObject
{
	GENERATED_BODY()

public:
	static UInworldProcessContext* Create(UObject* Owner, const TMap<FName, FInworldDataHandle>& GraphRuntimeData);
	static void SetProcessContext(inworld::ProcessContext* in_process_context);
	static inworld::ProcessContext* GetProcessContext();

	inworld::StatusOr_TextEmbedderInterface GetTextEmbedder(const FString& Id) const;

	UFUNCTION(BlueprintCallable, Category = "Inworld|ProcessContext")
	void SetRuntimeData(FName Key, FInworldDataHandle Data)
	{
		RuntimeDataMap.Add(Key, Data);
	}
	UFUNCTION(BlueprintCallable, Category = "Inworld|ProcessContext")
	bool GetRuntimeData(FName Key, FInworldDataHandle& OutData) const;

	UFUNCTION(BlueprintCallable, Category = "Inworld|ProcessContext")
	bool GetRuntimeBool(FName Key, bool& OutBool) const;

	UFUNCTION(BlueprintCallable, Category = "Inworld|ProcessContext")
	bool GetRuntimeInt(FName Key, int32& OutInt) const;

	UFUNCTION(BlueprintCallable, Category = "Inworld|ProcessContext")
	bool GetRuntimeFloat(FName Key, float& OutFloat) const;

	UFUNCTION(BlueprintCallable, Category = "Inworld|ProcessContext",
		meta = (DeterminesOutputType = "ObjectClass", DynamicOutputParam = "OutObject"))
	bool GetRuntimeObject(FName Key, TSubclassOf<UObject> ObjectClass, UObject*& OutObject) const;

	template <typename T, typename = std::enable_if_t<std::is_pointer_v<T> && std::is_base_of_v<UObject, std::remove_pointer_t<T>>>>
	T GetRuntimeObject(FName Key) const
	{
		UObject* OutObject{nullptr};
		const bool bSuccess = GetRuntimeObject(Key, std::remove_pointer_t<T>::StaticClass(), OutObject);
		if (bSuccess)
		{
			return Cast<std::remove_pointer_t<T>>(OutObject);
		}
		return nullptr;
	}

	template <typename T>
	std::enable_if_t<std::is_base_of_v<FInworldData, T>, bool> GetRuntimeStruct(FName Key, T& OutData) const
	{
		if (RuntimeDataMap.Contains(Key))
		{
			const FInworldDataHandle& RuntimeData{RuntimeDataMap[Key]};
			if (const T* Data = RuntimeData.Unwrap<T>(); Data != nullptr)
			{
				OutData = *Data;
				return true;
			}
		}
		return false;
	}

	template <typename T>
	std::enable_if_t<!std::is_base_of_v<FInworldData, T>, bool> GetRuntimeStruct(const FName& Key, T& OutData) const
	{
		if (RuntimeDataMap.Contains(Key))
		{
			const FInworldDataHandle& RuntimeData{RuntimeDataMap[Key]};
			if (FInworldData_Struct* Struct = RuntimeData.Unwrap<FInworldData_Struct>();
				Struct != nullptr && Struct->GetData(OutData))
			{
				return true;
			}
		}
		return false;
	}

private:
	TMap<FName, FInworldDataHandle> RuntimeDataMap;
	friend class UInworldBlueprintFunctionLibrary;
};
