terra-llm/Assets/AssetBank/Editor/SchemaConverter/Processors/HierarchyProcessor.cs

using System.Linq;
using Newtonsoft.Json.Linq;
using System.Collections.Generic;
using AssetBank.Editor.SchemaConverter.Data.Input;
using AssetBank.Editor.SchemaConverter.Data.Output;

namespace AssetBank.Editor.SchemaConverter.Processors
{
    /// <summary>
    /// Handles the logic for building the scene hierarchy from GameObject nodes.
    /// </summary>
    public class HierarchyProcessor
    {
        private readonly List<string> _warnings = new();
        private IReadOnlyList<OriginalNode> _allNodes;
        private Dictionary<string, HierarchicalNode> _unappendedPrefabs;
        private HashSet<string> _visitedGameObjects;
        private Dictionary<string, HierarchicalNode> _nodeMap;

        private static readonly HashSet<string> s_RedundantTransformKeys = new()
        {
            "m_Father", "m_Children", "m_PrefabInstance", "m_CorrespondingSourceObject", "m_PrefabAsset"
        };

        public (List<HierarchicalNode> rootNodes, List<string> warnings, HashSet<string> consumedAnchorIds) Process(
            IReadOnlyList<OriginalNode> allNodes,
            IReadOnlyDictionary<string, OriginalNode> prefabInstances,
            List<string> prefabGuids)
        {
            _allNodes = allNodes;
            _visitedGameObjects = new HashSet<string>();
            _nodeMap = new Dictionary<string, HierarchicalNode>();
            
            var consumedAnchorIds = new HashSet<string>();
            var rootHierarchicalNodes = new List<HierarchicalNode>();
            var allGameObjects = allNodes.Where(n => n.type_id == "1").ToList();

            // 1. Identify all Transforms that are linked to a PrefabInstance and create a lookup for them.
            var allTransforms = allNodes.Where(n => n.type_id == "4").ToList();
            _unappendedPrefabs = new Dictionary<string, HierarchicalNode>();
            foreach (var transformNode in allTransforms)
            {
                var prefabInstanceId = transformNode.data?["Transform"]?["m_PrefabInstance"]?["fileID"]?.ToString();
                if (!string.IsNullOrEmpty(prefabInstanceId) && prefabInstances.TryGetValue(prefabInstanceId, out var prefabInstance))
                {
                    var guid = prefabInstance.data?["PrefabInstance"]?["m_SourcePrefab"]?["guid"]?.ToString();
                    var guidIndex = prefabGuids.IndexOf(guid);

                    if (guidIndex != -1)
                    {
                        if (!_nodeMap.TryGetValue(transformNode.anchor_id, out var prefabNode))
                        {
                            prefabNode = new HierarchicalNode
                            {
                                anchor_id = transformNode.anchor_id,
                                type_id = transformNode.type_id,
                            };
                            _nodeMap[transformNode.anchor_id] = prefabNode;
                        }
                        prefabNode.prefab_guid_index = (ulong)guidIndex;
                        PopulateNodeFields(prefabNode, transformNode);
                        _unappendedPrefabs[transformNode.anchor_id] = prefabNode;
                        consumedAnchorIds.Add(transformNode.anchor_id);
                    }
                    else if (!string.IsNullOrEmpty(guid))
                    {
                        _warnings.Add($"Found prefab with GUID '{guid}' but it was not in the provided GUID list.");
                    }
                }
            }

            // 2. Identify and create nodes for all explicit root GameObjects.
            foreach (var goNode in allGameObjects)
            {
                consumedAnchorIds.Add(goNode.anchor_id);
                var transformComponent = GetTransformComponent(goNode);
                var fatherId = GetFatherIdFromTransform(transformComponent);

                if (fatherId == "0")
                {
                    if (!_nodeMap.TryGetValue(goNode.anchor_id, out var rootNode))
                    {
                        rootNode = new HierarchicalNode { anchor_id = goNode.anchor_id, type_id = goNode.type_id };
                        PopulateNodeFields(rootNode, goNode);
                        _nodeMap[goNode.anchor_id] = rootNode;
                    }
                    rootHierarchicalNodes.Add(rootNode);
                    _visitedGameObjects.Add(rootNode.anchor_id);
                }
            }

            // 3. Recursively build the hierarchy for each root GameObject.
            // Create a copy for iteration as the collection will be modified.
            foreach (var rootNode in rootHierarchicalNodes.ToList())
            {
                BuildHierarchyRecursive(rootNode.anchor_id, allGameObjects);
            }
            
            // 4. Iteratively process all unvisited GameObjects.
            bool madeChangesInPass;
            do
            {
                madeChangesInPass = false;
                var unvisitedGameObjects = allGameObjects.Where(go => !_visitedGameObjects.Contains(go.anchor_id)).ToList();

                foreach (var goNode in unvisitedGameObjects)
                {
                    var transformComponent = GetTransformComponent(goNode);
                    var fatherId = GetFatherIdFromTransform(transformComponent);
                    var parentNode = FindParentNode(fatherId);

                    if (parentNode != null)
                    {
                        if (!_nodeMap.TryGetValue(goNode.anchor_id, out var childNode))
                        {
                            childNode = new HierarchicalNode { anchor_id = goNode.anchor_id, type_id = goNode.type_id };
                            PopulateNodeFields(childNode, goNode);
                            _nodeMap[childNode.anchor_id] = childNode;
                        }
                        
                        if (!parentNode.children.Contains(childNode))
                        {
                            parentNode.children.Add(childNode);
                        }
                        _visitedGameObjects.Add(childNode.anchor_id);
                        
                        BuildHierarchyRecursive(childNode.anchor_id, allGameObjects);
                        madeChangesInPass = true;
                    }
                }
            } while (madeChangesInPass);

            // 5. Any remaining unvisited nodes are true orphans; treat them as roots.
            var finalUnvisited = allGameObjects.Where(go => !_visitedGameObjects.Contains(go.anchor_id)).ToList();
            foreach (var goNode in finalUnvisited)
            {
                if (!_nodeMap.TryGetValue(goNode.anchor_id, out var rootNode))
                {
                    rootNode = new HierarchicalNode { anchor_id = goNode.anchor_id, type_id = goNode.type_id };
                    PopulateNodeFields(rootNode, goNode);
                    _nodeMap[rootNode.anchor_id] = rootNode;
                }
                
                rootNode.is_orphan = true;
                rootHierarchicalNodes.Add(rootNode);
                _visitedGameObjects.Add(rootNode.anchor_id);
                _warnings.Add($"GameObject '{goNode.anchor_id}' was treated as an orphan root.");
                
                BuildHierarchyRecursive(rootNode.anchor_id, allGameObjects);
            }
            
            // 5b. Prune false roots: remove any nodes from the root list that are actually children of other nodes.
            var allChildren = new HashSet<HierarchicalNode>();
            var queue = new Queue<HierarchicalNode>(rootHierarchicalNodes);
            var visitedForPruning = new HashSet<HierarchicalNode>();

            while (queue.Count > 0)
            {
                var node = queue.Dequeue();
                if (!visitedForPruning.Add(node)) continue;

                foreach (var child in node.children)
                {
                    allChildren.Add(child);
                    queue.Enqueue(child);
                }
            }
            rootHierarchicalNodes.RemoveAll(node => allChildren.Contains(node));
            
            // 6. Update nodes that are implicitly prefabs.
            UpdateImplicitPrefabInstances(rootHierarchicalNodes, prefabInstances, prefabGuids);
            
            // 7. Final Cleanup Pass
            CleanupNodes(rootHierarchicalNodes);

            return (rootHierarchicalNodes, _warnings, consumedAnchorIds);
        }

