MassPrefabRestorer

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

public class MassPrefabRestorer : EditorWindow
{
    // A support class to "snapshot" the state of the original object
    private class ObjectSnapshot
    {
        public Vector3 Position;
        public Quaternion Rotation;
        public Vector3 Scale;
        public Transform Parent;
        public int SiblingIndex;
        public List<Transform> OriginalChildren;
        public Material[] Materials; // Array to support multi-material renderers
    }

    private DefaultAsset sourceFolder;
    private string prefixToRemove = "";
    private string suffixDelimiter = " (";
    private bool searchInSubfolders = true;

    [MenuItem("Tools/Mass Prefab Restorer")]
    public static void ShowWindow()
    {
        var window = GetWindow<MassPrefabRestorer>(utility: true, title: "Mass Prefab Restorer");
        window.minSize = new Vector2(450, 260);
        window.maxSize = new Vector2(450, 260);
    }

    void OnGUI()
    {
        EditorGUILayout.LabelField("Mass Prefab Restorer", EditorStyles.boldLabel);
        EditorGUILayout.HelpBox("Intelligently replaces objects, preserving transforms, unique children, and material overrides.", MessageType.Info);

        EditorGUILayout.Space();

        EditorGUILayout.LabelField("1. Prefab Folder Settings", EditorStyles.boldLabel);
        sourceFolder = (DefaultAsset)EditorGUILayout.ObjectField("Prefab Source Folder", sourceFolder, typeof(DefaultAsset), false);
        searchInSubfolders = EditorGUILayout.Toggle("Search in subfolders", searchInSubfolders);

        EditorGUILayout.Space();

        EditorGUILayout.LabelField("2. Name Cleanup Settings", EditorStyles.boldLabel);
        prefixToRemove = EditorGUILayout.TextField("Prefix to remove", prefixToRemove);
        suffixDelimiter = EditorGUILayout.TextField("Suffix Delimiter", suffixDelimiter);

        EditorGUILayout.Space();

        if (GUILayout.Button("Start Object Replacement", GUILayout.Height(40)))
        {
            ReplaceObjects();
        }
    }

    void ReplaceObjects()
    {
        if (sourceFolder == null) { EditorUtility.DisplayDialog("Error", "Please select a source folder.", "OK"); return; }
        string folderPath = AssetDatabase.GetAssetPath(sourceFolder);
        if (!AssetDatabase.IsValidFolder(folderPath)) { EditorUtility.DisplayDialog("Error", "The selected asset is not a valid folder.", "OK"); return; }
        SearchOption searchOption = searchInSubfolders ? SearchOption.AllDirectories : SearchOption.TopDirectoryOnly;
        string[] prefabPaths = Directory.GetFiles(folderPath, "*.prefab", searchOption);
        if (prefabPaths.Length == 0) { EditorUtility.DisplayDialog("No Prefabs Found", "No .prefab files were found in the selected folder.", "OK"); return; }
        var prefabDict = new Dictionary<string, GameObject>();
        foreach (var path in prefabPaths)
        {
            var prefabName = Path.GetFileNameWithoutExtension(path);
            var prefab = AssetDatabase.LoadAssetAtPath<GameObject>(path);
            if (prefab != null && !prefabDict.ContainsKey(prefabName)) { prefabDict.Add(prefabName, prefab); }
        }

        Scene targetScene;
        int replacedCount = 0;
        var notFoundLog = new List<string>();

        Undo.SetCurrentGroupName("Mass Prefab Restore");
        var group = Undo.GetCurrentGroup();

        var currentPrefabStage = PrefabStageUtility.GetCurrentPrefabStage();
        if (currentPrefabStage != null && currentPrefabStage.scene.IsValid())
        {
            targetScene = currentPrefabStage.scene;
            GameObject prefabRoot = currentPrefabStage.prefabContentsRoot;
            Debug.Log("In Prefab Stage: Processing children of the root object only.");
            List<Transform> childrenOfRoot = prefabRoot.transform.Cast<Transform>().ToList();
            foreach (var child in childrenOfRoot)
            {
                ProcessObjectRecursively(child.gameObject, prefabDict, ref replacedCount, notFoundLog, targetScene);
            }
        }
        else
        {
            targetScene = SceneManager.GetActiveScene();
            if (!targetScene.IsValid()) { EditorUtility.DisplayDialog("Error", "Could not get a valid scene to operate on.", "OK"); return; }
            Debug.Log("In a regular scene: Processing all root objects.");
            var rootObjects = new List<GameObject>();
            targetScene.GetRootGameObjects(rootObjects);
            if (rootObjects.Count == 0) { EditorUtility.DisplayDialog("Empty Scene", "No objects found in the scene hierarchy.", "OK"); return; }
            foreach (var rootObj in rootObjects)
            {
                ProcessObjectRecursively(rootObj, prefabDict, ref replacedCount, notFoundLog, targetScene);
            }
        }

        Undo.CollapseUndoOperations(group);
        var message = $"Operation complete.\n\nReplaced: {replacedCount}\nPrefabs not found: {notFoundLog.Count}";
        EditorUtility.DisplayDialog("Result", message, "OK");
        Debug.Log($"<color=green><b>--- Mass Prefab Restore Complete ---</b></color> (Scene: {targetScene.name})");
    }

