AppleHillsProduction/Assets/Editor/PrefabVariantGeneratorWindow.cs

using UnityEditor;
using UnityEngine;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Text;

namespace Editor
{
    public class PrefabVariantGeneratorWindow : EditorWindow
    {
        // Renderer configuration class to track sprite renderers and their assigned sprites
        [System.Serializable]
        private class RendererConfig
        {
            public string Path;                               // Hierarchy path to the renderer
            public string Name;                               // Display name for the renderer
            public SpriteRenderer Renderer;                   // Reference to the actual renderer
            public Sprite CurrentSprite;                      // Current sprite in the renderer
            public List<Sprite> AssignedSprites = new List<Sprite>(); // Sprites to use for variants
            public bool Enabled = true;                       // Whether to include in variant generation
            public bool Expanded = true;                      // UI expanded state
            public Vector2 ScrollPosition;                    // Scroll position for sprite list
        }

        // Main fields
        private GameObject sourcePrefab;
        private GameObject previousSourcePrefab;
        private List<RendererConfig> detectedRenderers = new List<RendererConfig>();
        private Vector2 mainScrollPosition;
        private string variantSaveFolder = "Assets/Prefabs/Variants";
        private string namingPattern = "{0}_{1}"; // Default: {0} = prefab name, {1} = first renderer sprite
        private bool userChangedSavePath = false;
        private int estimatedVariantCount = 0;
        private int maxSafeVariantCount = 100; // Warn above this number
        private GUIStyle boldFoldoutStyle;
        private bool showDefaultHelp = true;

        // Editor window setup
        [MenuItem("Tools/Sprite Variant Generator")]
        public static void ShowWindow()
        {
            var window = GetWindow<PrefabVariantGeneratorWindow>("Prefab Variant Generator");
            window.minSize = new Vector2(500, 600);
        }

        private void OnEnable()
        {
            // Initialize styles on enable to avoid null reference issues
            boldFoldoutStyle = new GUIStyle(EditorStyles.foldout)
            {
                fontStyle = FontStyle.Bold
            };
        }

