terra-llm/Assets/LLM/source/utils/virtual_tree_builder.py

import sys
from pathlib import Path
import ruamel.yaml

# Add the utils directory to the Python path
utils_path = Path(__file__).parent
sys.path.append(str(utils_path))

from yaml_utils import load_unity_yaml, convert_to_plain_python_types

def build_virtual_tree(file_path, object_map, guid_map):
    """
    Parses a Unity scene/prefab's data to build a nested dictionary 
    representing the GameObject hierarchy with resolved component names and correct children.
    """
    nodes = {}
    transform_to_gameobject = {}
    child_to_parent_transform = {}

    # --- Pass 1: Create all nodes and map all relationships ---
    for file_id, obj_data in object_map.items():
        if 'GameObject' in obj_data:
            go_info = obj_data['GameObject']
            
            resolved_components = []
            for comp_ref in go_info.get('m_Component', []):
                comp_id = str(comp_ref['component']['fileID'])
                if comp_id not in object_map:
                    continue

                comp_data = object_map[comp_id]
                if 'MonoBehaviour' in comp_data:
                    script_guid = comp_data['MonoBehaviour'].get('m_Script', {}).get('guid')
                    if script_guid:
                        resolved_components.append({"script": {"guid": script_guid}})
                else:
                    comp_name = next((key for key in comp_data if key != 'm_ObjectHideFlags'), None)
                    if comp_name:
                        resolved_components.append(comp_name)

            nodes[file_id] = {
                'fileID': file_id,
                'm_Name': go_info.get('m_Name'),
                'm_IsActive': go_info.get('m_IsActive', 1),
                'm_TagString': go_info.get('m_TagString'),
                'm_Layer': go_info.get('m_Layer'),
                'm_Components': resolved_components,
                'm_Children': []
            }
        
        elif 'Transform' in obj_data:
            # Map the parent-child relationship between transforms
            father_id = str(obj_data['Transform'].get('m_Father', {}).get('fileID', '0'))
            if father_id != '0':
                child_to_parent_transform[file_id] = father_id
            
            # Map the transform back to its owning GameObject
            gameobject_id = str(obj_data['Transform'].get('m_GameObject', {}).get('fileID', '0'))
            if gameobject_id != '0':
                transform_to_gameobject[file_id] = gameobject_id

    # --- Pass 2: Build the hierarchy ---
    root_nodes = []
    for transform_id, gameobject_id in transform_to_gameobject.items():
        if gameobject_id in nodes:
            parent_transform_id = child_to_parent_transform.get(transform_id)
            if parent_transform_id and parent_transform_id in transform_to_gameobject:
                parent_gameobject_id = transform_to_gameobject[parent_transform_id]
                if parent_gameobject_id in nodes:
                    # This is a child, append it to its parent
                    nodes[parent_gameobject_id]['m_Children'].append(nodes[gameobject_id])
            else:
                # This is a root GameObject
                root_nodes.append(nodes[gameobject_id])

    return root_nodes