    void ProcessObjectRecursively(GameObject currentObject, Dictionary<string, GameObject> prefabDict, ref int replacedCount, List<string> notFoundLog, Scene targetScene)
    {
        if (PrefabUtility.GetPrefabInstanceStatus(currentObject) == PrefabInstanceStatus.Connected)
        {
            var children = currentObject.transform.Cast<Transform>().ToList();
            foreach (var child in children) ProcessObjectRecursively(child.gameObject, prefabDict, ref replacedCount, notFoundLog, targetScene);
            return;
        }

        var originalName = currentObject.name;
        var cleanedName = CleanObjectName(originalName);

        if (!string.IsNullOrEmpty(cleanedName) && prefabDict.TryGetValue(cleanedName, out var prefab))
        {
            ReplaceAndStopRecursion(currentObject, prefab, targetScene);
            replacedCount++;
            Debug.Log($"<color=cyan>Replaced:</color> '{originalName}' -> with prefab '{cleanedName}'");
        }
        else
        {
            if (originalName.Contains(suffixDelimiter)) { notFoundLog.Add($"Prefab '{cleanedName}' not found for object '{originalName}'"); }
            var children = currentObject.transform.Cast<Transform>().ToList();
            foreach (var child in children) ProcessObjectRecursively(child.gameObject, prefabDict, ref replacedCount, notFoundLog, targetScene);
        }
    }

    string CleanObjectName(string originalName)
    {
        string cleanedName = originalName;
        int delimiterIndex = cleanedName.IndexOf(suffixDelimiter);
        if (delimiterIndex >= 0) { cleanedName = cleanedName.Substring(0, delimiterIndex); }
        if (!string.IsNullOrEmpty(prefixToRemove) && cleanedName.StartsWith(prefixToRemove)) { cleanedName = cleanedName.Substring(prefixToRemove.Length); }
        return cleanedName.Trim();
    }

    void ReplaceAndStopRecursion(GameObject originalObj, GameObject prefab, Scene targetScene)
    {
        // 1. CAPTURE STATE into the snapshot
        var snapshot = new ObjectSnapshot
        {
            Position = originalObj.transform.position,
            Rotation = originalObj.transform.rotation,
            Scale = originalObj.transform.localScale,
            Parent = originalObj.transform.parent,
            SiblingIndex = originalObj.transform.GetSiblingIndex(),
            OriginalChildren = originalObj.transform.Cast<Transform>().ToList()
        };

        var originalRenderer = originalObj.GetComponent<MeshRenderer>();
        if (originalRenderer != null)
        {
            snapshot.Materials = originalRenderer.sharedMaterials;
        }

        // 2. INSTANTIATE the new prefab
        GameObject newInstance = (GameObject)PrefabUtility.InstantiatePrefab(prefab, targetScene);
        Undo.RegisterCreatedObjectUndo(newInstance, "Create Prefab Instance");

        // 3. APPLY STATE from the snapshot
        newInstance.transform.SetParent(snapshot.Parent, worldPositionStays: true);
        newInstance.transform.SetPositionAndRotation(snapshot.Position, snapshot.Rotation);
        newInstance.transform.localScale = snapshot.Scale;
        newInstance.transform.SetSiblingIndex(snapshot.SiblingIndex);
        newInstance.name = prefab.name;

        // --- START: FINAL, ROBUST Smart Material Restoration Logic ---
        var newRenderer = newInstance.GetComponent<MeshRenderer>();
        if (newRenderer != null && snapshot.Materials != null) // Check snapshot.Materials != null, but allow length 0
        {
            Undo.RecordObject(newRenderer, "Restore Material Overrides");

            // Get the default materials from the newly instantiated prefab.
            var defaultPrefabMaterials = newRenderer.sharedMaterials;
            // Create a brand new array to hold the final material setup.
            var finalMaterials = new Material[defaultPrefabMaterials.Length];

            // Determine the number of materials to check from the original object.
            int materialsToCheck = Mathf.Min(snapshot.Materials.Length, finalMaterials.Length);

            // 1. Iterate through the materials that can be potentially restored.
            for (int i = 0; i < materialsToCheck; i++)
            {
                // THE KEY FIX: If the original material is not null, it's a valid override. Use it.
                if (snapshot.Materials[i] != null)
                {
                    finalMaterials[i] = snapshot.Materials[i];
                }
                // If the original material IS null (purple object), keep the prefab's default material.
                else
                {
                    finalMaterials[i] = defaultPrefabMaterials[i];
                }
            }

            // 2. If the new prefab has MORE slots, fill the rest with their defaults.
            for (int i = materialsToCheck; i < finalMaterials.Length; i++)
            {
                finalMaterials[i] = defaultPrefabMaterials[i];
            }

            // Assign the newly constructed, repaired array.
            newRenderer.sharedMaterials = finalMaterials;

            // Log appropriate messages.
            if (snapshot.Materials.Length == finalMaterials.Length)
            {
                Debug.Log($"<color=lime>Materials restored for:</color> '{newInstance.name}'");
            }
            else
            {
                Debug.LogWarning($"Material slot count mismatch for '{newInstance.name}'. Partially restored materials. (Original had {snapshot.Materials.Length}, New has {finalMaterials.Length}). Please review the object.");
            }
        }
        // --- END: FINAL, ROBUST Smart Material Restoration Logic ---

        var newPrefabChildrenNames = new HashSet<string>(newInstance.transform.Cast<Transform>().Select(c => c.gameObject.name));
        foreach (var originalChild in snapshot.OriginalChildren)
        {
            if (!newPrefabChildrenNames.Contains(originalChild.gameObject.name))
            {
                Debug.Log($"<color=orange>Preserved unique child:</color> '{originalChild.name}' and re-parented under '{newInstance.name}'.");
                Undo.SetTransformParent(originalChild, newInstance.transform, "Preserve Unique Child");
            }
        }

        Undo.DestroyObjectImmediate(originalObj);
    }
}