// Icon Studio - Mathew Simon

#include "IconStudioSubsystem.h"

#include "IconStudioPreset.h"
#include "IconStudioRenderer.h"
#include "IconStudioSession.h"
#include "IconStudioSettings.h"

#include "AssetRegistry/AssetRegistryModule.h"
#include "AssetRegistry/IAssetRegistry.h"
#include "ContentBrowserModule.h"
#include "Engine/StaticMesh.h"
#include "Engine/Texture2D.h"
#include "IContentBrowserSingleton.h"
#include "ImageCore.h"
#include "ImageUtils.h"
#include "Materials/MaterialInterface.h"
#include "Misc/PackageName.h"
#include "Misc/Paths.h"
#include "Misc/ScopedSlowTask.h"
#include "ObjectTools.h"
#include "UObject/SavePackage.h"
#include "UObject/UnrealType.h"

DEFINE_LOG_CATEGORY_STATIC(LogIconStudio, Log, All);

#define LOCTEXT_NAMESPACE "IconStudio"

namespace IconStudioLocal
{
	/** Lit une propriete objet par son nom, en verifiant le type attendu. */
	UObject* GetObjectProperty(const UObject* Owner, FName PropertyName, const UClass* ExpectedClass)
	{
		if (!Owner || PropertyName.IsNone())
		{
			return nullptr;
		}

		FObjectPropertyBase* Property = FindFProperty<FObjectPropertyBase>(Owner->GetClass(), PropertyName);
		if (!Property)
		{
			return nullptr;
		}

		UObject* Value = Property->GetObjectPropertyValue_InContainer(Owner);
		return (Value && Value->IsA(ExpectedClass)) ? Value : nullptr;
	}

	/** Lit un nom lisible, quel que soit le type textuel de la propriete. */
	FString GetTextualProperty(const UObject* Owner, FName PropertyName)
	{
		if (!Owner || PropertyName.IsNone())
		{
			return FString();
		}

		const UClass* Class = Owner->GetClass();
		if (const FNameProperty* NameProperty = FindFProperty<FNameProperty>(Class, PropertyName))
		{
			const FName Value = NameProperty->GetPropertyValue_InContainer(Owner);
			return Value.IsNone() ? FString() : Value.ToString();
		}
		if (const FStrProperty* StringProperty = FindFProperty<FStrProperty>(Class, PropertyName))
		{
			return StringProperty->GetPropertyValue_InContainer(Owner);
		}
		if (const FTextProperty* TextProperty = FindFProperty<FTextProperty>(Class, PropertyName))
		{
			return TextProperty->GetPropertyValue_InContainer(Owner).ToString();
		}

		return FString();
	}
}

void UIconStudioSubsystem::Initialize(FSubsystemCollectionBase& Collection)
{
	Super::Initialize(Collection);

	Renderer = MakeUnique<FIconStudioRenderer>();

	// Regler une lumiere du preset doit rafraichir l'image tout de suite, or le
	// preset est edite dans son propre Details View, que l'outil ne possede pas.
	// Ecouter le signal global est le seul point d'accroche qui couvre les deux
	// panneaux -- d'ou le filtrage strict juste apres, cette delegation etant
	// emise pour toute propriete modifiee dans l'editeur.
	PropertyChangedHandle = FCoreUObjectDelegates::OnObjectPropertyChanged.AddUObject(this, &UIconStudioSubsystem::HandleObjectPropertyChanged);
}

void UIconStudioSubsystem::Deinitialize()
{
	FCoreUObjectDelegates::OnObjectPropertyChanged.Remove(PropertyChangedHandle);
	PropertyChangedHandle.Reset();

	Renderer.Reset();
	Session = nullptr;

	Super::Deinitialize();
}

UIconStudioSession* UIconStudioSubsystem::GetSession()
{
	if (!Session)
	{
		Session = NewObject<UIconStudioSession>(this, NAME_None, RF_Transient);
		Session->OnSessionChanged.AddUObject(this, &UIconStudioSubsystem::HandleSessionChanged);
	}
	return Session;
}

void UIconStudioSubsystem::HandleSessionChanged()
{
	RefreshPreview();
}

void UIconStudioSubsystem::HandleObjectPropertyChanged(UObject* Object, FPropertyChangedEvent& Event)
{
	if (!Session || !Object)
	{
		return;
	}

	// Seul le preset en cours nous interesse : la session, elle, previent deja
	// par son propre delegue.
	if (Object == Session->Preset)
	{
		RefreshPreview();
	}
}