        private void OnGUI()
        {
            // Initialize styles if needed
            if (boldFoldoutStyle == null)
            {
                boldFoldoutStyle = new GUIStyle(EditorStyles.foldout)
                {
                    fontStyle = FontStyle.Bold
                };
            }

            EditorGUILayout.LabelField("Prefab Variant Generator", EditorStyles.boldLabel);
            EditorGUILayout.HelpBox("Create multiple prefab variants with different sprites assigned to renderers.", MessageType.Info);

            mainScrollPosition = EditorGUILayout.BeginScrollView(mainScrollPosition);

            // Source Prefab Selection
            EditorGUILayout.Space();
            EditorGUILayout.LabelField("Step 1: Select Source Prefab", EditorStyles.boldLabel);

            // Store previous selection to detect changes
            GameObject newSourcePrefab = (GameObject)EditorGUILayout.ObjectField("Source Prefab", sourcePrefab, typeof(GameObject), false);

            // Check if prefab selection changed
            if (newSourcePrefab != previousSourcePrefab)
            {
                sourcePrefab = newSourcePrefab;
                previousSourcePrefab = newSourcePrefab;

                // Auto-set save folder to match source prefab's directory if a valid prefab is selected
                if (sourcePrefab != null && !userChangedSavePath)
                {
                    string prefabPath = AssetDatabase.GetAssetPath(sourcePrefab);
                    if (!string.IsNullOrEmpty(prefabPath))
                    {
                        variantSaveFolder = Path.GetDirectoryName(prefabPath).Replace("\\", "/");
                    }
                }

                // Find sprite renderers in the prefab
                FindRenderersInPrefab();

                // Clear default help once a prefab is selected
                if (sourcePrefab != null)
                {
                    showDefaultHelp = false;
                }
            }

            // Warn if not a prefab
            if (sourcePrefab != null && !PrefabUtility.IsPartOfPrefabAsset(sourcePrefab) && !PrefabUtility.IsPartOfPrefabInstance(sourcePrefab))
            {
                EditorGUILayout.HelpBox("Please select a prefab asset.", MessageType.Warning);
            }

            // Display default help if no prefab selected
            if (showDefaultHelp && sourcePrefab == null)
            {
                EditorGUILayout.HelpBox(
                    "This tool lets you create prefab variants with different sprites.\n\n" +
                    "1. Select a source prefab\n" +
                    "2. Assign sprites to each detected sprite renderer\n" +
                    "3. Generate all combinations as prefab variants", 
                    MessageType.Info
                );
            }

            // Only show the rest if a valid prefab is selected
            if (sourcePrefab != null)
            {
                // Renderer sections
                EditorGUILayout.Space();
                EditorGUILayout.LabelField("Step 2: Configure Sprite Renderers", EditorStyles.boldLabel);

                if (detectedRenderers.Count == 0)
                {
                    EditorGUILayout.HelpBox("No sprite renderers found in prefab. A new renderer will be created.", MessageType.Info);
                }
                else
                {
                    EditorGUILayout.HelpBox($"{detectedRenderers.Count} sprite renderer{(detectedRenderers.Count > 1 ? "s" : "")} found in prefab.", MessageType.Info);
                }

                // Display each renderer configuration
                for (int i = 0; i < detectedRenderers.Count; i++)
                {
                    DrawRendererSection(detectedRenderers[i], i);
                }

                // Update estimated variant count
                UpdateVariantCount();

                // Output settings
                EditorGUILayout.Space();
                EditorGUILayout.LabelField("Step 3: Output Settings", EditorStyles.boldLabel);

                // Save folder
                EditorGUILayout.BeginHorizontal();
                EditorGUILayout.PrefixLabel("Save Folder");
                EditorGUILayout.SelectableLabel(variantSaveFolder, EditorStyles.textField, GUILayout.Height(EditorGUIUtility.singleLineHeight));
                if (GUILayout.Button("Select...", GUILayout.Width(80)))
                {
                    string newFolder = PrefabEditorUtility.SelectFolder(variantSaveFolder, "Prefabs/Variants");
                    if (newFolder != variantSaveFolder)
                    {
                        variantSaveFolder = newFolder;
                        userChangedSavePath = true; // Mark that user manually changed the path
                    }
                }
                EditorGUILayout.EndHorizontal();

                // Add a reset button if user changed the path and a valid prefab is selected
                if (userChangedSavePath && sourcePrefab != null)
                {
                    string prefabPath = AssetDatabase.GetAssetPath(sourcePrefab);
                    if (!string.IsNullOrEmpty(prefabPath))
                    {
                        if (GUILayout.Button("Reset Path to Prefab Directory"))
                        {
                            variantSaveFolder = Path.GetDirectoryName(prefabPath).Replace("\\", "/");
                            userChangedSavePath = false;
                        }
                    }
                }

                // Naming pattern field
                EditorGUILayout.BeginHorizontal();
                EditorGUILayout.PrefixLabel("Naming Pattern");
                namingPattern = EditorGUILayout.TextField(namingPattern);
                EditorGUILayout.EndHorizontal();

                // Help text for naming pattern
                StringBuilder helpText = new StringBuilder("Naming placeholders:\n");
                helpText.AppendLine("{0} = Prefab name");
                for (int i = 0; i < detectedRenderers.Count; i++)
                {
                    helpText.AppendLine($"{{{i+1}}} = {detectedRenderers[i].Name} sprite name");
                }
                EditorGUILayout.HelpBox(helpText.ToString(), MessageType.Info);

                // Variant count display
                string variantCountText = $"Will generate {estimatedVariantCount} variant{(estimatedVariantCount != 1 ? "s" : "")}";
                if (estimatedVariantCount > maxSafeVariantCount)
                {
                    EditorGUILayout.HelpBox($"Warning: {variantCountText}. This might take some time.", MessageType.Warning);
                }
                else if (estimatedVariantCount > 0)
                {
                    EditorGUILayout.HelpBox(variantCountText, MessageType.Info);
                }
                else
                {
                    EditorGUILayout.HelpBox("Please assign at least one sprite to each enabled renderer to generate variants.", MessageType.Warning);
                }

                // Generate button
                EditorGUILayout.Space();
                GUI.enabled = estimatedVariantCount > 0;
                if (GUILayout.Button("Generate Prefab Variants", GUILayout.Height(30)))
                {
                    // Show warning for large numbers of variants
                    if (estimatedVariantCount > maxSafeVariantCount)
                    {
                        bool proceed = EditorUtility.DisplayDialog(
                            "Generate Many Variants?",
                            $"You are about to generate {estimatedVariantCount} prefab variants. This might take some time and use significant disk space. Continue?",
                            "Generate",
                            "Cancel"
                        );

                        if (!proceed) return;
                    }

                    GeneratePrefabVariants();
                }
                GUI.enabled = true;
            }

            EditorGUILayout.EndScrollView();
        }

