terra-llm/Assets/IntelligentProjectAnalyzer/Serializers/PrefabSerializer.cs

using UnityEngine;
using UnityEditor;
using System.Collections.Generic;

namespace IntelligentProjectAnalyzer.Serializers
{
    #region Serializable Data Structures

    /// <summary>
    /// The root object for the serialized prefab JSON.
    /// Contains the flat list of all unique components and the root of the hierarchy.
    /// </summary>
    [System.Serializable]
    public class SerializablePrefab
    {
        public List<SerializableComponent> components = new();
        public SerializableGameObject hierarchy;
    }

    /// <summary>
    /// Represents a single, unique component type found in the prefab.
    /// Its index in the main list serves as its ID.
    /// </summary>
    [System.Serializable]
    public class SerializableComponent
    {
        public string type;
    }

    /// <summary>
    /// Represents a single GameObject in the prefab's hierarchy.
    /// This structure is recursive.
    /// </summary>
    [System.Serializable]
    public class SerializableGameObject
    {
        public string Name;
        // A list of integers pointing to the indices of components in the SerializablePrefab's main component list.
        public List<int> ComponentIndices = new();
        public List<SerializableGameObject> Children = new();
    }

    #endregion

    /// <summary>
    /// A static class that provides functionality to serialize a Unity GameObject/Prefab
    /// into a custom JSON format that describes its hierarchy and component structure.
    /// </summary>
    public static class PrefabSerializer
    {
        /// <summary>
        /// A test menu item to allow easy serialization of a selected prefab.
        /// </summary>
        [MenuItem("Assets/Serialize Selected Prefab to JSON")]
        public static void SerializeSelectedPrefab()
        {
            var selectedObject = Selection.activeObject as GameObject;

            var json = Serialize(selectedObject);
            Debug.Log($"--- Serialized Prefab: {selectedObject?.name} ---\n{json}");
        }

        /// <summary>
        /// Serializes a GameObject into a JSON string with a specific format.
        /// </summary>
        /// <param name="prefabRoot">The root GameObject of the prefab to serialize.</param>
        /// <returns>A formatted JSON string representing the prefab's structure.</returns>
        public static string Serialize(GameObject prefabRoot)
        {
            if (prefabRoot == null) return "{}";
            
            // Ensure we are working with the root prefab asset, not an instance in the scene.
            if (!PrefabUtility.IsPartOfPrefabAsset(prefabRoot))
            {
                Debug.LogError("The selected object is not a prefab asset. Please select a prefab from the Project window.");
                return "{}";
            }

            // This dictionary will map a component type name to its unique index.
            var componentTypeToIndexMap = new Dictionary<string, int>();
            var serializablePrefab = new SerializablePrefab();

            // Start the recursive traversal from the root GameObject.
            serializablePrefab.hierarchy = TraverseHierarchy(prefabRoot.transform, componentTypeToIndexMap, serializablePrefab.components);

            return JsonUtility.ToJson(serializablePrefab, true);
        }
        
        public static bool TrySerialize(GameObject prefabRoot, out SerializablePrefab serializablePrefab)
        {
            serializablePrefab = new SerializablePrefab();
            
            if (prefabRoot == null) return false;
            
            // Ensure we are working with the root prefab asset, not an instance in the scene.
            if (!PrefabUtility.IsPartOfPrefabAsset(prefabRoot))
            {
                Debug.LogError("The selected object is not a prefab asset. Please select a prefab from the Project window.");
                return false;
            }

            // This dictionary will map a component type name to its unique index.
            var componentTypeToIndexMap = new Dictionary<string, int>();

            // Start the recursive traversal from the root GameObject.
            serializablePrefab.hierarchy = TraverseHierarchy(prefabRoot.transform, componentTypeToIndexMap, serializablePrefab.components);
            
            return serializablePrefab != null;
        }

        /// <summary>
        /// Recursively traverses the transform hierarchy to build the serializable structure.
        /// </summary>
        /// <param name="currentTransform">The transform of the GameObject currently being processed.</param>
        /// <param name="componentMap">The lookup dictionary for component types to their indices.</param>
        /// <param name="componentList">The master list of unique component types.</param>
        /// <returns>A SerializableGameObject representing the current transform and its children.</returns>
        private static SerializableGameObject TraverseHierarchy(Transform currentTransform, Dictionary<string, int> componentMap, List<SerializableComponent> componentList)
        {
            var serializableGo = new SerializableGameObject
            {
                Name = currentTransform.name
            };

            // Get all components on the current GameObject.
            var componentsOnGo = currentTransform.GetComponents<Component>();
            foreach (var component in componentsOnGo)
            {
                // Some components can be null if a script is missing.
                if (component == null) continue;

                var componentType = component.GetType().FullName;
                
                if (componentType == null)
                {
                    Debug.LogError($"Failed to serialize component on GameObject {currentTransform.name} because its type is null. Skipping this component.");
                    continue;
                }

                // If we haven't seen this component type before, add it to our master list.
                if (!componentMap.ContainsKey(componentType))
                {
                    componentMap[componentType] = componentList.Count;
                    componentList.Add(new SerializableComponent { type = componentType });
                }

                // Add the index of this component type to this GameObject's list.
                serializableGo.ComponentIndices.Add(componentMap[componentType]);
            }

            // Recursively process all children.
            foreach (Transform child in currentTransform)
            {
                serializableGo.Children.Add(TraverseHierarchy(child, componentMap, componentList));
            }

            return serializableGo;
        }
    }
}