UShaderLab/Source/UShaderLabEditor/Private/ShaderLabMaterialInstanceFactory.cpp

// Copyright FlecsProj. All Rights Reserved.

#include "ShaderLabMaterialInstanceFactory.h"

#include "AssetRegistry/AssetData.h"
#include "AssetThumbnail.h"
#include "AssetTypeCategories.h"
#include "Editor.h"
#include "Materials/Material.h"
#include "Misc/App.h"
#include "ShaderCompiler.h"
#include "ShaderLabMaterialInstanceConstant.h"
#include "ShaderLabMaterialRegistry.h"
#include "ShaderLabModel.h"
#include "Widgets/Input/SButton.h"
#include "Widgets/Input/SSearchBox.h"
#include "Widgets/Layout/SBorder.h"
#include "Widgets/Layout/SBox.h"
#include "Widgets/SBoxPanel.h"
#include "Widgets/SWindow.h"
#include "Widgets/Text/STextBlock.h"
#include "Widgets/Views/SListView.h"
#include "Widgets/Views/STableRow.h"

#define LOCTEXT_NAMESPACE "ShaderLabMaterialInstanceFactory"

namespace
{
	constexpr int32 GThumbnailSize = 64;

	/** One row in the base-material picker: the in-memory base + display strings + a lazily-made thumbnail. */
	struct FShaderLabBaseEntry
	{
		UMaterial* Material = nullptr;
		FString Name;
		/** "Surface / Opaque" etc., from the parsed model settings (shown as the row subtitle). */
		FString DomainBlend;
		/** Live thumbnail, created on first display (on-demand) and reused across row recycling. */
		TSharedPtr<FAssetThumbnail> Thumbnail;
	};
	using FBaseEntryPtr = TSharedPtr<FShaderLabBaseEntry>;

	FString DomainToString(EShaderLabDomain Domain)
	{
		switch (Domain)
		{
		case EShaderLabDomain::Surface:     return TEXT("Surface");
		case EShaderLabDomain::PostProcess: return TEXT("PostProcess");
		case EShaderLabDomain::UI:          return TEXT("UI");
		case EShaderLabDomain::Decal:       return TEXT("Decal");
		}
		checkNoEntry();
		return FString();
	}

	FString BlendToString(EShaderLabBlendMode Blend)
	{
		switch (Blend)
		{
		case EShaderLabBlendMode::Opaque:      return TEXT("Opaque");
		case EShaderLabBlendMode::Masked:      return TEXT("Masked");
		case EShaderLabBlendMode::Translucent: return TEXT("Translucent");
		case EShaderLabBlendMode::Additive:    return TEXT("Additive");
		case EShaderLabBlendMode::Modulate:    return TEXT("Modulate");
		}
		checkNoEntry();
		return FString();
	}

	/**
	 * Modal picker over the discovered ShaderLab base materials, modelled on the engine's class-picker
	 * dialog: a search box above a flat, single-selection list with a live preview thumbnail per row.
	 *
	 * Thumbnails render asynchronously and on-demand: each base is a shader-map-less /Script shell, so
	 * MaterialInstanceThumbnailRenderer renders it via a child instance whose shader map compiles in the
	 * background; SListView only realizes visible rows, so only those compile/render (scroll-to-compile).
	 *
	 * The catch: FSlateApplication::AddModalWindow's loop ticks Slate but NOT the editor engine, so the
	 * FAssetThumbnailPool (a FTickableEditorObject) and shader-compile finalization never advance on
	 * their own here — a stock live thumbnail would freeze on its first black frame. We therefore drive
	 * both from a Slate active timer (which DOES fire inside the modal loop): see PumpThumbnails.
	 */
	class SShaderLabBasePicker : public SCompoundWidget
	{
	public:
		SLATE_BEGIN_ARGS(SShaderLabBasePicker) {}
			SLATE_ARGUMENT(TSharedPtr<SWindow>, ParentWindow)
		SLATE_END_ARGS()

