module-navigation-tool/Editor/Toolbar/Tools/ToolScenePicker.cs

using System;
using System.Collections.Generic;
using System.IO;
using System.Runtime.CompilerServices;
using System.Text;
using JetBrains.Annotations;
using UnityEditor;
using UnityEditor.SceneManagement;
using UnityEngine;
using UnityEngine.SceneManagement;

namespace Module.NavigationTool.Editor.Toolbar
{
    [UsedImplicitly]
    internal sealed class ToolScenePicker : AbstractToolbarDrawer
    {
        public override bool Visible => ToolbarScenePickerSettings.IsSceneEnabled;
        public override bool Enabled => !EditorApplication.isPlaying && !EditorApplication.isPlayingOrWillChangePlaymode;
        public override EToolbarPlacement Placement => EToolbarPlacement.Right;
        public override int Priority => (int)EToolbarPriority.Medium;

        private static readonly List<int> SELECTED_INDICES = new List<int>();
        private static readonly StringBuilder STRING_BUILDER = new StringBuilder();
        private static string SCENE_LABEL = string.Empty;
        private static string[] SCENE_NAMES = new string[0];
        private static string[] SCENE_SHORT_NAMES = new string[0];
        private static string[] SCENE_PATHS = new string[0];
        private static readonly GUIContent LABEL_SCENE_ADDITIVE = new GUIContent(string.Empty, "Additive scene loading (toogle to Single)");
        private static readonly GUIContent LABEL_SCENE_SINGLE = new GUIContent(string.Empty, "Single scene loading (toggle to Additive)");
        private static readonly GUIContent LABEL_SCENE_CREATE = new GUIContent(string.Empty, "Create a new scene");
        private const float BUTTON_WIDTH = 24.0f;
        
        private static bool IS_DIRTY = true;