void UIconStudioSubsystem::Status(const FString& Message, bool bIsError)
{
	if (bIsError)
	{
		UE_LOG(LogIconStudio, Warning, TEXT("%s"), *Message);
	}
	else
	{
		UE_LOG(LogIconStudio, Log, TEXT("%s"), *Message);
	}

	OnStatus.Broadcast(Message);
}

UStaticMesh* UIconStudioSubsystem::ResolveMesh(UObject* Subject, UMaterialInterface*& OutMaterialOverride) const
{
	OutMaterialOverride = nullptr;

	if (!Subject)
	{
		return nullptr;
	}

	// Cas direct : on previsualise un mesh pour lui-meme.
	if (UStaticMesh* AsMesh = Cast<UStaticMesh>(Subject))
	{
		return AsMesh;
	}

	const UIconStudioSettings* Settings = GetDefault<UIconStudioSettings>();

	for (const FName& PropertyName : Settings->MeshPropertyNames)
	{
		if (UObject* Found = IconStudioLocal::GetObjectProperty(Subject, PropertyName, UStaticMesh::StaticClass()))
		{
			OutMaterialOverride = Cast<UMaterialInterface>(
				IconStudioLocal::GetObjectProperty(Subject, Settings->MaterialOverridePropertyName, UMaterialInterface::StaticClass()));
			return Cast<UStaticMesh>(Found);
		}
	}

	return nullptr;
}

FString UIconStudioSubsystem::BuildAssetName(UObject* Subject, const UIconStudioPreset* Preset) const
{
	const UIconStudioSettings* Settings = GetDefault<UIconStudioSettings>();

	FString BaseName = IconStudioLocal::GetTextualProperty(Subject, Settings->DisplayNamePropertyName);
	if (BaseName.IsEmpty())
	{
		// Retomber sur le nom de l'asset, en retirant son prefixe de convention :
		// DA_Branch donnerait sinon T_Icon_DA_Branch.
		BaseName = Subject->GetName();
		int32 UnderscoreIndex = INDEX_NONE;
		if (BaseName.FindChar(TEXT('_'), UnderscoreIndex) && UnderscoreIndex <= 3)
		{
			BaseName.MidInline(UnderscoreIndex + 1);
		}
	}

	return ObjectTools::SanitizeObjectName(Preset->AssetPrefix + BaseName);
}

UTexture2D* UIconStudioSubsystem::CreateOrUpdateTexture(const FIconStudioImage& Image, const FString& PackagePath, const FString& AssetName, TArray<UPackage*>& OutTouchedPackages)
{
	const FString PackageName = PackagePath / AssetName;
	if (!FPackageName::IsValidLongPackageName(PackageName))
	{
		Status(FString::Printf(TEXT("Invalid output path: %s"), *PackageName), /*bIsError=*/true);
		return nullptr;
	}

	// Regenerer une icone doit MODIFIER la texture existante, jamais en creer une
	// seconde : tout ce qui la reference -- les DataAssets, les widgets -- doit
	// continuer de pointer au meme endroit.
	const FString ObjectPath = PackageName + TEXT(".") + AssetName;
	UObject* Existing = LoadObject<UObject>(nullptr, *ObjectPath, nullptr, LOAD_NoWarn | LOAD_Quiet);
	if (Existing && !Existing->IsA<UTexture2D>())
	{
		Status(FString::Printf(TEXT("%s already exists and is not a texture."), *ObjectPath), /*bIsError=*/true);
		return nullptr;
	}

	UPackage* Package = Existing ? Existing->GetOutermost() : CreatePackage(*PackageName);
	if (!Package)
	{
		return nullptr;
	}

	UTexture2D* Texture = Cast<UTexture2D>(Existing);
	const bool bIsNew = (Texture == nullptr);
	if (bIsNew)
	{
		Texture = NewObject<UTexture2D>(Package, *AssetName, RF_Public | RF_Standalone);
	}

	Texture->Modify();
	// FColor est stocke en BGRA en memoire : la source de texture prend donc les
	// pixels tels quels, sans conversion.
	Texture->Source.Init(Image.Size, Image.Size, 1, 1, TSF_BGRA8, reinterpret_cast<const uint8*>(Image.Pixels.GetData()));

	// Reglages d'une texture d'interface : compression qui preserve l'alpha,
	// aucun mip -- une icone est affichee a taille fixe, ses mips ne serviraient
	// qu'a la rendre floue -- et pas de streaming, pour qu'elle soit la des la
	// premiere image du HUD.
	Texture->SRGB = true;
	Texture->CompressionSettings = TC_EditorIcon;
	Texture->MipGenSettings = TMGS_NoMipmaps;
	Texture->LODGroup = TEXTUREGROUP_UI;
	Texture->NeverStream = true;
	Texture->Filter = TF_Bilinear;
	Texture->PostEditChange();

	if (bIsNew)
	{
		FAssetRegistryModule::AssetCreated(Texture);
	}

	Package->MarkPackageDirty();
	OutTouchedPackages.AddUnique(Package);

	return Texture;
}