        private void BuildHierarchyRecursive(string parentAnchorId, IReadOnlyList<OriginalNode> allGameObjects)
        {
            if (!_nodeMap.TryGetValue(parentAnchorId, out var parentNode)) return;

            var parentOriginalNode = _allNodes.FirstOrDefault(n => n.anchor_id == parentAnchorId);
            if (parentOriginalNode == null) return;
            
            var parentTransformToken = GetTransformComponent(parentOriginalNode);
            if (parentOriginalNode.type_id == "4")
            {
                parentTransformToken = parentOriginalNode.data;
            }

            var childrenIds = parentTransformToken?["Transform"]?["m_Children"]?
                .Select(c => c["fileID"]?.ToString())
                .Where(id => !string.IsNullOrEmpty(id))
                .ToList();

            if (childrenIds == null) return;

            foreach (var childId in childrenIds)
            {
                var childOriginalNode = allGameObjects.FirstOrDefault(go => GetTransformId(go) == childId);

                if (childOriginalNode != null)
                {
                    if (!_nodeMap.TryGetValue(childOriginalNode.anchor_id, out var childNode))
                    {
                        childNode = new HierarchicalNode { anchor_id = childOriginalNode.anchor_id, type_id = childOriginalNode.type_id };
                        PopulateNodeFields(childNode, childOriginalNode);
                        _nodeMap[childNode.anchor_id] = childNode;
                    }
                    
                    if (!parentNode.children.Contains(childNode))
                    {
                        parentNode.children.Add(childNode);
                    }
                    
                    if (_visitedGameObjects.Add(childNode.anchor_id))
                    {
                        BuildHierarchyRecursive(childNode.anchor_id, allGameObjects);
                    }
                }
                else if (_unappendedPrefabs.TryGetValue(childId, out var prefabChildNode))
                {
                    if (!parentNode.children.Contains(prefabChildNode))
                    {
                        parentNode.children.Add(prefabChildNode);
                    }
                    _unappendedPrefabs.Remove(childId);
                    BuildHierarchyRecursive(prefabChildNode.anchor_id, allGameObjects);
                }
            }
        }
        
