terra-llm/Assets/LLM/source/scene_processor.py

import argparse
import sys
import json
from pathlib import Path

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

from yaml_utils import load_unity_yaml, convert_to_plain_python_types

class UnitySceneProcessor:
    def __init__(self, guid_map):
        self.guid_map = guid_map
        self.object_map = {}
        self.nodes = {}
        self.prefab_nodes = {}
        self.transform_to_gameobject = {}
        self.gameobject_to_transform = {}
        self.transform_children = {}
        self.stripped_gameobjects = {}
        self.processed_relationships = set()

    def load_prefab_data(self, prefab_guid):
        """Load and parse prefab data from GUID"""
        if not prefab_guid or prefab_guid not in self.guid_map:
            return {}
        
        prefab_path = self.guid_map[prefab_guid]
        try:
            documents = load_unity_yaml(prefab_path)
            if not documents:
                return {}
            
            raw_object_map = {int(doc.anchor.value): doc for doc in documents if hasattr(doc, 'anchor') and doc.anchor is not None}
            return {file_id: convert_to_plain_python_types(obj) for file_id, obj in raw_object_map.items()}
        except Exception:
            return {}

    def calculate_deep_sibling_index(self, scene_transform_id, prefab_guid):
        """Calculate deep sibling index for scene objects that are children of prefab objects"""
        scene_transform = self.object_map.get(scene_transform_id, {}).get('Transform', {})
        scene_root_order = scene_transform.get('m_RootOrder', 0)
        
        # Get parent transform (should be stripped)
        parent_transform_id = scene_transform.get('m_Father', {}).get('fileID')
        if not parent_transform_id:
            return str(scene_root_order)
        
        parent_transform = self.object_map.get(parent_transform_id, {}).get('Transform', {})
        corresponding_source = parent_transform.get('m_CorrespondingSourceObject', {})
        prefab_transform_id = corresponding_source.get('fileID')
        
        if not prefab_transform_id:
            return str(scene_root_order)
        
        # Load prefab data and traverse hierarchy
        prefab_data = self.load_prefab_data(prefab_guid)
        if not prefab_data:
            return str(scene_root_order)
        
        # Build sibling index by traversing up the prefab hierarchy
        sibling_indices = []
        current_transform_id = prefab_transform_id
        
        while current_transform_id and current_transform_id in prefab_data:
            transform_data = prefab_data.get(current_transform_id, {}).get('Transform', {})
            if not transform_data:
                break
            
            root_order = transform_data.get('m_RootOrder', 0)
            sibling_indices.insert(0, str(root_order))
            
            # Move to parent
            parent_id = transform_data.get('m_Father', {}).get('fileID')
            if not parent_id or parent_id == 0:
                break
            current_transform_id = parent_id
        
        # Add the scene object's own root order at the end
        sibling_indices.append(str(scene_root_order))
        return '-'.join(sibling_indices)

    def process_first_pass(self):
        """First pass: Build relationship maps and create basic nodes"""
        stripped_transforms = {}
        stripped_gameobjects = {}
        prefab_instances = {}
        
        for file_id, obj_data in self.object_map.items():
            if 'GameObject' in obj_data:
                go_info = obj_data['GameObject']
                
                # Always create a node for a GameObject.
                # If it's part of a prefab, we'll link it later.
                self.nodes[file_id] = {
                    'fileID': str(file_id),
                    'm_Name': go_info.get('m_Name', 'Unknown'),
                    'm_IsActive': go_info.get('m_IsActive', 1),
                    'm_TagString': go_info.get('m_TagString', 'Untagged'),
                    'm_Layer': go_info.get('m_Layer', 0),
                    'components': [],
                    'children': []
                }
                
                # If it's a stripped GameObject, track it for component linking
                is_stripped = any('stripped' in str(key) for key in obj_data.keys() if hasattr(key, '__str__'))
                if is_stripped:
                    prefab_instance_id = go_info.get('m_PrefabInstance', {}).get('fileID')
                    if prefab_instance_id:
                        stripped_gameobjects[file_id] = {
                            'prefab_instance_id': prefab_instance_id,
                            'm_CorrespondingSourceObject': go_info.get('m_CorrespondingSourceObject', {})
                        }
                
                # If it's the root GameObject of a prefab instance, add a link for merging
                prefab_instance_id = go_info.get('m_PrefabInstance', {}).get('fileID')
                if prefab_instance_id and not is_stripped:
                    self.nodes[file_id]['prefab_instance_id'] = prefab_instance_id
                
            elif 'PrefabInstance' in obj_data:
                prefab_info = obj_data['PrefabInstance']
                source_prefab = prefab_info.get('m_SourcePrefab', {})
                modifications = prefab_info.get('m_Modification', {})
                
                # Extract m_IsActive, m_Name, and m_RootOrder from modifications
                m_is_active = 1  # default
                m_name = None
                m_root_order = 999999 # default
                for mod in modifications.get('m_Modifications', []):
                    if mod.get('propertyPath') == 'm_IsActive':
                        m_is_active = mod.get('value', 1)
                    elif mod.get('propertyPath') == 'm_Name':
                        m_name = mod.get('value')
                    elif mod.get('propertyPath') == 'm_RootOrder':
                        m_root_order = mod.get('value', 999999)

                prefab_instances[file_id] = {
                    'source_prefab_guid': source_prefab.get('guid'),
                    'm_TransformParent': modifications.get('m_TransformParent', {}).get('fileID'),
                    'm_IsActive': m_is_active,
                    'm_Name': m_name,
                    'm_RootOrder': m_root_order
                }
                
            elif 'Transform' in obj_data:
                transform_info = obj_data['Transform']
                
                # Check if this is a stripped transform (has m_PrefabInstance)
                prefab_instance_id = transform_info.get('m_PrefabInstance', {}).get('fileID')
                if prefab_instance_id:
                    stripped_transforms[file_id] = {
                        'prefab_instance_id': prefab_instance_id,
                        'm_CorrespondingSourceObject': transform_info.get('m_CorrespondingSourceObject', {})
                    }
                
                # Build transform mappings
                gameobject_id = transform_info.get('m_GameObject', {}).get('fileID')
                if gameobject_id:
                    self.transform_to_gameobject[file_id] = gameobject_id
                    self.gameobject_to_transform[gameobject_id] = file_id
                
                # Build parent-child relationships
                parent_id = transform_info.get('m_Father', {}).get('fileID')
                if parent_id and parent_id != 0:
                    if parent_id not in self.transform_children:
                        self.transform_children[parent_id] = []
                    self.transform_children[parent_id].append(file_id)
        
        # Create prefab nodes
        for prefab_id, prefab_info in prefab_instances.items():
            self.prefab_nodes[prefab_id] = {
                'fileID': str(prefab_id),
                'source_prefab_guid': prefab_info['source_prefab_guid'],
                'm_IsActive': prefab_info['m_IsActive'],
                'm_RootOrder': prefab_info.get('m_RootOrder', 999999),
                'components': [],  # Will be populated with override components
                'children': []
            }
            if prefab_info.get('m_Name'):
                self.prefab_nodes[prefab_id]['m_Name'] = prefab_info['m_Name']
        
        # Store for second pass
        self.stripped_transforms = stripped_transforms
        self.stripped_gameobjects = stripped_gameobjects
        self.prefab_instances = prefab_instances

    def process_second_pass(self):
        """Second pass: Build hierarchy relationships"""
        
        # 1. Handle standard GameObject parent-child relationships
        for go_id, node in self.nodes.items():
            transform_id = self.gameobject_to_transform.get(go_id)
            if not transform_id:
                continue
            
            # Get child transforms
            child_transform_ids = self.transform_children.get(transform_id, [])
            for child_transform_id in child_transform_ids:
                
                # Check if child transform is stripped (part of prefab)
                if child_transform_id in self.stripped_transforms:
                    stripped_info = self.stripped_transforms[child_transform_id]
                    prefab_instance_id = stripped_info['prefab_instance_id']
                    
                    if prefab_instance_id in self.prefab_nodes:
                        relationship_key = f"{go_id}->{prefab_instance_id}"
                        if relationship_key not in self.processed_relationships:
                            node['children'].append(self.prefab_nodes[prefab_instance_id])
                            self.processed_relationships.add(relationship_key)
                else:
                    # Regular GameObject child
                    child_go_id = self.transform_to_gameobject.get(child_transform_id)
                    if child_go_id and child_go_id in self.nodes:
                        relationship_key = f"{go_id}->{child_go_id}"
                        if relationship_key not in self.processed_relationships:
                            node['children'].append(self.nodes[child_go_id])
                            self.processed_relationships.add(relationship_key)
        
        # 2. Handle prefab-to-parent relationships
        for prefab_id, prefab_info in self.prefab_instances.items():
            parent_transform_id = prefab_info['m_TransformParent']
            if not parent_transform_id or parent_transform_id == 0:
                continue  # Root prefab, will be handled in get_root_objects
            
            # Find parent GameObject or parent prefab
            parent_go_id = self.transform_to_gameobject.get(parent_transform_id)
            if parent_go_id and parent_go_id in self.nodes:
                # Prefab child of GameObject
                relationship_key = f"{parent_go_id}->{prefab_id}"
                if relationship_key not in self.processed_relationships:
                    self.nodes[parent_go_id]['children'].append(self.prefab_nodes[prefab_id])
                    self.processed_relationships.add(relationship_key)
            elif parent_transform_id in self.stripped_transforms:
                # Prefab child of another prefab
                stripped_info = self.stripped_transforms[parent_transform_id]
                parent_prefab_id = stripped_info['prefab_instance_id']
                
                if parent_prefab_id in self.prefab_nodes:
                    relationship_key = f"{parent_prefab_id}->{prefab_id}"
                    if relationship_key not in self.processed_relationships:
                        self.prefab_nodes[parent_prefab_id]['children'].append(self.prefab_nodes[prefab_id])
                        self.processed_relationships.add(relationship_key)
        
        # 3. Handle scene objects as children of prefabs (with deep sibling index)
        for go_id, node in self.nodes.items():
            transform_id = self.gameobject_to_transform.get(go_id)
            if not transform_id:
                continue
            
            # Find parent transform
            parent_transform_id = None
            for parent_id, children in self.transform_children.items():
                if transform_id in children:
                    parent_transform_id = parent_id
                    break
            
            if parent_transform_id and parent_transform_id in self.stripped_transforms:
                stripped_info = self.stripped_transforms[parent_transform_id]
                prefab_instance_id = stripped_info['prefab_instance_id']
                
                if prefab_instance_id in self.prefab_nodes:
                    # Calculate deep sibling index
                    prefab_guid = self.prefab_instances[prefab_instance_id]['source_prefab_guid']
                    deep_sibling_index = self.calculate_deep_sibling_index(transform_id, prefab_guid)
                    if deep_sibling_index:
                        node['deep_sibling_index'] = deep_sibling_index
                    
                    # Add as child of prefab
                    relationship_key = f"{prefab_instance_id}->{go_id}"
                    if relationship_key not in self.processed_relationships:
                        self.prefab_nodes[prefab_instance_id]['children'].append(node)
                        self.processed_relationships.add(relationship_key)

    def process_third_pass(self):
        """Third pass: Extract components for both GameObjects and PrefabInstances"""
        for file_id, obj_data in self.object_map.items():
            # Skip GameObjects, Transforms, and PrefabInstances
            if any(key in obj_data for key in ['GameObject', 'Transform', 'PrefabInstance']):
                continue
            
            # Find the GameObject this component belongs to
            gameobject_id = None
            for key, data in obj_data.items():
                if isinstance(data, dict) and 'm_GameObject' in data:
                    gameobject_id = data['m_GameObject'].get('fileID')
                    break
            
            if not gameobject_id:
                continue
            
            component_name = next((key for key in obj_data if key not in ['m_ObjectHideFlags']), 'Unknown')
            
            # Handle MonoBehaviour with script GUID
            if component_name == 'MonoBehaviour':
                script_info = obj_data['MonoBehaviour'].get('m_Script', {})
                script_guid = script_info.get('guid')
                if script_guid:
                    component_entry = script_guid
                else:
                    component_entry = "MonoBehaviour"
            else:
                component_entry = component_name
            
            # Check if this component belongs to a regular GameObject
            if gameobject_id in self.nodes:
                if component_entry not in self.nodes[gameobject_id]['components']:
                    self.nodes[gameobject_id]['components'].append(component_entry)
            
            # Check if this component belongs to a stripped GameObject (prefab override)
            elif gameobject_id in self.stripped_gameobjects:
                stripped_info = self.stripped_gameobjects[gameobject_id]
                prefab_instance_id = stripped_info['prefab_instance_id']
                
                if prefab_instance_id in self.prefab_nodes:
                    if component_entry not in self.prefab_nodes[prefab_instance_id]['components']:
                        self.prefab_nodes[prefab_instance_id]['components'].append(component_entry)
            
            # If neither, this component belongs to an unknown GameObject (potential orphan)
            else:
                print(f"Warning: Component {component_name} (ID: {file_id}) references unknown GameObject {gameobject_id}")

    def verification_pass(self):
        """
        Verifies and fixes parent-child relationships that might have been missed
        by using a direct child->parent lookup. This acts as a patch for the
        less reliable reverse-lookup used in process_second_pass.
        """
        for go_id, node in list(self.nodes.items()):
            transform_id = self.gameobject_to_transform.get(go_id)
            if not transform_id:
                continue

            transform_info = self.object_map.get(transform_id, {}).get('Transform', {})
            parent_transform_id = transform_info.get('m_Father', {}).get('fileID')

            if not parent_transform_id or parent_transform_id == 0:
                continue  # This is a root object, no parent to verify.

            # Case 1: Parent is a PrefabInstance (via a stripped transform)
            if parent_transform_id in self.stripped_transforms:
                stripped_info = self.stripped_transforms[parent_transform_id]
                parent_prefab_id = stripped_info['prefab_instance_id']
                if parent_prefab_id in self.prefab_nodes:
                    relationship_key = f"{parent_prefab_id}->{go_id}"
                    if relationship_key not in self.processed_relationships:
                        # Add deep sibling index if it's a child of a prefab
                        prefab_guid = self.prefab_instances[parent_prefab_id]['source_prefab_guid']
                        deep_sibling_index = self.calculate_deep_sibling_index(transform_id, prefab_guid)
                        if deep_sibling_index:
                            node['deep_sibling_index'] = deep_sibling_index
                        
                        self.prefab_nodes[parent_prefab_id]['children'].append(node)
                        self.processed_relationships.add(relationship_key)

            # Case 2: Parent is a regular GameObject
            else:
                parent_go_id = self.transform_to_gameobject.get(parent_transform_id)
                if parent_go_id and parent_go_id in self.nodes:
                    relationship_key = f"{parent_go_id}->{go_id}"
                    if relationship_key not in self.processed_relationships:
                        self.nodes[parent_go_id]['children'].append(node)
                        self.processed_relationships.add(relationship_key)

    def merge_prefab_data_pass(self):
        """Fourth pass: Merge GameObject data into their corresponding PrefabInstance nodes"""
        nodes_to_delete = []
        for go_id, go_node in self.nodes.items():
            prefab_instance_id = go_node.get('prefab_instance_id')
            if not prefab_instance_id:
                continue

            if prefab_instance_id in self.prefab_nodes:
                # Merge the GameObject data into the PrefabInstance node
                # The prefab node's existing data takes precedence
                prefab_node = self.prefab_nodes[prefab_instance_id]
                
                # Create a copy of the gameobject node to avoid modifying during iteration
                go_node_copy = go_node.copy()
                
                # Remove keys that should not be merged or are already handled
                del go_node_copy['fileID']
                del go_node_copy['prefab_instance_id']
                
                # Update prefab_node with go_node_copy data
                # This will overwrite keys in prefab_node with values from go_node_copy
                # if they exist in both. We want the opposite.
                
                # Let's do it manually to ensure prefab data is kept
                for key, value in go_node_copy.items():
                    if key not in prefab_node:
                        prefab_node[key] = value
                
                # Specifically handle children and components to merge them
                prefab_node['children'].extend(go_node_copy.get('children', []))
                
                # Merge components, avoiding duplicates
                existing_components = set(str(c) for c in prefab_node.get('components', []))
                for comp in go_node_copy.get('components', []):
                    if str(comp) not in existing_components:
                        prefab_node['components'].append(comp)
                
                # Rename 'components' to 'addedComponents' for clarity
                if 'components' in prefab_node:
                    prefab_node['addedComponents'] = prefab_node.pop('components')

                # Mark the original GameObject node for deletion
                nodes_to_delete.append(go_id)

        # Remove the now-redundant GameObject nodes
        for go_id in nodes_to_delete:
            del self.nodes[go_id]

    def get_root_objects(self):
        """Get all root-level objects, sorted by m_RootOrder, and handle orphans"""
        root_objects = []
        all_children_ids = set()
        
        # Collect all child IDs to identify orphans
        for go_id, node in self.nodes.items():
            self._collect_child_ids(node, all_children_ids)
        for prefab_id, prefab_node in self.prefab_nodes.items():
            self._collect_child_ids(prefab_node, all_children_ids)
        
        # Find GameObjects that have no parent and collect their m_RootOrder
        gameobject_roots = []
        for go_id, node in self.nodes.items():
            transform_id = self.gameobject_to_transform.get(go_id)
            if not transform_id:
                # No transform - likely orphan
                if go_id not in all_children_ids:
                    node['is_orphan'] = True
                    gameobject_roots.append((node, 999999))  # Orphans go to end
                continue
            
            # Check if this transform has a parent
            has_parent = False
            for parent_id, children in self.transform_children.items():
                if transform_id in children:
                    has_parent = True
                    break
            
            # Also check if it's a child of any prefab
            is_prefab_child = False
            for parent_transform_id in self.stripped_transforms:
                if transform_id in self.transform_children.get(parent_transform_id, []):
                    is_prefab_child = True
                    break
            
            if not has_parent and not is_prefab_child:
                if go_id not in all_children_ids:
                    # Get m_RootOrder from transform
                    root_order = self._get_root_order(transform_id)
                    gameobject_roots.append((node, root_order))
                else:
                    # This is an orphan that was somehow referenced but not properly parented
                    node['is_orphan'] = True
                    gameobject_roots.append((node, 999999))  # Orphans go to end
        
        # Find root prefab instances and collect their m_RootOrder
        prefab_roots = []
        for prefab_id, prefab_info in self.prefab_instances.items():
            parent_transform_id = prefab_info['m_TransformParent']
            if not parent_transform_id or parent_transform_id == 0:
                if prefab_id in self.prefab_nodes:
                    if prefab_id not in all_children_ids:
                        # Get m_RootOrder directly from the prefab node
                        root_order = self.prefab_nodes[prefab_id].get('m_RootOrder', 999999)
                        prefab_roots.append((self.prefab_nodes[prefab_id], root_order))
                    else:
                        # Orphan prefab
                        self.prefab_nodes[prefab_id]['is_orphan'] = True
                        prefab_roots.append((self.prefab_nodes[prefab_id], 999999))  # Orphans go to end
        
        # Check for completely disconnected GameObjects (orphans)
        for go_id, node in self.nodes.items():
            if go_id not in all_children_ids and not any(obj[0].get('fileID') == str(go_id) for obj in gameobject_roots):
                node['is_orphan'] = True
                gameobject_roots.append((node, 999999))  # Orphans go to end
        
        # Check for completely disconnected PrefabInstances (orphans)
        for prefab_id, prefab_node in self.prefab_nodes.items():
            if prefab_id not in all_children_ids and not any(obj[0].get('fileID') == str(prefab_id) for obj in prefab_roots):
                prefab_node['is_orphan'] = True
                prefab_roots.append((prefab_node, 999999))  # Orphans go to end
        
        # Combine and sort all root objects by m_RootOrder
        all_roots = gameobject_roots + prefab_roots
        all_roots.sort(key=lambda x: x[1])  # Sort by m_RootOrder (second element of tuple)
        
        return [obj[0] for obj in all_roots]  # Return only the objects, not the tuples

    def _get_root_order(self, transform_id):
        """Get m_RootOrder from a transform"""
        if transform_id not in self.object_map:
            return 999999  # Default for missing transforms
        
        transform_data = self.object_map[transform_id].get('Transform', {})
        return transform_data.get('m_RootOrder', 999999)

    def _collect_child_ids(self, node, child_ids_set):
        """Recursively collect all child IDs from a node tree"""
        for child in node.get('children', []):
            child_id = child.get('fileID')
            if child_id:
                # Convert to int for comparison with our keys
                try:
                    child_ids_set.add(int(child_id))
                except ValueError:
                    pass
            self._collect_child_ids(child, child_ids_set)

    def cleanup_pass(self, nodes):
        """
        Recursively cleans up temporary or internal properties from the final node structure.
        """
        # List of properties to remove from the final output
        cleanup_keys = ['m_RootOrder']

        for node in nodes:
            for key in cleanup_keys:
                if key in node:
                    del node[key]
            
            if 'children' in node and node['children']:
                self.cleanup_pass(node['children'])

    def process_file(self, file_path):
        """Main processing method"""
        # Load and parse the file
        documents = load_unity_yaml(file_path)
        if not documents:
            return []

        # Build object map
        raw_object_map = {int(doc.anchor.value): doc for doc in documents if hasattr(doc, 'anchor') and doc.anchor is not None}
        self.object_map = {file_id: convert_to_plain_python_types(obj) for file_id, obj in raw_object_map.items()}

        # Process in passes
        self.process_first_pass()
        self.process_second_pass()
        self.verification_pass()
        self.process_third_pass()
        self.merge_prefab_data_pass()

        # Get the final, sorted root objects
        root_objects = self.get_root_objects()
        
        # Run the final cleanup pass
        self.cleanup_pass(root_objects)

        return root_objects