		void Construct(const FArguments& InArgs)
		{
			ParentWindow = InArgs._ParentWindow;
			check(ParentWindow.IsValid());

			// Local pool: we tick it ourselves (the editor won't, inside the modal loop).
			ThumbnailPool = MakeShared<FAssetThumbnailPool>(/*InNumInPool*/ 32);

			BuildEntries();
			FilteredEntries = AllEntries;

			ChildSlot
			[
				SNew(SVerticalBox)

				// Search box.
				+ SVerticalBox::Slot().AutoHeight().Padding(8.f, 8.f, 8.f, 4.f)
				[
					SNew(SSearchBox)
					.HintText(LOCTEXT("SearchHint", "Search base materials..."))
					.OnTextChanged(this, &SShaderLabBasePicker::OnSearchTextChanged)
				]

				// Scrollable list of bases.
				+ SVerticalBox::Slot().FillHeight(1.f).Padding(8.f, 4.f)
				[
					SNew(SBorder)
					.BorderImage(FAppStyle::GetBrush("ToolPanel.GroupBorder"))
					.Padding(2.f)
					[
						SAssignNew(ListView, SListView<FBaseEntryPtr>)
						.ListItemsSource(&FilteredEntries)
						.SelectionMode(ESelectionMode::Single)
						.OnGenerateRow(this, &SShaderLabBasePicker::OnGenerateRow)
						.OnMouseButtonDoubleClick(this, &SShaderLabBasePicker::OnItemDoubleClicked)
						.OnSelectionChanged(this, &SShaderLabBasePicker::OnSelectionChanged)
					]
				]

				// Select / Cancel.
				+ SVerticalBox::Slot().AutoHeight().HAlign(HAlign_Right).Padding(8.f)
				[
					SNew(SHorizontalBox)
					+ SHorizontalBox::Slot().AutoWidth().Padding(0, 0, 4, 0)
					[
						SNew(SButton)
						.Text(LOCTEXT("Select", "Select"))
						.IsEnabled(this, &SShaderLabBasePicker::IsSelectEnabled)
						.OnClicked(this, &SShaderLabBasePicker::OnSelectClicked)
					]
					+ SHorizontalBox::Slot().AutoWidth()
					[
						SNew(SButton)
						.Text(LOCTEXT("Cancel", "Cancel"))
						.OnClicked(this, &SShaderLabBasePicker::OnCancelClicked)
					]
				]
			];

			// Drive shader-compile finalization + thumbnail rendering every frame for the modal's life.
			RegisterActiveTimer(0.f, FWidgetActiveTimerDelegate::CreateSP(this, &SShaderLabBasePicker::PumpThumbnails));
		}

		/** The chosen base material, or nullptr if the user cancelled / closed the window. */
		UMaterial* GetChosen() const { return Chosen; }

	private:
		void BuildEntries()
		{
			FShaderLabMaterialRegistry& Registry = FShaderLabMaterialRegistry::Get();
			const TMap<FName, FShaderLabModel>& Models = Registry.GetRegisteredModels();

			// Sort by name for a stable, flat layout.
			TArray<FName> Names;
			Models.GenerateKeyArray(Names);
			Names.Sort([](const FName& A, const FName& B) { return A.Compare(B) < 0; });

			for (const FName& Name : Names)
			{
				UMaterial* Base = Registry.FindMaterial(Name.ToString());
				check(Base); // Registered models always have a created material (RegisterFromModel contract).
				const FShaderLabModel& Model = Models.FindChecked(Name);

				FBaseEntryPtr Entry = MakeShared<FShaderLabBaseEntry>();
				Entry->Material = Base;
				Entry->Name = Base->GetName();
				Entry->DomainBlend = FString::Printf(TEXT("%s / %s"),
					*DomainToString(Model.Settings.Domain), *BlendToString(Model.Settings.BlendMode));
				AllEntries.Add(Entry);
			}
		}

		/** Finalize background shader compiles and render any queued/realtime thumbnails (see class note). */
		EActiveTimerReturnType PumpThumbnails(double, float DeltaTime)
		{
			if (GShaderCompilingManager)
			{
				GShaderCompilingManager->ProcessAsyncResults(/*bLimitExecutionTime*/ true, /*bBlockOnGlobalShaderCompletion*/ false);
			}
			check(ThumbnailPool.IsValid());
			ThumbnailPool->Tick(DeltaTime);
			return EActiveTimerReturnType::Continue;
		}

		TSharedRef<ITableRow> OnGenerateRow(FBaseEntryPtr Entry, const TSharedRef<STableViewBase>& OwnerTable)
		{
			check(Entry.IsValid());

			// On-demand: SListView only generates visible rows, so a base's thumbnail (and thus its
			// background shader compile) only starts once the user actually scrolls it into view.
			if (!Entry->Thumbnail.IsValid())
			{
				Entry->Thumbnail = MakeShared<FAssetThumbnail>(FAssetData(Entry->Material), GThumbnailSize, GThumbnailSize, ThumbnailPool);
				Entry->Thumbnail->SetRealTime(true); // keep refreshing past the first (still-compiling) frame
			}

			return SNew(STableRow<FBaseEntryPtr>, OwnerTable)
			[
				SNew(SHorizontalBox)
				+ SHorizontalBox::Slot().AutoWidth().VAlign(VAlign_Center).Padding(4.f)
				[
					SNew(SBox).WidthOverride(GThumbnailSize).HeightOverride(GThumbnailSize)
					[
						Entry->Thumbnail->MakeThumbnailWidget()
					]
				]
				+ SHorizontalBox::Slot().FillWidth(1.f).VAlign(VAlign_Center).Padding(8.f, 0.f)
				[
					SNew(SVerticalBox)
					+ SVerticalBox::Slot().AutoHeight()
					[
						SNew(STextBlock)
						.Text(FText::FromString(Entry->Name))
						.Font(FAppStyle::GetFontStyle("PropertyWindow.BoldFont"))
					]
					+ SVerticalBox::Slot().AutoHeight()
					[
						SNew(STextBlock)
						.Text(FText::FromString(Entry->DomainBlend))
						.ColorAndOpacity(FSlateColor::UseSubduedForeground())
					]
				]
			];
		}

