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

#include "Assets/InworldEditorTextEmbedder.h"

#include "AssetRegistry/AssetRegistryModule.h"
#include "Assets/InworldTextEmbeddingsFactory.h"
#include "Async/Async.h"
#include "ContentBrowserModule.h"
#include "IContentBrowserSingleton.h"
#include "InworldRuntimeSettings.h"
#include "Kismet2/BlueprintEditorUtils.h"
#include "Primitives/TextEmbedder/InworldTextEmbedderTasks.h"
#include "Primitives/TextEmbedder/InworldTextEmbedding.h"
#include "Tasks/Task.h"

#define LOCTEXT_NAMESPACE "InworldEditorTextEmbedder"

UInworldEditorTextEmbedder* UInworldEditorTextEmbedder::CreateInworldEditorTextEmbedder(const FString& ConfigId)
{
	UInworldEditorTextEmbedder* TextEmbedder = NewObject<UInworldEditorTextEmbedder>();
	TextEmbedder->TextEmbeddingsFactory = NewObject<UInworldTextEmbeddingsFactory>();
	TextEmbedder->TextEmbedder = UInworldTextEmbedder::CreateInworldTextEmbedder(ConfigId);
	TextEmbedder->CachedConfigId = ConfigId;
	return TextEmbedder;
}

void UInworldEditorTextEmbedder::EmbedAsync(const TArray<FString>& Text)
{
	if (TextEmbedder != nullptr)
	{
		TextEmbedder->EmbedAsync(Text, FOnTextEmbeddedNative::CreateUObject(this, &UInworldEditorTextEmbedder::OnTextEmbedded));
	}
	else
	{
		UE_LOG(LogInworldTextEmbedder, Error, TEXT("Failed to embed text: Text Embedder is null."));
	}
}

void UInworldEditorTextEmbedder::OnTextEmbedded(bool bSuccess, const TArray<FInworldTextEmbedding>& TextEmbeddings) const
{
	if (bSuccess)
	{
		UInworldTextEmbeddingAsset* TextEmbeddingAsset = CreateTextEmbeddingsAsset();
		if (TextEmbeddingAsset)
		{
			TextEmbeddingAsset->TextEmbeddings = TextEmbeddings;
			TextEmbeddingAsset->SelectedConfigId = CachedConfigId;
			TextEmbeddingAsset->CachedConfigIdForLastEmbedding = CachedConfigId;
			TextEmbeddingAsset->CachedTextForLastEmbedding = TextEmbeddingAsset->GetSourceTextArray();
		}
		else
		{
			UE_LOG(LogInworldTextEmbedder, Error, TEXT("Failed to create Text Embedding asset."));
		}
	}
}

void UInworldEditorTextEmbedder::ParseTextArrayString(const FString& TextArrayString, FOnTextArrayStringParsed Callback)
{
	UE::Tasks::Launch(UE_SOURCE_LOCATION,
		[TextArrayString, Callback]()
		{
			UE::Tasks::TTask<TArray<FString>> ParsedTexts = inworld::Tasks::TextEmbedder::ParseTextArrayString(TextArrayString);
			TArray<FString> ParsedText = ParsedTexts.GetResult();

			AsyncTask(ENamedThreads::GameThread,
				[ParsedText, Callback]()
				{
					if (ParsedText.Num() > 0)
					{
						Callback.ExecuteIfBound(true, ParsedText);
					}
					else
					{
						UE_LOG(LogInworldTextEmbedder, Error, TEXT("Failed to parse batch, no text entries found."));
						Callback.ExecuteIfBound(false, {});
					}
				});
		});
}

UInworldTextEmbeddingAsset* UInworldEditorTextEmbedder::CreateTextEmbeddingsAsset() const
{
	const FContentBrowserModule& ContentBrowserModule = FModuleManager::LoadModuleChecked<FContentBrowserModule>("ContentBrowser");
	FContentBrowserItemPath Path = ContentBrowserModule.Get().GetCurrentPath();

	UClass* Class = UInworldTextEmbeddingAsset::StaticClass();
	const FString ClassName = FBlueprintEditorUtils::GetClassNameWithoutSuffix(Class);

	FSaveAssetDialogConfig SaveAssetDialogConfig;
	SaveAssetDialogConfig.DialogTitleOverride = LOCTEXT("SaveAssetDialogTitle", "Save Asset As");
	if (Path.HasInternalPath())
	{
		SaveAssetDialogConfig.DefaultPath = Path.GetInternalPathString();
	}
	SaveAssetDialogConfig.DefaultAssetName = FString::Printf(TEXT("Inworld_Text_Embedding_New"));
	SaveAssetDialogConfig.ExistingAssetPolicy = ESaveAssetDialogExistingAssetPolicy::AllowButWarn;
	SaveAssetDialogConfig.AssetClassNames.Add(FTopLevelAssetPath(Class->GetPathName()));

	const FString SaveObjectPath = ContentBrowserModule.Get().CreateModalSaveAssetDialog(SaveAssetDialogConfig);
	if (!SaveObjectPath.IsEmpty())
	{
		const FString SavePackageName = FPackageName::ObjectPathToPackageName(SaveObjectPath);
		const FString SavePackagePath = FPaths::GetPath(SavePackageName);
		const FString SaveAssetName = FPaths::GetBaseFilename(SavePackageName);

		UPackage* Package = CreatePackage(*SavePackageName);
		if (ensure(Package))
		{
			UObject* NewAsset = TextEmbeddingsFactory->FactoryCreateNew(UInworldTextEmbeddingAsset::StaticClass(), Package,
				FName(*SaveAssetName), RF_Public | RF_Transactional | RF_Standalone, nullptr, GWarn);

			if (UInworldTextEmbeddingAsset* DefaultTextEmbeddings = Cast<UInworldTextEmbeddingAsset>(NewAsset))
			{
				FAssetRegistryModule::AssetCreated(NewAsset);
				Package->MarkPackageDirty();

				UE_LOG(LogInworldTextEmbedder, Log, TEXT("New Inworld Text Embedding Created."));

				return DefaultTextEmbeddings;
			}
		}
	}
	return nullptr;
}
#undef LOCTEXT_NAMESPACE