def main():
    parser = argparse.ArgumentParser(description="Process Unity scene/prefab files into JSON representation")
    parser.add_argument("--input", required=True, help="Path to input .unity or .prefab file")
    parser.add_argument("--guid-map", required=True, help="Path to guid_map.json file")
    parser.add_argument("--output", required=True, help="Path to output .json file")
    
    args = parser.parse_args()
    
    input_path = Path(args.input)
    guid_map_path = Path(args.guid_map)
    output_path = Path(args.output)
    
    if not input_path.exists():
        print(f"Error: Input file {input_path} does not exist", file=sys.stderr)
        sys.exit(1)
    
    if not guid_map_path.exists():
        print(f"Error: GUID map file {guid_map_path} does not exist", file=sys.stderr)
        sys.exit(1)
    
    # Load GUID map
    try:
        with open(guid_map_path, 'r', encoding='utf-8') as f:
            guid_map = json.load(f)
    except Exception as e:
        print(f"Error loading GUID map: {e}", file=sys.stderr)
        sys.exit(1)
    
    # Process the file
    processor = UnitySceneProcessor(guid_map)
    try:
        result = processor.process_file(input_path)
        
        # Ensure output directory exists
        output_path.parent.mkdir(parents=True, exist_ok=True)
        
        # Write output
        with open(output_path, 'w', encoding='utf-8') as f:
            json.dump(result, f, indent=2, ensure_ascii=False)
        
        print(f"Successfully processed {input_path.name} -> {output_path}")
        
    except Exception as e:
        print(f"Error processing file {input_path}: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()