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

#include "Assets/InworldTextEmbeddingAssetDetailsCustomization.h"

#include "DetailCategoryBuilder.h"
#include "DetailLayoutBuilder.h"
#include "DetailWidgetRow.h"
#include "IPropertyUtilities.h"
#include "InworldRuntimeSettings.h"
#include "Primitives/TextEmbedder/InworldTextEmbedding.h"
#include "Widgets/Images/SThrobber.h"
#include "Widgets/Input/SButton.h"
#include "Widgets/Input/SComboBox.h"
#include "Widgets/Text/STextBlock.h"

#define LOCTEXT_NAMESPACE "FInworldTextEmbeddingAssetDetailsCustomization"

TSharedRef<IDetailCustomization> FInworldTextEmbeddingAssetDetailsCustomization::MakeInstance()
{
	return MakeShareable(new FInworldTextEmbeddingAssetDetailsCustomization());
}

void FInworldTextEmbeddingAssetDetailsCustomization::CustomizeDetails(IDetailLayoutBuilder& DetailBuilder)
{
	TArray<TWeakObjectPtr<UObject>> SelectedObjects;
	DetailBuilder.GetObjectsBeingCustomized(SelectedObjects);

	// We only support customizing one asset at a time for this operation
	if (SelectedObjects.Num() != 1)
	{
		return;
	}

	CustomizedAsset = Cast<UInworldTextEmbeddingAsset>(SelectedObjects[0].Get());
	if (!CustomizedAsset.IsValid())
	{
		return;
	}

	// Populate the options for our combobox
	GenerateConfigIdOptions();

	// Disable editing of the text property while embeddings are being generated
	{
		TSharedPtr<IPropertyHandle> InworldTextEmbeddingsHandle =
			DetailBuilder.GetProperty(GET_MEMBER_NAME_CHECKED(UInworldTextEmbeddingAsset, InworldTextEmbeddings));
		check(InworldTextEmbeddingsHandle.IsValid());

		IDetailPropertyRow* StructPropertyRow = DetailBuilder.EditDefaultProperty(InworldTextEmbeddingsHandle);

		if (StructPropertyRow)
		{
			StructPropertyRow->IsEnabled(TAttribute<bool>(this, &FInworldTextEmbeddingAssetDetailsCustomization::IsEditingEnabled));
		}
		else
		{
			UE_LOG(LogInworldTextEmbedder, Warning, TEXT("Failed to find property row for InworldTextEmbeddings"));
		}
	}

	// Cache the property handle for our selected config ID
	SelectedConfigIdHandle = DetailBuilder.GetProperty(GET_MEMBER_NAME_CHECKED(UInworldTextEmbeddingAsset, SelectedConfigId));
	check(SelectedConfigIdHandle.IsValid());

	// If no config is selected, try to set it to "Default" if that option exists
	if (CustomizedAsset->SelectedConfigId.IsEmpty())
	{
		for (const TSharedPtr<FString>& Option : ConfigIdOptions)
		{
			if (Option.IsValid() && *Option == TEXT("Default"))
			{
				SelectedConfigIdHandle->SetValue(TEXT("Default"));
				UE_LOG(LogInworldTextEmbedder, Log, TEXT("Asset '%s' has no config selected. Automatically selecting 'Default'."),
					*CustomizedAsset->GetName());

				break;
			}
		}
	}

	DetailBuilder.HideProperty(SelectedConfigIdHandle);

	IDetailCategoryBuilder& Category =
		DetailBuilder.EditCategory("Embedding Actions", FText::GetEmpty(), ECategoryPriority::Important);

	Category.AddCustomRow(FText::FromString(TEXT("Generate Embeddings")))
		.NameContent()[SNew(STextBlock).Text(FText::FromString(TEXT("Configuration"))).Font(IDetailLayoutBuilder::GetDetailFont())]
		.ValueContent()
			[SNew(SHorizontalBox)
				// ++ SComboBox Slot
				+
				SHorizontalBox::Slot().FillWidth(1.0f).VAlign(
					VAlign_Center)[SNew(SComboBox<TSharedPtr<FString>>)
									   .OptionsSource(&ConfigIdOptions)
									   .IsEnabled(this, &FInworldTextEmbeddingAssetDetailsCustomization::IsEditingEnabled)
									   .OnComboBoxOpening(this, &FInworldTextEmbeddingAssetDetailsCustomization::OnComboBoxOpening)
									   .OnSelectionChanged(this, &FInworldTextEmbeddingAssetDetailsCustomization::OnConfigIdChanged)
									   .OnGenerateWidget_Lambda(
										   [this](TSharedPtr<FString> InItem)
										   {
											   // Generate tooltip for each option in the dropdown list every time it's triggered to
											   // make sure it's up to date
											   const FString ConfigId = *InItem;
											   return SNew(STextBlock)
												   .Text(FText::FromString(ConfigId))
												   .ToolTipText(TAttribute<FText>::Create(TAttribute<FText>::FGetter::CreateLambda(
													   [this, ConfigId]() -> FText
													   {
														   return this->GetConfigIdTooltip(ConfigId);
													   })));
										   })[SNew(STextBlock)
												  .Text(this, &FInworldTextEmbeddingAssetDetailsCustomization::GetSelectedConfigId)
												  .ToolTipText_Lambda(	  // Generate tooltip when hovering over the combobox
																		  // without opening it
													  [this]() -> FText
													  {
														  const FString SelectedId = this->GetSelectedConfigId().ToString();

														  if (!SelectedId.IsEmpty() &&
															  ConfigIdOptions.ContainsByPredicate(
																  [&SelectedId](const TSharedPtr<FString>& Option)
																  {
																	  return *Option == SelectedId;
																  }))
														  {
															  return this->GetConfigIdTooltip(SelectedId);
														  }
														  return FText::GetEmpty();
													  })]]
				// ++ Button Slot
				+ SHorizontalBox::Slot()
					  .AutoWidth()
					  .Padding(4.0f, 0.0f)
					  .VAlign(VAlign_Center)[SNew(SButton)
												 .IsEnabled(this, &FInworldTextEmbeddingAssetDetailsCustomization::IsEditingEnabled)
												 .OnClicked(this,
													 &FInworldTextEmbeddingAssetDetailsCustomization::OnGenerateEmbeddingsClicked)
												 .Content()[SNew(STextBlock).Text(FText::FromString(TEXT("Generate Embeddings")))]]
				// ++ Throbber Slot
				+ SHorizontalBox::Slot()
					  .AutoWidth()
					  .Padding(4.0f, 0.0f)
					  .VAlign(VAlign_Center)[SNew(SThrobber).Visibility(
						  this, &FInworldTextEmbeddingAssetDetailsCustomization::GetThrobberVisibility)]];

	// Force a refresh to update the enabled state of the text property
	DetailBuilder.GetPropertyUtilities()->ForceRefresh();
}

