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

#include "Tools/InworldTextEmbedder.h"

#include "Async/Async.h"
#include "Async/TaskGraphInterfaces.h"
#include "InworldRuntimeSettings.h"
#include "Primitives/TextEmbedder/InworldTextEmbedderTasks.h"
#include "Primitives/TextEmbedder/InworldTextEmbedding.h"

THIRD_PARTY_INCLUDES_START
#include "InworldRuntimeLibrary/inworld.h"
THIRD_PARTY_INCLUDES_END

#define LOCTEXT_NAMESPACE "InworldTextEmbedder"

UInworldTextEmbedder* UInworldTextEmbedder::CreateInworldTextEmbedder(const FString& ConfigId)
{
	UInworldTextEmbedder* TextEmbedder = NewObject<UInworldTextEmbedder>();
	const UInworldRuntimeSettings* Settings = GetDefault<UInworldRuntimeSettings>();

	if (const FInworldTextEmbedderCreationConfig* CreationConfig = Settings->TextEmbedderCreationConfig.Find(ConfigId))
	{
		inworld::TextEmbedderCreationConfig config;
		if (CreationConfig->ComputeHost == EInworldComputeHost::LOCAL)
		{
			inworld::LocalTextEmbedderConfig localConfig;
			inworld::Convert(CreationConfig->LocalCreationConfig, localConfig);
			config.set_local_config(localConfig);
		}
		else
		{
			inworld::RemoteTextEmbedderConfig remoteConfig;
			inworld::Convert(CreationConfig->RemoteCreationConfig, remoteConfig);
			config.set_remote_config(remoteConfig);
		}
		inworld::StatusOr_TextEmbedderInterface statusOr_TextEmbedderInterface =
			inworld::TextEmbedderFactory().CreateTextEmbedder(config);

		if (statusOr_TextEmbedderInterface.ok())
		{
			TextEmbedder->TextEmbedderInterface =
				MakeShared<inworld::TextEmbedderInterface>(statusOr_TextEmbedderInterface.value());
			return TextEmbedder;
		}
		else
		{
			inworld::Status status = statusOr_TextEmbedderInterface.status();
			FString ErrorMessage = FString(UTF8_TO_TCHAR(status.ToString().c_str()));
			UE_LOG(LogInworldTextEmbedder, Error, TEXT("Failed to create Text Embedder Interface. Status: %s"), *ErrorMessage);
		}
	}
	else
	{
		UE_LOG(LogInworldTextEmbedder, Error, TEXT("Failed to find Text Embedder creation config: %s."), *ConfigId);
	}
	UE_LOG(LogInworldTextEmbedder, Error, TEXT("Failed to create Text Embedder."));
	return nullptr;
}

void UInworldTextEmbedder::EmbedAsync(const TArray<FString>& Text, FOnTextEmbedded Callback)
{
	FOnTextEmbeddedNative NativeCallback;
	if (Callback.IsBound())
	{
		NativeCallback.BindUFunction(Callback.GetUObject(), Callback.GetFunctionName());
	}
	EmbedAsync(Text, NativeCallback);
}

void UInworldTextEmbedder::EmbedAsync(const TArray<FString>& Text, FOnTextEmbeddedNative Callback) const
{
	TSharedPtr<inworld::TextEmbedderInterface> TempTextEmbedderInterface = TextEmbedderInterface;

	UE::Tasks::Launch(UE_SOURCE_LOCATION,
		[TempTextEmbedderInterface, Text, Callback]()
		{
			UE::Tasks::TTask<TOptional<TArray<FInworldTextEmbedding>>> TextEmbeddingTask =
				inworld::Tasks::TextEmbedder::EmbedTask(TempTextEmbedderInterface, Text);

			TOptional<TArray<FInworldTextEmbedding>> TextEmbeddings = TextEmbeddingTask.GetResult();
			AsyncTask(ENamedThreads::GameThread,
				[TextEmbeddings, Callback]()
				{
					if (TextEmbeddings.IsSet())
					{
						Callback.ExecuteIfBound(true, TextEmbeddings.GetValue());
					}
					else
					{
						Callback.ExecuteIfBound(false, {});
					}
				});
		});
}

bool UInworldTextEmbedder::Embed(const TArray<FString>& Text, TArray<FInworldTextEmbedding>& OutTextEmbeddings)
{
	UE::Tasks::TTask<TOptional<TArray<FInworldTextEmbedding>>> TextEmbeddingTask =
		inworld::Tasks::TextEmbedder::EmbedTask(TextEmbedderInterface, Text);

	TOptional<TArray<FInworldTextEmbedding>> TextEmbeddings = TextEmbeddingTask.GetResult();
	if (TextEmbeddings.IsSet())
	{
		OutTextEmbeddings = TextEmbeddings.GetValue();
		return true;
	}
	OutTextEmbeddings = {};
	return false;
}

#undef LOCTEXT_NAMESPACE
