terra-llm/Assets/LLM/source/parsers/trs_processor.py

import sys
from pathlib import Path

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

from yaml_utils import load_unity_yaml, convert_to_plain_python_types
from hierarchy_utils import HierarchyParser

class TRSSceneProcessor:
    """
    A specialized processor that uses the robust multi-pass system from the mid-level
    extractor to build a correct scene hierarchy, then extracts only TRS data.
    """
    def __init__(self, guid_map):
        self.guid_map = guid_map
        self.object_map = {}
        self.prefab_cache = {}
        self.nodes = {}
        self.prefab_nodes = {}
        self.transform_to_gameobject = {}
        self.gameobject_to_transform = {}
        self.transform_children = {}
        self.stripped_gameobjects = {}
        self.stripped_transforms = {}
        self.prefab_instances = {}
        self.processed_relationships = set()

    def load_documents(self, file_path):
        """Loads and parses the yaml documents from a scene/prefab file."""
        documents = load_unity_yaml(file_path)
        if not documents:
            self.object_map = {}
            return False
        
        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()}
        return True

    def _load_prefab_data(self, prefab_guid):
        """Load and cache prefab data from its GUID."""
        if prefab_guid in self.prefab_cache:
            return self.prefab_cache[prefab_guid]
        
        if not prefab_guid or prefab_guid not in self.guid_map:
            self.prefab_cache[prefab_guid] = None
            return None

        prefab_path = self.guid_map[prefab_guid]
        try:
            documents = load_unity_yaml(prefab_path)
            if not documents:
                self.prefab_cache[prefab_guid] = None
                return None
            
            raw_object_map = {int(doc.anchor.value): doc for doc in documents if hasattr(doc, 'anchor') and doc.anchor is not None}
            prefab_data = {file_id: convert_to_plain_python_types(obj) for file_id, obj in raw_object_map.items()}
            self.prefab_cache[prefab_guid] = prefab_data
            return prefab_data
        except Exception:
            self.prefab_cache[prefab_guid] = None
            return None

    def build_relationship_maps(self):
        """Pass 1: Build all necessary maps for hierarchy construction."""
        for file_id, obj_data in self.object_map.items():
            if 'GameObject' in obj_data:
                go_info = obj_data['GameObject']
                self.nodes[file_id] = {'fileID': str(file_id), 'm_Name': go_info.get('m_Name', 'Unknown'), 'children': []}
                if any('stripped' in str(k) for k in obj_data.keys()):
                    self.stripped_gameobjects[file_id] = go_info
                prefab_instance_id = go_info.get('m_PrefabInstance', {}).get('fileID')
                if prefab_instance_id and not file_id in self.stripped_gameobjects:
                    self.nodes[file_id]['prefab_instance_id'] = prefab_instance_id

            elif 'PrefabInstance' in obj_data:
                prefab_info = obj_data['PrefabInstance']
                self.prefab_instances[file_id] = prefab_info
                
                # Get name from modifications or fallback to prefab asset filename
                mod_name = next((mod['value'] for mod in prefab_info.get('m_Modification', {}).get('m_Modifications', []) if mod.get('propertyPath') == 'm_Name'), None)
                
                if not mod_name:
                    guid = prefab_info.get('m_SourcePrefab', {}).get('guid')
                    if guid and guid in self.guid_map:
                        prefab_path = Path(self.guid_map[guid])
                        mod_name = prefab_path.stem
                    else:
                        mod_name = 'Unknown'

                self.prefab_nodes[file_id] = {'fileID': str(file_id), 'm_Name': mod_name, 'children': []}

            elif 'Transform' in obj_data:
                transform_info = obj_data['Transform']
                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
                
                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)
                
                if transform_info.get('m_PrefabInstance', {}).get('fileID'):
                    self.stripped_transforms[file_id] = transform_info

    def build_hierarchy(self):
        """Pass 2: Connect nodes into a hierarchy."""
        # Link standard GameObjects
        for go_id, node in self.nodes.items():
            transform_id = self.gameobject_to_transform.get(go_id)
            if not transform_id: continue
            
            child_transform_ids = self.transform_children.get(transform_id, [])
            for child_transform_id in child_transform_ids:
                child_go_id = self.transform_to_gameobject.get(child_transform_id)
                if child_go_id and child_go_id in self.nodes:
                    node['children'].append(self.nodes[child_go_id])

        # Link prefab instances to their parents
        for prefab_id, prefab_info in self.prefab_instances.items():
            parent_transform_id = prefab_info.get('m_Modification', {}).get('m_TransformParent', {}).get('fileID')
            if not parent_transform_id: continue

            parent_go_id = self.transform_to_gameobject.get(parent_transform_id)
            if parent_go_id and parent_go_id in self.nodes:
                self.nodes[parent_go_id]['children'].append(self.prefab_nodes[prefab_id])
            elif parent_transform_id in self.stripped_transforms:
                parent_prefab_instance_id = self.stripped_transforms[parent_transform_id].get('m_PrefabInstance', {}).get('fileID')
                if parent_prefab_instance_id in self.prefab_nodes:
                    self.prefab_nodes[parent_prefab_instance_id]['children'].append(self.prefab_nodes[prefab_id])

        # Link scene objects that are children of prefabs
        for go_id, node in 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 parent_transform_id in self.stripped_transforms:
                parent_prefab_instance_id = self.stripped_transforms[parent_transform_id].get('m_PrefabInstance', {}).get('fileID')
                if parent_prefab_instance_id in self.prefab_nodes:
                    self.prefab_nodes[parent_prefab_instance_id]['children'].append(node)

    def merge_prefab_info(self):
        """Pass 3: Merge root GameObject data into PrefabInstance nodes."""
        nodes_to_delete = []
        for go_id, go_node in self.nodes.items():
            prefab_instance_id = go_node.get('prefab_instance_id')
            if prefab_instance_id and prefab_instance_id in self.prefab_nodes:
                prefab_node = self.prefab_nodes[prefab_instance_id]
                prefab_info = self.prefab_instances[prefab_instance_id]
                
                # Name is now set in build_relationship_maps. We just merge the rest.
                prefab_node['source_prefab_guid'] = prefab_info.get('m_SourcePrefab', {}).get('guid')
                prefab_node['children'].extend(go_node.get('children', []))
                nodes_to_delete.append(go_id)

        for go_id in nodes_to_delete:
            del self.nodes[go_id]

    def extract_trs_data(self, nodes):
        """Pass 4: Recursively find and attach TRS data to each node."""
        for node in nodes:
            file_id = int(node['fileID'])
            
            # Default TRS values
            trs = {"position": {"x": 0, "y": 0, "z": 0}, "rotation": {"x": 0, "y": 0, "z": 0, "w": 1}, "scale": {"x": 1, "y": 1, "z": 1}}

            if 'source_prefab_guid' in node: # It's a prefab instance
                prefab_instance_info = self.prefab_instances.get(file_id, {})
                modifications = prefab_instance_info.get('m_Modification', {}).get('m_Modifications', [])
                
                # Find the root transform of the prefab to get base TRS
                prefab_data = self._load_prefab_data(node['source_prefab_guid'])
                if prefab_data:
                    for obj in prefab_data.values():
                        if 'Transform' in obj and obj['Transform'].get('m_Father', {}).get('fileID') == 0:
                            base_transform = obj['Transform']
                            trs['position'] = base_transform.get('m_LocalPosition', trs['position'])
                            trs['rotation'] = base_transform.get('m_LocalRotation', trs['rotation'])
                            trs['scale'] = base_transform.get('m_LocalScale', trs['scale'])
                            break
                
                # Apply modifications
                for mod in modifications:
                    prop = mod.get('propertyPath', '')

                    # Skip any property that is not part of a Transform
                    if not ('m_LocalPosition' in prop or 'm_LocalRotation' in prop or 'm_LocalScale' in prop):
                        continue
                    
                    try:
                        val = float(mod.get('value', 0))
                    except (ValueError, TypeError):
                        # If the value is not a valid float, skip this modification
                        print(f"Warning: Could not convert value '{mod.get('value')}' for property '{prop}' to float. Skipping.", file=sys.stderr)
                        continue

                    if prop == 'm_LocalPosition.x': trs['position']['x'] = val
                    elif prop == 'm_LocalPosition.y': trs['position']['y'] = val
                    elif prop == 'm_LocalPosition.z': trs['position']['z'] = val
                    elif prop == 'm_LocalRotation.x': trs['rotation']['x'] = val
                    elif prop == 'm_LocalRotation.y': trs['rotation']['y'] = val
                    elif prop == 'm_LocalRotation.z': trs['rotation']['z'] = val
                    elif prop == 'm_LocalRotation.w': trs['rotation']['w'] = val
                    elif prop == 'm_LocalScale.x': trs['scale']['x'] = val
                    elif prop == 'm_LocalScale.y': trs['scale']['y'] = val
                    elif prop == 'm_LocalScale.z': trs['scale']['z'] = val

            else: # It's a regular GameObject
                transform_id = self.gameobject_to_transform.get(file_id)
                if transform_id and transform_id in self.object_map:
                    transform_info = self.object_map[transform_id].get('Transform', {})
                    trs['position'] = transform_info.get('m_LocalPosition', trs['position'])
                    trs['rotation'] = transform_info.get('m_LocalRotation', trs['rotation'])
                    trs['scale'] = transform_info.get('m_LocalScale', trs['scale'])
            
            node['trs'] = trs
            
            if 'children' in node and node['children']:
                self.extract_trs_data(node['children'])

    def cleanup_hierarchy(self, nodes):
        """Pass 5: Remove intermediate data, leaving only the desired fields."""
        final_nodes = []
        for node in nodes:
            clean_node = {
                'name': node.get('m_Name', 'Unknown'),
                'trs': node.get('trs', {}),
                'children': self.cleanup_hierarchy(node.get('children', []))
            }
            if 'source_prefab_guid' in node:
                clean_node['source_prefab_guid'] = node['source_prefab_guid']
            final_nodes.append(clean_node)
        return final_nodes

    def process(self):
        """Main processing pipeline."""
        if not self.object_map:
            return []

        # --- Run all passes ---
        self.build_relationship_maps()
        self.build_hierarchy()
        self.merge_prefab_info()
        
        # Get root nodes for the final passes
        parser = HierarchyParser(self.object_map)
        root_object_ids = parser.get_root_object_ids()
        
        all_nodes = {**self.nodes, **self.prefab_nodes}
        root_nodes = [all_nodes[file_id] for file_id, _ in root_object_ids if file_id in all_nodes]

        self.extract_trs_data(root_nodes)
        final_hierarchy = self.cleanup_hierarchy(root_nodes)

        return final_hierarchy