From 898296125424820576c5fcc5e7bedeb1b8bbb5f4 Mon Sep 17 00:00:00 2001
From: Menotdan <32620310+Menotdan@users.noreply.github.com>
Date: Mon, 23 Oct 2023 22:08:50 -0400
Subject: [PATCH] Add custom material option to Material Preview.

---
 .../Viewport/Previews/MaterialPreview.cs      | 74 ++++++++++++++++---
 1 file changed, 63 insertions(+), 11 deletions(-)

diff --git a/Source/Editor/Viewport/Previews/MaterialPreview.cs b/Source/Editor/Viewport/Previews/MaterialPreview.cs
index 5520d6c4c..f05c8413b 100644
--- a/Source/Editor/Viewport/Previews/MaterialPreview.cs
+++ b/Source/Editor/Viewport/Previews/MaterialPreview.cs
@@ -7,6 +7,9 @@ using FlaxEngine.GUI;
 using FlaxEditor.Viewport.Widgets;
 using FlaxEditor.GUI.ContextMenu;
 using Object = FlaxEngine.Object;
+using FlaxEditor.CustomEditors.Editors;
+using FlaxEditor.GUI;
+using FlaxEditor.Scripting;
 
 namespace FlaxEditor.Viewport.Previews
 {
@@ -49,6 +52,8 @@ namespace FlaxEditor.Viewport.Previews
         private Image _guiMaterialControl;
         private readonly MaterialBase[] _postFxMaterialsCache = new MaterialBase[1];
         private ContextMenu _modelWidgetButtonMenu;
+        private AssetPicker _customModelPicker;
+        private Model _customModel;
 
         /// <summary>
         /// Gets or sets the material asset to preview. It can be <see cref="FlaxEngine.Material"/> or <see cref="FlaxEngine.MaterialInstance"/>.
@@ -74,15 +79,72 @@ namespace FlaxEditor.Viewport.Previews
             get => _selectedModelIndex;
             set
             {
+                if (value == -1) // Using Custom Model
+                {
+                    return;
+                }
+
                 if (value < 0 || value > Models.Length)
                     throw new ArgumentOutOfRangeException();
 
+                if (_customModelPicker != null)
+                {
+                    _customModelPicker.SelectedAsset = null;
+                }
                 _selectedModelIndex = value;
                 _previewModel.Model = FlaxEngine.Content.LoadAsyncInternal<Model>("Editor/Primitives/" + Models[value]);
                 _previewModel.Transform = Transforms[value];
             }
         }
 
+        // Used to automatically update which entry is checked.
+        // TODO: Maybe a better system with predicate bool checks could be used?
+        private void ResetModelContextMenu()
+        {
+            _modelWidgetButtonMenu.ItemsContainer.DisposeChildren();
+
+            // Fill out all models
+            for (int i = 0; i < Models.Length; i++)
+            {
+                var index = i;
+                var button = _modelWidgetButtonMenu.AddButton(Models[index]);
+                button.ButtonClicked += _ => SelectedModelIndex = index;
+                button.Checked = SelectedModelIndex == index && _customModel == null;
+                button.Tag = index;
+            }
+
+            _modelWidgetButtonMenu.AddSeparator();
+            _customModelPicker = new AssetPicker(new ScriptType(typeof(Model)), Float2.Zero);
+
+            // Label Button
+            var customModelPickerLabel = _modelWidgetButtonMenu.AddButton("Custom Model:");
+            customModelPickerLabel.CloseMenuOnClick = false;
+            customModelPickerLabel.Checked = _customModel != null;
+
+            // Container button
+            var customModelPickerButton = _modelWidgetButtonMenu.AddButton("");
+
+            customModelPickerButton.Height = _customModelPicker.Height + 4;
+            customModelPickerButton.CloseMenuOnClick = false;
+            _customModelPicker.Parent = customModelPickerButton;
+            _customModelPicker.SelectedAsset = _customModel;
+            _customModelPicker.SelectedItemChanged += () =>
+            {
+                _customModel = _customModelPicker.SelectedAsset as Model;
+                if (_customModelPicker.SelectedAsset == null)
+                {
+                    SelectedModelIndex = 0;
+                    ResetModelContextMenu();
+                    return;
+                }
+
+                _previewModel.Model = _customModel;
+                _previewModel.Transform = Transforms[0];
+                SelectedModelIndex = -1;
+                ResetModelContextMenu();
+            };
+        }
+
         /// <summary>
         /// Initializes a new instance of the <see cref="MaterialPreview"/> class.
         /// </summary>
@@ -107,17 +169,7 @@ namespace FlaxEditor.Viewport.Previews
                 {
                     if (!control.Visible)
                         return;
-                    _modelWidgetButtonMenu.ItemsContainer.DisposeChildren();
-
-                    // Fill out all models
-                    for (int i = 0; i < Models.Length; i++)
-                    {
-                        var index = i;
-                        var button = _modelWidgetButtonMenu.AddButton(Models[index]);
-                        button.ButtonClicked += _ => SelectedModelIndex = index;
-                        button.Checked = SelectedModelIndex == index;
-                        button.Tag = index;
-                    }
+                    ResetModelContextMenu();
                 };
                 new ViewportWidgetButton("Model", SpriteHandle.Invalid, _modelWidgetButtonMenu)
                 {