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

#include "Primitives/TextEmbedder/InworldTextEmbedding.h"

#if WITH_EDITOR
#include "Editor.h"
#include "InworldRuntimeSettings.h"
#include "InworldRuntimeSubsystem.h"
#include "Misc/MessageDialog.h"
#include "Subsystems/EditorAssetSubsystem.h"
#include "Tools/InworldTextEmbedder.h"
#include "UObject/ObjectSaveContext.h"

#define LOCTEXT_NAMESPACE "InworldTextEmbeddingAsset"

void UInworldTextEmbeddingAsset::PostLoad()
{
	Super::PostLoad();

	// TODO: Remove in future release
	// Allows for old assets to update easily
	if (!InworldTextEmbeddings.Text.IsEmpty())
	{
		for (const FString& Text : InworldTextEmbeddings.Text)
		{
			TextEmbeddings.Add({Text, {}, {}});
		}
		InworldTextEmbeddings = {};
	}

	CachedTextForLastEmbedding = GetSourceTextArray();
}

void UInworldTextEmbeddingAsset::PostEditChangeProperty(FPropertyChangedEvent& PropertyChangedEvent)
{
	Super::PostEditChangeProperty(PropertyChangedEvent);
	const FName PropertyName = (PropertyChangedEvent.Property != nullptr) ? PropertyChangedEvent.Property->GetFName() : NAME_None;
	if (PropertyName == GET_MEMBER_NAME_CHECKED(FInworldTextEmbedding, Text) ||
		PropertyName == GET_MEMBER_NAME_CHECKED(UInworldTextEmbeddingAsset, SelectedConfigId))
	{
		if (GetSourceTextArray() != CachedTextForLastEmbedding || SelectedConfigId != CachedConfigIdForLastEmbedding)
		{
			ClearCalculatedData();
			MarkPackageDirty();
		}
	}
}

void UInworldTextEmbeddingAsset::PreSave(FObjectPreSaveContext SaveContext)
{
	Super::PreSave(SaveContext);

	if (SelectedConfigId.IsEmpty())
	{
		FText DialogTitle = LOCTEXT("ConfigRequiredTitle", "Configuration Required");
		FText DialogText = FText::Format(
			LOCTEXT("ConfigRequiredMessage_Save", "A configuration must be selected for asset '{0}' before it can be saved."),
			FText::FromString(this->GetName()));

		FMessageDialog::Open(EAppMsgType::Ok, DialogText, DialogTitle);

		if (UPackage* Package = GetOutermost())
		{
			Package->SetDirtyFlag(false);
		}
		return;
	}

	if (bIsGeneratingEmbeddings)
	{
		FText DialogTitle = FText::FromString(TEXT("Save Blocked"));
		FText DialogText = FText::Format(
			LOCTEXT("SaveBlockedMessage",
				"Embedding generation is in progress for asset '{0}'.\n\nPlease wait for it to complete before saving."),
			FText::FromString(this->GetName()));

		FMessageDialog::Open(EAppMsgType::Ok, DialogText, DialogTitle);
		return;
	}

	if (!bHasCalculatedData)
	{
		FText DialogTitle = FText::FromString(TEXT("Generate New Embeddings?"));
		FText DialogText = FText::Format(LOCTEXT("GenerateNewEmbeddingsMessage",
											 "The text in asset '{0}' has been modified, but new embeddings have not been "
											 "generated.\n\nWould you like to generate new embeddings and save the asset?\n"),
			FText::FromString(this->GetName()));

		EAppReturnType::Type DialogResult = FMessageDialog::Open(EAppMsgType::OkCancel, DialogText, DialogTitle);

		switch (DialogResult)
		{
			case EAppReturnType::Ok:
			{
				GenerateEmbeddingsWithCallback(
					[this](bool bSuccess)
					{
						if (bSuccess)
						{
							UE_LOG(LogInworldTextEmbedder, Log, TEXT("Asset '%s': Auto-saving after embedding generation."),
								*this->GetName());

							if (UEditorAssetSubsystem* EditorAssetSubsystem = GEditor->GetEditorSubsystem<UEditorAssetSubsystem>();
								EditorAssetSubsystem == nullptr || !EditorAssetSubsystem->SaveAsset(this->GetPathName()))
							{
								UE_LOG(LogInworldTextEmbedder, Error,
									TEXT("Asset '%s': Auto-save FAILED after embedding generation. Re-dirtying package."),
									*this->GetName());
								this->MarkPackageDirty();
							}
						}
					});

				// Cancel the current synchronous save. The async callback will handle saving the asset later.
				if (UPackage* Package = GetOutermost())
				{
					Package->SetDirtyFlag(false);
				}
				break;
			}
			case EAppReturnType::Cancel:
			{
				// Cancel the current save action, but re-dirty the package on the next tick
				// so the asset still appears as modified to the user.
				if (UPackage* Package = GetOutermost())
				{
					Package->SetDirtyFlag(false);
				}

				// Use a one-time ticker to mark the package dirty on the next tick
				FTSTicker::GetCoreTicker().AddTicker(FTickerDelegate::CreateLambda(
					[WeakThis = TWeakObjectPtr<UInworldTextEmbeddingAsset>(this)](float DeltaTime) -> bool
					{
						if (UInworldTextEmbeddingAsset* Asset = WeakThis.Get())
						{
							Asset->MarkPackageDirty();
						}

						return false;
					}));

				break;
			}
		}
	}
}

void UInworldTextEmbeddingAsset::StartEmbeddingGeneration()
{
	GenerateEmbeddingsWithCallback(
		[](bool bSuccess)
		{
		});
}