        private void DrawRendererSection(RendererConfig config, int index)
        {
            EditorGUILayout.BeginVertical(EditorStyles.helpBox);

            EditorGUILayout.BeginHorizontal();
            // Expand/collapse button
            config.Expanded = EditorGUILayout.Foldout(config.Expanded, "", boldFoldoutStyle);

            // Enable/disable toggle
            bool newEnabled = EditorGUILayout.Toggle(config.Enabled, GUILayout.Width(20));
            if (newEnabled != config.Enabled)
            {
                config.Enabled = newEnabled;
                UpdateVariantCount();
            }

            // Renderer name/title
            EditorGUILayout.LabelField(config.Name, EditorStyles.boldLabel);

            // Current sprite preview if available
            if (config.CurrentSprite != null)
            {
                GUILayout.Box(
                    AssetPreview.GetAssetPreview(config.CurrentSprite), 
                    GUILayout.Width(40), 
                    GUILayout.Height(40)
                );
            }

            EditorGUILayout.EndHorizontal();

            // Only show contents if expanded
            if (config.Expanded)
            {
                // Path display
                if (!string.IsNullOrEmpty(config.Path))
                {
                    EditorGUILayout.LabelField($"Path: {config.Path}", EditorStyles.miniLabel);
                }

                EditorGUI.BeginDisabledGroup(!config.Enabled);

                // Sprite selection controls
                EditorGUILayout.Space();
                EditorGUILayout.BeginHorizontal();

                // Drag and drop area for sprites
                EditorGUILayout.BeginVertical(EditorStyles.helpBox, GUILayout.Height(60));
                EditorGUILayout.LabelField("Drag and drop sprites here", EditorStyles.centeredGreyMiniLabel);

                Rect dropArea = GUILayoutUtility.GetRect(0, 40, GUILayout.ExpandWidth(true));
                HandleDragAndDrop(dropArea, config);

                EditorGUILayout.EndVertical();

                if (GUILayout.Button("Add Selected", GUILayout.Width(100), GUILayout.Height(60)))
                {
                    AddSelectedSpritesToConfig(config);
                }

                EditorGUILayout.EndHorizontal();

                if (GUILayout.Button("Clear Sprites"))
                {
                    config.AssignedSprites.Clear();
                    UpdateVariantCount();
                }

                // Display selected sprites
                EditorGUILayout.Space();
                EditorGUILayout.LabelField($"Selected Sprites ({config.AssignedSprites.Count}):", EditorStyles.miniBoldLabel);

                // Sprite list
                config.ScrollPosition = EditorGUILayout.BeginScrollView(config.ScrollPosition, GUILayout.Height(120));
                for (int i = config.AssignedSprites.Count - 1; i >= 0; i--)
                {
                    EditorGUILayout.BeginHorizontal();
                    config.AssignedSprites[i] = (Sprite)EditorGUILayout.ObjectField(
                        config.AssignedSprites[i], 
                        typeof(Sprite), 
                        false,
                        GUILayout.ExpandWidth(true)
                    );

                    // Preview sprite
                    if (config.AssignedSprites[i] != null)
                    {
                        GUILayout.Box(
                            AssetPreview.GetAssetPreview(config.AssignedSprites[i]), 
                            GUILayout.Width(40), 
                            GUILayout.Height(40)
                        );
                    }

                    if (GUILayout.Button("Remove", GUILayout.Width(60)))
                    {
                        config.AssignedSprites.RemoveAt(i);
                        UpdateVariantCount();
                    }

                    EditorGUILayout.EndHorizontal();
                }
                EditorGUILayout.EndScrollView();

                EditorGUI.EndDisabledGroup();
            }

            EditorGUILayout.EndVertical();
        }