        private static void UpdateImplicitPrefabInstances(
            IEnumerable<HierarchicalNode> nodes,
            IReadOnlyDictionary<string, OriginalNode> prefabInstances,
            List<string> prefabGuids)
        {
            if (nodes == null) return;

            foreach (var node in nodes)
            {
                // Only try to update the node if its prefab status is unknown.
                if (!node.prefab_guid_index.HasValue)
                {
                    string prefabInstanceFileId = null;
                
                    if (node.type_id == "1" && node.Fields.TryGetValue("m_Component", out var componentsToken) && componentsToken is JArray components)
                    {
                        foreach (var component in components)
                        {
                            var fileID = component?["Transform"]?["m_PrefabInstance"]?["fileID"]?.ToString();

                            if (!string.IsNullOrEmpty(fileID) && fileID != "0")
                            {
                                prefabInstanceFileId = fileID;
                                break;
                            }
                        }
                    }

                    if (prefabInstanceFileId != null)
                    {
                        if (prefabInstances.TryGetValue(prefabInstanceFileId, out var prefabInstanceNode))
                        {
                            var guid = prefabInstanceNode.data?["PrefabInstance"]?["m_SourcePrefab"]?["guid"]?.ToString();
                            if (!string.IsNullOrEmpty(guid))
                            {
                                var guidIndex = prefabGuids.IndexOf(guid);
                                if (guidIndex != -1)
                                {
                                    node.prefab_guid_index = (ulong)guidIndex;
                                }
                            }
                        }
                    }
                }
                
                // ALWAYS recurse, regardless of the parent's status.
                if (node.children.Any())
                {
                    UpdateImplicitPrefabInstances(node.children, prefabInstances, prefabGuids);
                }
            }
        }

        private static void PopulateNodeFields(HierarchicalNode node, OriginalNode originalNode)
        {
            JObject dataBlock = null;
            if (originalNode.type_id == "1")
            {
                dataBlock = originalNode.data?["GameObject"] as JObject;
            }
            else if (originalNode.type_id == "4")
            {
                dataBlock = originalNode.data;
            }
            
            if (dataBlock == null) return;
            
            foreach (var property in dataBlock.Properties())
            {
                node.Fields[property.Name] = property.Value;
            }
        }

        private static void CleanupNodes(IEnumerable<HierarchicalNode> nodes)
        {
            foreach (var node in nodes)
            {
                if (node.Fields.TryGetValue("Transform", out var transformToken) && transformToken is JObject transformObj)
                {
                    foreach (var key in s_RedundantTransformKeys)
                    {
                        transformObj.Remove(key);
                    }
                }

                if (node.Fields.TryGetValue("m_Component", out var componentsToken) && componentsToken is JArray components)
                {
                    foreach (var component in components)
                    {
                        if (component["Transform"] is JObject cTransform)
                        {
                            foreach (var key in s_RedundantTransformKeys)
                            {
                                cTransform.Remove(key);
                            }
                        }
                    }
                }

                if (node.children.Any())
                {
                    CleanupNodes(node.children);
                }
            }
        }

        private static JToken GetTransformComponent(OriginalNode goNode)
        {
            return goNode.data?["GameObject"]?["m_Component"]?.FirstOrDefault(c => c["Transform"] != null);
        }
        
        private static string GetTransformId(OriginalNode goNode)
        {
            return GetTransformComponent(goNode)?["anchor_id"]?.ToString();
        }

        private static string GetFatherIdFromTransform(JToken transformComponent)
        {
            return transformComponent?["Transform"]?["m_Father"]?["fileID"]?.ToString() ?? "0";
        }

        private HierarchicalNode FindParentNode(string fatherTransformId)
        {
            foreach (var node in _nodeMap.Values)
            {
                if (node.type_id == "1")
                {
                    var originalNode = _allNodes.FirstOrDefault(n => n.anchor_id == node.anchor_id);
                    if (originalNode == null) continue;
                    
                    var transformId = GetTransformId(originalNode);
                    if (transformId == fatherTransformId)
                    {
                        return node;
                    }
                }
                else if (node.type_id == "4")
                {
                    if (node.anchor_id == fatherTransformId)
                    {
                        return node;
                    }
                }
            }
            return null;
        }
    }
}