		void OnSearchTextChanged(const FText& Text)
		{
			const FString Filter = Text.ToString();
			FilteredEntries.Reset();
			for (const FBaseEntryPtr& Entry : AllEntries)
			{
				if (Filter.IsEmpty() || Entry->Name.Contains(Filter))
				{
					FilteredEntries.Add(Entry);
				}
			}
			check(ListView.IsValid());
			ListView->RequestListRefresh();
		}

		void OnSelectionChanged(FBaseEntryPtr Entry, ESelectInfo::Type) { Selected = Entry; }

		void OnItemDoubleClicked(FBaseEntryPtr Entry)
		{
			check(Entry.IsValid());
			Selected = Entry;
			Confirm();
		}

		bool IsSelectEnabled() const { return Selected.IsValid(); }

		FReply OnSelectClicked()
		{
			check(Selected.IsValid()); // Button is disabled without a selection.
			Confirm();
			return FReply::Handled();
		}

		FReply OnCancelClicked()
		{
			ParentWindow.Pin()->RequestDestroyWindow();
			return FReply::Handled();
		}

		void Confirm()
		{
			check(Selected.IsValid());
			Chosen = Selected->Material;
			ParentWindow.Pin()->RequestDestroyWindow();
		}

		TWeakPtr<SWindow> ParentWindow;
		TSharedPtr<FAssetThumbnailPool> ThumbnailPool;
		TSharedPtr<SListView<FBaseEntryPtr>> ListView;
		TArray<FBaseEntryPtr> AllEntries;
		TArray<FBaseEntryPtr> FilteredEntries;
		FBaseEntryPtr Selected;
		UMaterial* Chosen = nullptr;
	};

	/** Pop the modal base-material picker; returns the chosen base or nullptr (cancelled / none registered). */
	UMaterial* PickShaderLabBase()
	{
		// No registered bases -> nothing to pick (the factory aborts creation).
		if (FShaderLabMaterialRegistry::Get().GetRegisteredModels().IsEmpty())
		{
			return nullptr;
		}

		TSharedRef<SWindow> Window = SNew(SWindow)
			.Title(LOCTEXT("PickBaseTitle", "Choose ShaderLab base material"))
			.ClientSize(FVector2D(480, 560))
			.SupportsMaximize(false)
			.SupportsMinimize(false);

		// Hold the picker alive past window destruction so we can read the choice after the modal loop.
		TSharedRef<SShaderLabBasePicker> Picker = SNew(SShaderLabBasePicker).ParentWindow(Window);
		Window->SetContent(Picker);

		GEditor->EditorAddModalWindow(Window);
		return Picker->GetChosen();
	}
}

UShaderLabMaterialInstanceFactory::UShaderLabMaterialInstanceFactory()
{
	SupportedClass = UShaderLabMaterialInstanceConstant::StaticClass();
	bCreateNew = true;
	bEditAfterNew = true;
}

uint32 UShaderLabMaterialInstanceFactory::GetMenuCategories() const
{
	return EAssetTypeCategories::Materials;
}

FText UShaderLabMaterialInstanceFactory::GetDisplayName() const
{
	return LOCTEXT("DisplayName", "ShaderLab Material Instance");
}

bool UShaderLabMaterialInstanceFactory::ConfigureProperties()
{
	// A commandlet/automation can set InitialParent directly and skip the (headless-unavailable) dialog.
	if (InitialParent)
	{
		return true;
	}
	if (FApp::IsUnattended() || IsRunningCommandlet() || !GEditor)
	{
		return false; // No parent and no UI to pick one.
	}

	InitialParent = PickShaderLabBase();
	return InitialParent != nullptr; // false = user cancelled / no bases -> abort creation.
}

UObject* UShaderLabMaterialInstanceFactory::FactoryCreateNew(
	UClass* InClass, UObject* InParent, FName InName, EObjectFlags Flags, UObject* Context, FFeedbackContext* Warn)
{
	UShaderLabMaterialInstanceConstant* Instance =
		NewObject<UShaderLabMaterialInstanceConstant>(InParent, InClass, InName, Flags);

	if (InitialParent)
	{
		Instance->SetParentEditorOnly(InitialParent);
		Instance->PostEditChange();
	}
	return Instance;
}

#undef LOCTEXT_NAMESPACE