bool FInworldTextEmbeddingAssetDetailsCustomization::IsEditingEnabled() const
{
	if (CustomizedAsset.IsValid())
	{
		return !CustomizedAsset->bIsGeneratingEmbeddings;
	}
	return false;
}

EVisibility FInworldTextEmbeddingAssetDetailsCustomization::GetThrobberVisibility() const
{
	if (CustomizedAsset.IsValid() && CustomizedAsset->bIsGeneratingEmbeddings)
	{
		return EVisibility::Visible;
	}
	return EVisibility::Collapsed;
}

FReply FInworldTextEmbeddingAssetDetailsCustomization::OnGenerateEmbeddingsClicked()
{
	if (CustomizedAsset.IsValid())
	{
		CustomizedAsset->StartEmbeddingGeneration();
	}
	return FReply::Handled();
}

void FInworldTextEmbeddingAssetDetailsCustomization::GenerateConfigIdOptions()
{
	ConfigIdOptions.Empty();
	const TArray<FString> ConfigIds = UInworldTextEmbeddingAsset::GetAvailableConfigIds();
	for (const FString& Id : ConfigIds)
	{
		ConfigIdOptions.Add(MakeShared<FString>(Id));
	}
}

void FInworldTextEmbeddingAssetDetailsCustomization::OnConfigIdChanged(
	TSharedPtr<FString> NewSelection, ESelectInfo::Type SelectInfo)
{
	if (CustomizedAsset.IsValid() && NewSelection.IsValid())
	{
		if (SelectInfo != ESelectInfo::OnNavigation)
		{
			SelectedConfigIdHandle->SetValue(*NewSelection);
		}
	}
}

FText FInworldTextEmbeddingAssetDetailsCustomization::GetSelectedConfigId() const
{
	if (CustomizedAsset.IsValid())
	{
		// If the current selection is not a valid option, display a hint
		bool bIsValidOption = false;
		for (const auto& Option : ConfigIdOptions)
		{
			// If the current selection is not a valid option, display a hint
			if (Option.IsValid() && *Option == CustomizedAsset->SelectedConfigId)
			{
				bIsValidOption = true;
				break;
			}
		}

		if (bIsValidOption)
		{
			return FText::FromString(CustomizedAsset->SelectedConfigId);
		}

		if (!CustomizedAsset->SelectedConfigId.IsEmpty())
		{
			FText InvalidTextFormat = LOCTEXT("InvalidConfigIdFormat", "{0} (Invalid)");
			return FText::Format(InvalidTextFormat, FText::FromString(CustomizedAsset->SelectedConfigId));
		}
	}

	return LOCTEXT("SelectConfigIdHint", "Select Configuration...");
}

FText FInworldTextEmbeddingAssetDetailsCustomization::GetConfigIdTooltip(FString ConfigId) const
{
	const UInworldRuntimeSettings* Settings = GetDefault<UInworldRuntimeSettings>();
	if (Settings)
	{
		if (const FInworldTextEmbedderCreationConfig* Config = Settings->TextEmbedderCreationConfig.Find(ConfigId))
		{
			if (Config->ComputeHost == EInworldComputeHost::REMOTE)
			{
				return FText::FromString(FString::Printf(TEXT("Host: Remote\nProvider: %s\nModel: %s"),
					*Config->RemoteCreationConfig.Provider, *Config->RemoteCreationConfig.Model));
			}
			else	// Handle LOCAL case
			{
				FString ComputeDevice =
					(Config->LocalCreationConfig.LocalComputeHost == EInworldLocalComputeHost::CPU) ? TEXT("CPU") : TEXT("GPU");

				return FText::FromString(
					FString::Printf(TEXT("Host: Local (%s)\nModel: %s"), *ComputeDevice, *Config->LocalCreationConfig.Model));
			}
		}
	}
	return FText::FromString(TEXT("Configuration details not found."));
}

void FInworldTextEmbeddingAssetDetailsCustomization::OnComboBoxOpening()
{
	GenerateConfigIdOptions();
}

#undef LOCTEXT_NAMESPACE