        private void HandleDragAndDrop(Rect dropArea, RendererConfig config)
        {
            Event evt = Event.current;
            switch (evt.type)
            {
                case EventType.DragUpdated:
                case EventType.DragPerform:
                    if (!dropArea.Contains(evt.mousePosition))
                        break;

                    DragAndDrop.visualMode = DragAndDropVisualMode.Copy;

                    if (evt.type == EventType.DragPerform)
                    {
                        DragAndDrop.AcceptDrag();
                        bool added = false;

                        foreach (var draggedObject in DragAndDrop.objectReferences)
                        {
                            if (draggedObject is Sprite sprite)
                            {
                                if (!config.AssignedSprites.Contains(sprite))
                                {
                                    config.AssignedSprites.Add(sprite);
                                    added = true;
                                }
                            }
                            else if (draggedObject is Texture2D texture)
                            {
                                // Try to get sprites from texture
                                string texturePath = AssetDatabase.GetAssetPath(texture);
                                var sprites = AssetDatabase.LoadAllAssetsAtPath(texturePath)
                                    .OfType<Sprite>()
                                    .ToArray();

                                foreach (var s in sprites)
                                {
                                    if (!config.AssignedSprites.Contains(s))
                                    {
                                        config.AssignedSprites.Add(s);
                                        added = true;
                                    }
                                }
                            }
                        }

                        if (added)
                        {
                            UpdateVariantCount();
                        }

                        evt.Use();
                    }
                    break;
            }
        }

        private void AddSelectedSpritesToConfig(RendererConfig config)
        {
            var selectedObjects = Selection.objects;
            bool added = false;

            foreach (var obj in selectedObjects)
            {
                if (obj is Sprite sprite)
                {
                    if (!config.AssignedSprites.Contains(sprite))
                    {
                        config.AssignedSprites.Add(sprite);
                        added = true;
                    }
                }
                else if (obj is Texture2D texture)
                {
                    // Try to get sprites from texture
                    string texturePath = AssetDatabase.GetAssetPath(texture);
                    var sprites = AssetDatabase.LoadAllAssetsAtPath(texturePath)
                        .OfType<Sprite>()
                        .ToArray();

                    foreach (var s in sprites)
                    {
                        if (!config.AssignedSprites.Contains(s))
                        {
                            config.AssignedSprites.Add(s);
                            added = true;
                        }
                    }
                }
            }

            if (added)
            {
                UpdateVariantCount();
            }
        }

        private void FindRenderersInPrefab()
        {
            detectedRenderers.Clear();

            if (sourcePrefab == null) return;

            // Get all renderers in prefab (including children)
            GameObject instance = null;

            try
            {
                instance = (GameObject)PrefabUtility.InstantiatePrefab(sourcePrefab);
                SpriteRenderer[] renderers = instance.GetComponentsInChildren<SpriteRenderer>(true);

                for (int i = 0; i < renderers.Length; i++)
                {
                    var renderer = renderers[i];
                    string path = GetRelativePath(instance.transform, renderer.transform);
                    string name = renderer.gameObject.name;

                    // For root object, use "Main"
                    if (string.IsNullOrEmpty(path))
                    {
                        name = "Main";
                    }
                    // For objects with the same name, add index
                    else if (renderers.Count(r => r.gameObject.name == renderer.gameObject.name) > 1)
                    {
                        name = $"{name} ({i+1})";
                    }

                    detectedRenderers.Add(new RendererConfig
                    {
                        Path = path,
                        Name = name,
                        Renderer = renderer,
                        CurrentSprite = renderer.sprite,
                        AssignedSprites = renderer.sprite != null ? 
                            new List<Sprite> { renderer.sprite } : 
                            new List<Sprite>()
                    });
                }
            }
            finally
            {
                if (instance != null)
                    DestroyImmediate(instance);
            }

            // If no renderers found, create a default entry
            if (detectedRenderers.Count == 0)
            {
                detectedRenderers.Add(new RendererConfig
                {
                    Path = "",
                    Name = "Main",
                    Renderer = null,
                    CurrentSprite = null,
                    AssignedSprites = new List<Sprite>()
                });
            }

            UpdateVariantCount();
        }