bool UIconStudioSubsystem::AssignIcon(UObject* Subject, UTexture2D* Texture, TArray<UPackage*>& OutTouchedPackages) const
{
	if (!Subject || !Texture)
	{
		return false;
	}

	const UIconStudioSettings* Settings = GetDefault<UIconStudioSettings>();
	if (Settings->IconPropertyName.IsNone())
	{
		return false;
	}

	FObjectPropertyBase* Property = FindFProperty<FObjectPropertyBase>(Subject->GetClass(), Settings->IconPropertyName);
	if (!Property || !Property->PropertyClass || !Texture->IsA(Property->PropertyClass))
	{
		return false;
	}

	Subject->Modify();
	Property->SetObjectPropertyValue_InContainer(Subject, Texture);

	// Sans PostEditChangeProperty, un DataAsset qui valide ses donnees ou met en
	// cache quelque chose ne verrait jamais passer la modification.
	FPropertyChangedEvent ChangedEvent(Property);
	Subject->PostEditChangeProperty(ChangedEvent);
	Subject->MarkPackageDirty();

	OutTouchedPackages.AddUnique(Subject->GetOutermost());
	return true;
}

bool UIconStudioSubsystem::WritePng(const FIconStudioImage& Image, const FString& AssetName) const
{
	const FString Directory = FPaths::ProjectSavedDir() / TEXT("IconStudio");
	const FString FullPath = Directory / (AssetName + TEXT(".png"));

	FImage Output;
	Output.Init(Image.Size, Image.Size, ERawImageFormat::BGRA8, EGammaSpace::sRGB);
	FMemory::Memcpy(Output.RawData.GetData(), Image.Pixels.GetData(), Image.Pixels.Num() * sizeof(FColor));

	return FImageUtils::SaveImageByExtension(*FullPath, Output);
}

UTexture2D* UIconStudioSubsystem::RefreshPreview()
{
	UIconStudioSession* const Current = GetSession();

	// Un seul point de sortie qui previent toujours, meme en cas d'echec : c'est
	// ce signal qui sert aussi a la fenetre pour raccrocher son panneau de
	// preset. S'en dispenser quand il n'y a rien a rendre laisserait l'interface
	// afficher les reglages du preset precedent.
	UTexture2D* Preview = nullptr;

	if (Current->Preset && Current->Subject && Renderer.IsValid())
	{
		UMaterialInterface* MaterialOverride = nullptr;
		if (UStaticMesh* const Mesh = ResolveMesh(Current->Subject, MaterialOverride))
		{
			FIconStudioImage Image;
			if (Renderer->Render(Mesh, MaterialOverride, Current->Preset, Current->Pass, Image))
			{
				Preview = Renderer->UpdatePreviewTexture(Image);
			}
			else
			{
				Status(TEXT("Render failed."), /*bIsError=*/true);
			}
		}
		else
		{
			Status(FString::Printf(TEXT("%s exposes no static mesh to render."), *Current->Subject->GetName()), /*bIsError=*/true);
		}
	}

	OnPreviewUpdated.Broadcast(Preview);
	return Preview;
}