void UInworldTextEmbeddingAsset::GenerateEmbeddingsWithCallback(TFunction<void(bool bSuccess)> OnComplete)
{
	if (GetSourceTextArray() == CachedTextForLastEmbedding && SelectedConfigId == CachedConfigIdForLastEmbedding &&
		!TextEmbeddings.IsEmpty() && bHasCalculatedData)
	{
		UE_LOG(LogInworldTextEmbedder, Log,
			TEXT("Asset '%s': Text has not changed since last successful embedding. Skipping generation."), *this->GetName());
		OnComplete(true);
		return;
	}

	if (bIsGeneratingEmbeddings)
	{
		UE_LOG(LogInworldTextEmbedder, Warning, TEXT("Asset '%s': Embedding generation is already in progress."), *this->GetName());
		OnComplete(false);
		return;
	}

	bIsGeneratingEmbeddings = true;
	UE_LOG(LogInworldTextEmbedder, Log, TEXT("Asset '%s': Starting embedding generation..."), *this->GetName());

	if (SelectedConfigId.IsEmpty())
	{
		FText DialogTitle = LOCTEXT("ConfigRequiredTitle", "Configuration Required");
		FText DialogText = FText::Format(LOCTEXT("ConfigRequiredMessage_Generate",
											 "Please select a configuration for asset '{0}' before generating embeddings."),
			FText::FromString(this->GetName()));

		FMessageDialog::Open(EAppMsgType::Ok, DialogText, DialogTitle);

		UE_LOG(LogInworldTextEmbedder, Error,
			TEXT("Asset '%s': No embedder configuration selected. Please select a configuration in the Details panel before "
				 "generating."),
			*this->GetName());
		bIsGeneratingEmbeddings = false;
		OnComplete(false);
		return;
	}

	UInworldRuntimeSubsystem* RuntimeSubsystem = UInworldRuntimeSubsystem::Get();
	if (!RuntimeSubsystem)
	{
		UE_LOG(LogInworldTextEmbedder, Error, TEXT("Failed to get InworldRuntimeSubsystem."));
		bIsGeneratingEmbeddings = false;
		OnComplete(false);
		return;
	}

	UInworldTextEmbedder* TextEmbedder = RuntimeSubsystem->GetOrCreateTextEmbedder(SelectedConfigId);

	if (!TextEmbedder)
	{
		UE_LOG(LogInworldTextEmbedder, Error, TEXT("Failed to create InworldTextEmbedder."));
		bIsGeneratingEmbeddings = false;
		OnComplete(false);
		return;
	}

	// Capture the text state before the async operation to handle potential race conditions.
	const TArray<FString> TextToEmbed = GetSourceTextArray();

	// Protect the local TextEmbedder from being garbage collected during the async operation.
	TextEmbedder->AddToRoot();

	TWeakObjectPtr<UInworldTextEmbeddingAsset> WeakThis = this;
	TextEmbedder->EmbedAsync(TextToEmbed,
		FOnTextEmbeddedNative::CreateLambda(
			[WeakThis, OnComplete, TextEmbedder, TextToEmbed](bool bSuccess, const TArray<FInworldTextEmbedding>& ResultEmbeddings)
			{
				// Always remove from root once the operation is complete to prevent memory leaks.
				if (TextEmbedder)
				{
					TextEmbedder->RemoveFromRoot();
				}

				UInworldTextEmbeddingAsset* Asset = WeakThis.Get();
				if (!Asset)
				{
					OnComplete(false);
					return;
				}

				// Stop the UI throbber regardless of success or failure
				Asset->bIsGeneratingEmbeddings = false;

				if (bSuccess)
				{
					// Race condition check: discard results if the source text has changed since we started.
					if (Asset->GetSourceTextArray() != TextToEmbed)
					{
						UE_LOG(LogInworldTextEmbedder, Warning,
							TEXT("Asset '%s': Text changed during embedding generation. Discarding stale results."),
							*Asset->GetName());
						OnComplete(false);
						return;
					}

					UE_LOG(LogInworldTextEmbedder, Log, TEXT("Asset '%s' re-embedded successfully."), *Asset->GetName());

					Asset->TextEmbeddings = ResultEmbeddings;

					// Update the cache with the version of the text and config we just successfully processed.
					Asset->CachedTextForLastEmbedding = TextToEmbed;
					Asset->CachedConfigIdForLastEmbedding = Asset->SelectedConfigId;
					Asset->bHasCalculatedData = true;
				}
				else
				{
					UE_LOG(LogInworldTextEmbedder, Error, TEXT("Failed to re-embed text for asset '%s'."), *Asset->GetName());
					Asset->ClearCalculatedData();
				}

				Asset->MarkPackageDirty();
				OnComplete(bSuccess);
			}));
}

TArray<FString> UInworldTextEmbeddingAsset::GetAvailableConfigIds()
{
	TArray<FString> Keys;
	const UInworldRuntimeSettings* Settings = GetDefault<UInworldRuntimeSettings>();
	if (Settings)
	{
		Settings->TextEmbedderCreationConfig.GetKeys(Keys);
	}
	return Keys;
}

TArray<FString> UInworldTextEmbeddingAsset::GetSourceTextArray() const
{
	TArray<FString> SourceText;
	Algo::Transform(TextEmbeddings, SourceText,
		[](const FInworldTextEmbedding& TextEmbedding)
		{
			return TextEmbedding.Text;
		});
	return SourceText;
}

void UInworldTextEmbeddingAsset::ClearCalculatedData()
{
	for (FInworldTextEmbedding& TextEmbedding : TextEmbeddings)
	{
		TextEmbedding.NormalizedText = {};
		TextEmbedding.Values = {};
	}
	bHasCalculatedData = false;
}

#undef LOCTEXT_NAMESPACE

#endif	  // WITH_EDITOR
