terra-llm/Assets/LLM/source/orchestrators/extract_scene_trs.py

import argparse
import sys
import json
from pathlib import Path

# Add parent directories to the Python path to find utils and parsers
source_path = Path(__file__).parent.parent
sys.path.append(str(source_path / 'utils'))
sys.path.append(str(source_path / 'parsers'))

from file_utils import find_files_by_extension
from json_utils import write_json
from config_utils import load_config
from trs_processor import TRSSceneProcessor

def generate_guid_mappers(input_dir, output_dir, indent=None, shrink=False, ignored_folders=None):
    """
    Finds all .meta files and generates JSON files mapping GUIDs to asset paths.
    This is a necessary setup step to allow the TRS processor to find prefab files.
    """
    assets_dir = input_dir / "Assets"
    if not assets_dir.is_dir():
        print(f"Error: 'Assets' directory not found in '{input_dir}'", file=sys.stderr)
        return None

    print("--> Finding all .meta files for GUID mapping...")
    meta_files = find_files_by_extension(str(assets_dir), '.meta', ignored_folders=ignored_folders, project_root=input_dir)
    print(f"--> Found {len(meta_files)} .meta files.")

    guid_map = {}

    print("--> Parsing .meta files and mapping GUIDs to asset paths...")
    for meta_file_path_str in meta_files:
        meta_file_path = Path(meta_file_path_str)
        asset_file_path = Path(meta_file_path_str.rsplit('.meta', 1)[0])

        if not asset_file_path.is_file():
            continue

        guid = None
        try:
            with open(meta_file_path, 'r', encoding='utf-8') as f:
                for line in f:
                    if line.strip().startswith('guid:'):
                        guid = line.strip().split(':')[1].strip()
                        break
        except Exception as e:
            print(f"Warning: Could not read or parse guid from {meta_file_path}. {e}", file=sys.stderr)
            continue

        if guid:
            # Use the full, resolved path for the guid_map value
            guid_map[guid] = asset_file_path.resolve().as_posix()
            
    print("--> GUID mapping complete.")
    return guid_map


def main():
    """
    Main function to run the TRS-only data extraction for scenes.
    """
    parser = argparse.ArgumentParser(
        description="Generates a simplified JSON representation of scene hierarchies with only TRS data."
    )
    parser.add_argument(
        "--input",
        type=str,
        required=True,
        help="The root directory of the target Unity project."
    )
    parser.add_argument(
        "--output",
        type=str,
        required=True,
        help="The directory where the generated JSON files will be saved."
    )
    args = parser.parse_args()

    # --- Load Configuration ---
    config = load_config()
    ignored_folders = config.get('ignored_folders', [])
    shrink_json = config.get('shrink_json', False)
    indent_level = config.get('indentation_level', 4)

    input_dir = Path(args.input).resolve()
    output_dir = Path(args.output).resolve()

    if not input_dir.is_dir():
        print(f"Error: Input path '{input_dir}' is not a valid directory.", file=sys.stderr)
        sys.exit(1)

    # --- Setup Output Directory ---
    try:
        output_dir.mkdir(parents=True, exist_ok=True)
        print(f"Output will be saved to: {output_dir}")
    except OSError as e:
        print(f"Error: Could not create output directory '{output_dir}'. {e}", file=sys.stderr)
        sys.exit(1)

    assets_dir = input_dir / "Assets"
    if not assets_dir.is_dir():
        print(f"Warning: 'Assets' directory not found in '{input_dir}'. Skipping all processing.", file=sys.stderr)
        return

    print("\n--- Running Scene TRS Extraction ---")

    # --- Task 1: Generate GUID Map (for prefab lookups) ---
    print("\n[1/2] Generating GUID Mappers...")
    guid_map = generate_guid_mappers(
        input_dir, 
        output_dir, # Not saving mappers, just using the map in memory
        ignored_folders=ignored_folders
    )
    if guid_map is None:
        print("Error: Failed to generate GUID map. Aborting.", file=sys.stderr)
        sys.exit(1)
    print("--> GUID Mapper generation complete.")

    # --- Task 2: Find and Process Scene Files ---
    print("\n[2/2] Parsing Scene Hierarchies for TRS data...")
    
    scene_files = find_files_by_extension(str(assets_dir), '.unity', ignored_folders=ignored_folders, project_root=input_dir)
    
    if not scene_files:
        print("--> No scene files found to process.")
        return
        
    print(f"--> Found {len(scene_files)} scene file(s) to process.")

    total_files = len(scene_files)
    for i, file_path_str in enumerate(scene_files):
        file_path = Path(file_path_str)
        
        output_json_path = output_dir / f"{file_path.stem}.json"
        
        try:
            print(f"\n--- Processing {file_path.name} ({i+1}/{total_files}) ---")
            
            processor = TRSSceneProcessor(guid_map)
            
            print(f"    -> Loading and parsing file...")
            if not processor.load_documents(file_path):
                print(f"Warning: Could not load or parse {file_path.name}. Skipping.", file=sys.stderr)
                continue

            print(f"    -> Building hierarchy and extracting TRS data...")
            hierarchy = processor.process()
            
            # The hierarchy is the list of root objects, which will be the top-level JSON array.
            write_json(hierarchy, output_json_path, indent=indent_level, shrink=shrink_json)
            print(f"--> Successfully processed TRS hierarchy for {file_path.name} -> {output_json_path}")

        except Exception as e:
            print(f"Error processing TRS hierarchy for {file_path.name}: {e}", file=sys.stderr)

    print("\nScene TRS extraction complete.")

if __name__ == "__main__":
    main()