bool UIconStudioSubsystem::GenerateInternal(UObject* Subject, TArray<UPackage*>& OutTouchedPackages)
{
	UIconStudioSession* const Current = GetSession();

	const UIconStudioPreset* const Preset = EnsurePreset();
	if (!Preset)
	{
		Status(TEXT("No preset available."), /*bIsError=*/true);
		return false;
	}
	if (!Subject)
	{
		return false;
	}

	const UIconStudioSettings* Settings = GetDefault<UIconStudioSettings>();

	if (Current->bSkipIfIconAlreadySet)
	{
		if (IconStudioLocal::GetObjectProperty(Subject, Settings->IconPropertyName, UTexture2D::StaticClass()))
		{
			Status(FString::Printf(TEXT("%s already has an icon, skipped."), *Subject->GetName()));
			return false;
		}
	}

	UMaterialInterface* MaterialOverride = nullptr;
	UStaticMesh* const Mesh = ResolveMesh(Subject, MaterialOverride);
	if (!Mesh)
	{
		Status(FString::Printf(TEXT("%s exposes no static mesh, skipped."), *Subject->GetName()), /*bIsError=*/true);
		return false;
	}

	// La passe choisie dans la session s'applique aussi au lot : exporter les
	// base colors de toute une famille d'objets est un besoin reel. C'est bien
	// ce qu'on voit dans la previsualisation qui sera ecrit sur le disque.
	FIconStudioImage Image;
	if (!Renderer->Render(Mesh, MaterialOverride, Preset, Current->Pass, Image))
	{
		Status(FString::Printf(TEXT("Render failed for %s."), *Subject->GetName()), /*bIsError=*/true);
		return false;
	}

	const FString AssetName = BuildAssetName(Subject, Preset);

	if (Preset->bAlsoWritePngToDisk && !WritePng(Image, AssetName))
	{
		Status(FString::Printf(TEXT("Could not write the PNG for %s."), *AssetName), /*bIsError=*/true);
	}

	UTexture2D* const Texture = CreateOrUpdateTexture(Image, Preset->OutputDirectory.Path, AssetName, OutTouchedPackages);
	if (!Texture)
	{
		return false;
	}

	if (Current->bAssignIconToDataAsset && Subject != Mesh)
	{
		if (AssignIcon(Subject, Texture, OutTouchedPackages))
		{
			Status(FString::Printf(TEXT("%s -> %s (assigned)"), *Subject->GetName(), *AssetName));
			return true;
		}

		Status(FString::Printf(TEXT("%s -> %s (created, but no '%s' property to assign it to)"),
			*Subject->GetName(), *AssetName, *Settings->IconPropertyName.ToString()));
		return true;
	}

	Status(FString::Printf(TEXT("%s -> %s"), *Subject->GetName(), *AssetName));
	return true;
}

bool UIconStudioSubsystem::SavePackageToDisk(UPackage* Package)
{
	if (!Package)
	{
		return false;
	}

	const FString Filename = FPackageName::LongPackageNameToFilename(Package->GetName(), FPackageName::GetAssetPackageExtension());

	FSavePackageArgs SaveArgs;
	SaveArgs.TopLevelFlags = RF_Public | RF_Standalone;
	SaveArgs.SaveFlags = SAVE_NoError;
	SaveArgs.Error = GError;

	return UPackage::SavePackage(Package, nullptr, *Filename, SaveArgs);
}

void UIconStudioSubsystem::SavePackages(const TArray<UPackage*>& Packages) const
{
	if (!GetDefault<UIconStudioSettings>()->bSaveModifiedPackages)
	{
		return;
	}

	for (UPackage* Package : Packages)
	{
		SavePackageToDisk(Package);
	}
}

UIconStudioPreset* UIconStudioSubsystem::EnsurePreset()
{
	UIconStudioSession* const Current = GetSession();
	if (Current->Preset)
	{
		return Current->Preset;
	}

	const IAssetRegistry& AssetRegistry = FModuleManager::LoadModuleChecked<FAssetRegistryModule>(TEXT("AssetRegistry")).Get();

	// Tant que le registre balaie encore le projet, une recherche infructueuse
	// ne prouve rien : creer un preset maintenant en ferait un second a cote de
	// celui qu'on n'a pas encore vu.
	if (AssetRegistry.IsLoadingAssets())
	{
		Status(TEXT("Asset registry still scanning, try again in a second."), /*bIsError=*/true);
		return nullptr;
	}

	FARFilter Filter;
	Filter.ClassPaths.Add(UIconStudioPreset::StaticClass()->GetClassPathName());
	Filter.bRecursiveClasses = true;

	TArray<FAssetData> Existing;
	AssetRegistry.GetAssets(Filter, Existing);

	if (!Existing.IsEmpty())
	{
		Current->Preset = Cast<UIconStudioPreset>(Existing[0].GetAsset());
		if (Current->Preset)
		{
			Status(FString::Printf(TEXT("Using preset %s."), *Current->Preset->GetName()));
			return Current->Preset;
		}
	}

	// Aucun preset dans le projet : on en pose un, avec les valeurs par defaut
	// de la classe. C'est un asset normal, que l'auteur retouche, renomme ou
	// duplique ensuite comme n'importe quel autre.
	static const FString DefaultPackageName(TEXT("/Game/EditorTools/DA_IconPreset_Default"));
	static const FString DefaultAssetName(TEXT("DA_IconPreset_Default"));

	UPackage* const Package = CreatePackage(*DefaultPackageName);
	if (!Package)
	{
		Status(TEXT("Could not create the default preset package."), /*bIsError=*/true);
		return nullptr;
	}

	UIconStudioPreset* const Preset = NewObject<UIconStudioPreset>(Package, *DefaultAssetName, RF_Public | RF_Standalone);
	FAssetRegistryModule::AssetCreated(Preset);
	Package->MarkPackageDirty();
	SavePackageToDisk(Package);

	Current->Preset = Preset;
	Status(FString::Printf(TEXT("Created a default preset: %s"), *DefaultPackageName));

	return Preset;
}