        private static void Initialize()
        {
            if (!IS_DIRTY)
                return;
            
            SELECTED_INDICES.Clear();
            var listNames = new List<string>();
            var listShortNames = new List<string>();
            var listPaths = new List<string>();
            
            InitializeBuildSettingsScenes(listNames, listShortNames, listPaths);
            InitializeRemainingScenes(listNames, listShortNames, listPaths);
            
            for (var i = 0; i < SceneManager.sceneCount; i++)
            {
                Scene scene = SceneManager.GetSceneAt(i);
                int index = listPaths.IndexOf(scene.path);
                
                if (index != -1)
                    SELECTED_INDICES.Add(listPaths.IndexOf(scene.path));
            }
            
            SCENE_NAMES = listNames.ToArray();
            SCENE_SHORT_NAMES = listShortNames.ToArray();
            SCENE_PATHS = listPaths.ToArray();
            RefreshSceneLabel();
            IS_DIRTY = false;
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        private static void InitializeBuildSettingsScenes(List<string> listNames, List<string> listShortNames, List<string> listPaths)
        {
            EditorBuildSettingsScene[] scenes = EditorBuildSettings.scenes;

            for (var i = 0; i < scenes.Length; i++)
            {
                if (string.IsNullOrEmpty(scenes[i].path))
                    continue;

                listNames.Add(Path.GetFileNameWithoutExtension(scenes[i].path));
                listShortNames.Add(listNames[i]);
                listPaths.Add(scenes[i].path);
            }
        }
        
        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        private static void InitializeRemainingScenes(List<string> listNames, List<string> listShortNames, List<string> listPaths)
        {
            List<string> sortByAssetLabels = ToolbarScenePickerSettings.SceneAssetLabels;
            var assetLabelToScenes = new SceneSortElement[sortByAssetLabels.Count + 1];
            string[] guids = AssetDatabase.FindAssets("t:scene");
            
            for (var i = 0; i < guids.Length; i++)
            {
                string path = AssetDatabase.GUIDToAssetPath(guids[i]);
                
                if (listPaths.Contains(path) || !path.StartsWith("Assets"))
                    continue;

                var scene = AssetDatabase.LoadAssetAtPath<SceneAsset>(path);

                if (scene == null)
                    continue;
                
                string sceneName = path.Substring(7, path.Length - 13)
                                       .Replace('/', '\\');

                string[] assetLabels = AssetDatabase.GetLabels(scene);
                int index = -1;

                for (var j = 0; j < assetLabels.Length; j++)
                {
                    index = sortByAssetLabels.IndexOf(assetLabels[j]);

                    if (index != -1)
                        break;
                }

                if (index == -1)
                    index = sortByAssetLabels.Count;
                if (assetLabelToScenes[index] == null)
                    assetLabelToScenes[index] = new SceneSortElement();

                assetLabelToScenes[index].names.Add(sceneName);
                assetLabelToScenes[index].shortNames.Add(Path.GetFileName(sceneName));
                assetLabelToScenes[index].paths.Add(path);
            }

            for (var i = 0; i < assetLabelToScenes.Length; i++)
            {
                SceneSortElement e = assetLabelToScenes[i];

                if (e == null)
                    continue;

                listNames.Add(string.Empty);
                listShortNames.Add(string.Empty);
                listPaths.Add(string.Empty);
                
                listNames.AddRange(e.names);
                listShortNames.AddRange(e.shortNames);
                listPaths.AddRange(e.paths);
            }
        }

        public override void Update()
        {
            Initialize();
        }

        protected override void Draw(Rect rect)
        {
            Initialize();
            
            var rect0 = new Rect(rect.x, rect.y, rect.width - BUTTON_WIDTH * 2.0f, rect.height);
            var rect1 = new Rect(rect0.xMax, rect.y, BUTTON_WIDTH, rect.height);
            var rect2 = new Rect(rect1.xMax, rect.y, BUTTON_WIDTH, rect.height);
            DrawSceneDropDown(rect0);
            DrawSceneLoadToggle(rect1);
            DrawSceneAddButton(rect2);
        }

        private void DrawSceneDropDown(Rect rect)
        {
            if (GUI.Button(rect, SCENE_LABEL, styles.popup))
                ShowDropDown(rect);
        }

        private void DrawSceneLoadToggle(Rect rect)
        {
            bool isScenePickerSetAsAdditive = ToolbarScenePickerSettings.IsScenePickerSetAsAdditive;
            bool tempIsAdditive = EditorGUI.Toggle(rect, isScenePickerSetAsAdditive, styles.button);

            if (tempIsAdditive)
            {
                GUI.Label(rect, styles.iconSceneAdditive, styles.labelCenter);
                GUI.Label(rect, LABEL_SCENE_ADDITIVE, styles.labelCenter);
            }
            else
            {
                GUI.Label(rect, styles.iconSceneSingle, styles.labelCenter);
                GUI.Label(rect, LABEL_SCENE_SINGLE, styles.labelCenter);
            }

            if (tempIsAdditive != isScenePickerSetAsAdditive)
                ToolbarScenePickerSettings.IsScenePickerSetAsAdditive = tempIsAdditive;
        }

        private void DrawSceneAddButton(Rect rect)
        {
            if (GUI.Button(rect, styles.iconPlusSmall, styles.button))
            {
                try
                {
                    int selection = EditorUtility.DisplayDialogComplex("Create new scene", "How do you want to add the scene?", "Single", "Additive", "Cancel");

                    if (selection == 2)
                        return;
                
                    Scene scene = EditorSceneManager.NewScene(NewSceneSetup.EmptyScene, selection == 0 ? NewSceneMode.Single : NewSceneMode.Additive);
                
                    if (EditorUtility.DisplayDialog("Create new scene", "Save scene to disk?", "Yes", "No"))
                        EditorSceneManager.SaveScene(scene);
                }
                catch (Exception e)
                {
                    EditorUtility.DisplayDialog("Failed to add scene", e.Message, "Ok");
                }
            }
            
            GUI.Label(rect, LABEL_SCENE_CREATE, styles.labelCenter);
        }

        private static void ShowDropDown(Rect rect)
        {
            var menu = new GenericMenu();
                
            for (var i = 0; i < SCENE_NAMES.Length; i++)
            {
                int sceneIndex = i;
                menu.AddItem(new GUIContent(SCENE_NAMES[i]), SELECTED_INDICES.Contains(i), () => SelectScene(sceneIndex));
            }
            
            menu.DropDown(rect);
        }

        private static void SelectScene(int index)
        {
            string path = SCENE_PATHS[index];
            Scene scene = SceneManager.GetSceneByPath(path);
            bool isScenePickerSetAsAdditive = ToolbarScenePickerSettings.IsScenePickerSetAsAdditive;

            if (scene.isLoaded)
            {
                if (SceneManager.sceneCount == 1)
                    return;
                
                EditorSceneManager.CloseScene(scene, true);
                SELECTED_INDICES.Remove(index);
            }
            else if (isScenePickerSetAsAdditive)
            {
                EditorSceneManager.OpenScene(SCENE_PATHS[index], OpenSceneMode.Additive);
                SELECTED_INDICES.Add(index);
            }
            else
            {
                EditorSceneManager.OpenScene(SCENE_PATHS[index], OpenSceneMode.Single);
                SELECTED_INDICES.Clear();
                SELECTED_INDICES.Add(index);
            }
            
            RefreshSceneLabel();
        }

        private static void RefreshSceneLabel()
        {
            STRING_BUILDER.Clear();

            for (var i = 0; i < SELECTED_INDICES.Count; i++)
            {
                if (i > 0)
                    STRING_BUILDER.Append(", ");

                STRING_BUILDER.Append(SCENE_SHORT_NAMES[SELECTED_INDICES[i]]);
            }

            SCENE_LABEL = STRING_BUILDER.ToString();
        }
        
        public override float CalculateWidth()
        {
            return 100.0f + BUTTON_WIDTH * 2.0f;
        }

        public static void SetAsDirty()
        {
            IS_DIRTY = true;
        }

        private sealed class SceneSortElement
        {
            public readonly List<string> names = new List<string>();
            public readonly List<string> shortNames = new List<string>();
            public readonly List<string> paths = new List<string>();
        }
    }
}