        private string GetRelativePath(Transform root, Transform target)
        {
            if (target == root) return "";

            string path = target.name;
            Transform parent = target.parent;

            while (parent != null && parent != root)
            {
                path = parent.name + "/" + path;
                parent = parent.parent;
            }

            return path;
        }

        private void UpdateVariantCount()
        {
            // Calculate estimated variants
            estimatedVariantCount = 0;

            // Get only enabled renderers with at least one sprite
            var enabledConfigs = detectedRenderers
                .Where(r => r.Enabled && r.AssignedSprites.Count > 0)
                .ToList();

            if (enabledConfigs.Count > 0)
            {
                // Start with count of first renderer's sprites
                estimatedVariantCount = enabledConfigs[0].AssignedSprites.Count;

                // Multiply by subsequent renderers' sprite counts
                for (int i = 1; i < enabledConfigs.Count; i++)
                {
                    estimatedVariantCount *= enabledConfigs[i].AssignedSprites.Count;
                }
            }
        }

        private List<List<Sprite>> GenerateAllCombinations()
        {
            var enabledConfigs = detectedRenderers
                .Where(r => r.Enabled && r.AssignedSprites.Count > 0)
                .ToList();

            if (enabledConfigs.Count == 0)
                return new List<List<Sprite>>();

            // Initialize with first renderer's sprites
            var combinations = enabledConfigs[0].AssignedSprites
                .Select(s => new List<Sprite> { s })
                .ToList();

            // Add each subsequent renderer's sprites
            for (int i = 1; i < enabledConfigs.Count; i++)
            {
                var newCombinations = new List<List<Sprite>>();
                foreach (var combo in combinations)
                {
                    foreach (var sprite in enabledConfigs[i].AssignedSprites)
                    {
                        var newCombo = new List<Sprite>(combo) { sprite };
                        newCombinations.Add(newCombo);
                    }
                }
                combinations = newCombinations;
            }

            return combinations;
        }