bool UIconStudioSubsystem::GenerateForSubject(UObject* Subject)
{
	TArray<UPackage*> Touched;
	const bool bSuccess = GenerateInternal(Subject, Touched);
	SavePackages(Touched);
	return bSuccess;
}

bool UIconStudioSubsystem::GenerateCurrent()
{
	return GenerateForSubject(GetSession()->Subject);
}

int32 UIconStudioSubsystem::GenerateBatch()
{
	UIconStudioSession* const Current = GetSession();
	if (Current->BatchTargets.IsEmpty())
	{
		Status(TEXT("Batch list is empty."), /*bIsError=*/true);
		return 0;
	}

	TArray<UPackage*> Touched;
	int32 SuccessCount = 0;

	// Un lot de deux cents objets fige l'editeur plusieurs secondes : sans barre
	// de progression, l'utilisateur croit a un plantage et tue le processus au
	// milieu de l'ecriture des packages.
	FScopedSlowTask SlowTask(static_cast<float>(Current->BatchTargets.Num()), LOCTEXT("GeneratingIcons", "Generating icons..."));
	SlowTask.MakeDialog(/*bShowCancelButton=*/true);

	for (const TObjectPtr<UObject>& Target : Current->BatchTargets)
	{
		if (SlowTask.ShouldCancel())
		{
			Status(TEXT("Batch cancelled."));
			break;
		}

		SlowTask.EnterProgressFrame(1.f, FText::FromString(Target ? Target->GetName() : TEXT("...")));

		if (GenerateInternal(Target, Touched))
		{
			++SuccessCount;
		}
	}

	SavePackages(Touched);
	Status(FString::Printf(TEXT("Batch finished: %d / %d icons generated."), SuccessCount, Current->BatchTargets.Num()));

	return SuccessCount;
}

int32 UIconStudioSubsystem::AddContentBrowserSelectionToBatch()
{
	FContentBrowserModule& ContentBrowser = FModuleManager::LoadModuleChecked<FContentBrowserModule>(TEXT("ContentBrowser"));

	TArray<FAssetData> Selection;
	ContentBrowser.Get().GetSelectedAssets(Selection);

	UIconStudioSession* const Current = GetSession();
	int32 AddedCount = 0;

	for (const FAssetData& Asset : Selection)
	{
		// GetAsset charge l'objet : c'est voulu, on a besoin de lire ses
		// proprietes, et il faudrait le charger de toute facon pour le rendre.
		UObject* Loaded = Asset.GetAsset();
		if (!Loaded)
		{
			continue;
		}

		UMaterialInterface* Unused = nullptr;
		if (!ResolveMesh(Loaded, Unused))
		{
			continue;
		}

		if (!Current->BatchTargets.Contains(Loaded))
		{
			Current->BatchTargets.Add(Loaded);
			++AddedCount;
		}
	}

	Status(FString::Printf(TEXT("%d asset(s) added to the batch (%d total)."), AddedCount, Current->BatchTargets.Num()));
	return AddedCount;
}

bool UIconStudioSubsystem::SetSubjectFromContentBrowser()
{
	FContentBrowserModule& ContentBrowser = FModuleManager::LoadModuleChecked<FContentBrowserModule>(TEXT("ContentBrowser"));

	TArray<FAssetData> Selection;
	ContentBrowser.Get().GetSelectedAssets(Selection);

	for (const FAssetData& Asset : Selection)
	{
		UObject* Loaded = Asset.GetAsset();
		UMaterialInterface* Unused = nullptr;
		if (Loaded && ResolveMesh(Loaded, Unused))
		{
			GetSession()->Subject = Loaded;
			RefreshPreview();
			return true;
		}
	}

	Status(TEXT("No renderable asset selected in the Content Browser."), /*bIsError=*/true);
	return false;
}

void UIconStudioSubsystem::ClearBatch()
{
	GetSession()->BatchTargets.Reset();
	Status(TEXT("Batch list cleared."));
}

#undef LOCTEXT_NAMESPACE
