// Copyright FlecsProj. All Rights Reserved.

#include "ShaderLabMaterialInstanceFactory.h"

#include "AssetRegistry/AssetData.h"
#include "AssetThumbnail.h"
#include "AssetTypeCategories.h"
#include "Editor.h"
#include "Framework/Application/SlateApplication.h"
#include "Materials/Material.h"
#include "Misc/App.h"
#include "ShaderLabMaterialInstanceConstant.h"
#include "ShaderLabMaterialRegistry.h"
#include "Widgets/Input/SButton.h"
#include "Widgets/Layout/SBox.h"
#include "Widgets/Layout/SScrollBox.h"
#include "Widgets/Layout/SUniformWrapPanel.h"
#include "Widgets/SBoxPanel.h"
#include "Widgets/SWindow.h"
#include "Widgets/Text/STextBlock.h"

#define LOCTEXT_NAMESPACE "ShaderLabMaterialInstanceFactory"

namespace
{
	/** Modal picker of the discovered ShaderLab base materials, with preview thumbnails. */
	UMaterial* PickShaderLabBase()
	{
		FShaderLabMaterialRegistry& Registry = FShaderLabMaterialRegistry::Get();

		// Collect the registered in-memory base materials (sorted by name for a stable layout).
		TArray<UMaterial*> Bases;
		TArray<FString> Names;
		Registry.GetRegisteredModels().GenerateKeyArray(Names);
		Names.Sort();
		for (const FString& Name : Names)
		{
			if (UMaterial* Base = Registry.FindMaterial(Name))
			{
				Bases.Add(Base);
			}
		}
		if (Bases.IsEmpty())
		{
			return nullptr;
		}

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

		TSharedRef<FAssetThumbnailPool> ThumbnailPool = MakeShared<FAssetThumbnailPool>(64);
		TSharedRef<SUniformWrapPanel> Grid = SNew(SUniformWrapPanel).SlotPadding(FMargin(6.f));

		for (UMaterial* Base : Bases)
		{
			// Hold the thumbnail alive for the window's lifetime by capturing it in the button lambda.
			TSharedRef<FAssetThumbnail> Thumbnail = MakeShared<FAssetThumbnail>(FAssetData(Base), 96, 96, ThumbnailPool);
			Grid->AddSlot()
			[
				SNew(SButton)
				.HAlign(HAlign_Center)
				.VAlign(VAlign_Center)
				.ToolTipText(FText::FromString(Base->GetName()))
				.OnClicked_Lambda([&Chosen, Base, Window, Thumbnail]()
				{
					Chosen = Base;
					Window->RequestDestroyWindow();
					return FReply::Handled();
				})
				.Content()
				[
					SNew(SVerticalBox)
					+ SVerticalBox::Slot().AutoHeight().HAlign(HAlign_Center)
					[
						SNew(SBox).WidthOverride(96.f).HeightOverride(96.f)
						[
							Thumbnail->MakeThumbnailWidget()
						]
					]
					+ SVerticalBox::Slot().AutoHeight().HAlign(HAlign_Center).Padding(0, 4, 0, 0)
					[
						SNew(STextBlock).Text(FText::FromString(Base->GetName()))
					]
				]
			];
		}

		Window->SetContent(
			SNew(SScrollBox)
			+ SScrollBox::Slot().Padding(8.f)
			[
				Grid
			]);

		GEditor->EditorAddModalWindow(Window);
		return Chosen;
	}
}

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