        private void GeneratePrefabVariants()
        {
            if (sourcePrefab == null)
            {
                EditorUtility.DisplayDialog("Error", "Please select a source prefab.", "OK");
                return;
            }

            // Get enabled renderer configurations
            var enabledConfigs = detectedRenderers
                .Where(r => r.Enabled && r.AssignedSprites.Count > 0)
                .ToList();

            if (enabledConfigs.Count == 0)
            {
                EditorUtility.DisplayDialog("Error", "Please assign at least one sprite to a renderer.", "OK");
                return;
            }

            // Ensure the save folder exists
            EnsureFolderExists(variantSaveFolder);

            // Generate all sprite combinations
            var combinations = GenerateAllCombinations();

            string sourcePrefabPath = AssetDatabase.GetAssetPath(sourcePrefab);
            string prefabName = Path.GetFileNameWithoutExtension(sourcePrefabPath);
            int successCount = 0;

            // Show progress bar
            EditorUtility.DisplayProgressBar("Generating Prefab Variants", "Preparing...", 0f);

            try
            {
                // For each combination, create a prefab variant
                for (int i = 0; i < combinations.Count; i++)
                {
                    // Update progress
                    if (i % 5 == 0 || i == combinations.Count - 1)
                    {
                        float progress = (float)i / combinations.Count;
                        if (EditorUtility.DisplayCancelableProgressBar(
                            "Generating Prefab Variants", 
                            $"Creating variant {i+1} of {combinations.Count}",
                            progress))
                        {
                            // User canceled
                            break;
                        }
                    }

                    var combination = combinations[i];

                    // Generate variant name
                    string variantName = prefabName;
                    string[] spriteNames = new string[combination.Count];

                    for (int j = 0; j < combination.Count; j++)
                    {
                        spriteNames[j] = combination[j].name;
                    }

                    // Format with the naming pattern
                    object[] formatArgs = new object[spriteNames.Length + 1];
                    formatArgs[0] = prefabName;
                    for (int j = 0; j < spriteNames.Length; j++)
                    {
                        formatArgs[j + 1] = spriteNames[j];
                    }

                    try
                    {
                        variantName = string.Format(namingPattern, formatArgs);
                    }
                    catch (System.FormatException)
                    {
                        // Fallback if format fails
                        variantName = $"{prefabName}_{string.Join("_", spriteNames)}";
                    }

                    variantName = PrefabEditorUtility.SanitizeFileName(variantName);
                    string variantPath = Path.Combine(variantSaveFolder, variantName + ".prefab").Replace("\\", "/");

                    // Create the prefab variant
                    GameObject prefabInstance = (GameObject)PrefabUtility.InstantiatePrefab(sourcePrefab);

                    try
                    {
                        // Apply sprites to renderers
                        for (int j = 0; j < enabledConfigs.Count; j++)
                        {
                            SpriteRenderer renderer = null;

                            // Find the corresponding renderer in the instance
                            if (string.IsNullOrEmpty(enabledConfigs[j].Path))
                            {
                                // Root object
                                renderer = prefabInstance.GetComponent<SpriteRenderer>();
                                if (renderer == null)
                                {
                                    renderer = prefabInstance.AddComponent<SpriteRenderer>();
                                }
                            }
                            else
                            {
                                // Child object
                                Transform child = prefabInstance.transform.Find(enabledConfigs[j].Path);
                                if (child != null)
                                {
                                    renderer = child.GetComponent<SpriteRenderer>();
                                }
                            }

                            // Apply sprite if renderer was found
                            if (renderer != null)
                            {
                                renderer.sprite = combination[j];
                            }
                        }

                        // Create the prefab variant
                        GameObject prefabVariant = PrefabUtility.SaveAsPrefabAsset(prefabInstance, variantPath);

                        if (prefabVariant != null)
                        {
                            successCount++;
                        }
                    }
                    catch (System.Exception e)
                    {
                        Debug.LogError($"Error creating prefab variant: {e.Message}");
                    }
                    finally
                    {
                        // Clean up the instance
                        DestroyImmediate(prefabInstance);
                    }
                }
            }
            finally
            {
                EditorUtility.ClearProgressBar();
            }

            AssetDatabase.Refresh();

            if (successCount > 0)
            {
                EditorUtility.DisplayDialog(
                    "Prefab Variants Created", 
                    $"Successfully created {successCount} prefab variants in {variantSaveFolder}.", 
                    "OK"
                );

                // Open the folder in Project view
                var folderObject = AssetDatabase.LoadAssetAtPath<UnityEngine.Object>(variantSaveFolder);
                if (folderObject != null)
                {
                    Selection.activeObject = folderObject;
                    EditorGUIUtility.PingObject(folderObject);
                }
            }
            else
            {
                EditorUtility.DisplayDialog(
                    "Prefab Variants", 
                    "No prefab variants were created. Please check the console for errors.", 
                    "OK"
                );
            }
        }

        private void EnsureFolderExists(string folderPath)
        {
            if (!AssetDatabase.IsValidFolder(folderPath))
            {
                string[] folderParts = folderPath.Split('/');
                string currentPath = folderParts[0];

                for (int i = 1; i < folderParts.Length; i++)
                {
                    string folderName = folderParts[i];
                    string newPath = Path.Combine(currentPath, folderName);

                    if (!AssetDatabase.IsValidFolder(newPath))
                    {
                        AssetDatabase.CreateFolder(currentPath, folderName);
                    }

                    currentPath = newPath;
                }

                AssetDatabase.Refresh();
            }
        }
    }
}