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

import argparse
import sys
import json
from pathlib import Path

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

from file_utils import find_files_by_extension
from deep_parser import parse_scene_or_prefab
from json_utils import write_json
from yaml_utils import load_unity_yaml, convert_to_plain_python_types
from hierarchy_utils import HierarchyParser

def main():
    parser = argparse.ArgumentParser(description="Parses scenes and prefabs into a per-GameObject breakdown.")
    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 output will be saved.")
    parser.add_argument("--indent", type=int, default=None)
    parser.add_argument("--shrink-json", action='store_true')
    parser.add_argument("--ignored-folders", type=str, default='[]', help='JSON string of folders to ignore.')
    args = parser.parse_args()

    try:
        ignored_folders = json.loads(args.ignored_folders)
    except json.JSONDecodeError:
        print("Warning: Could not parse ignored-folders JSON string. Defaulting to empty list.", file=sys.stderr)
        ignored_folders = []

    input_dir = Path(args.input).resolve()
    output_dir = Path(args.output).resolve()
    assets_dir = input_dir / "Assets"

    if not assets_dir.is_dir():
        return

    scene_files = find_files_by_extension(str(assets_dir), '.unity', ignored_folders=ignored_folders, project_root=input_dir)
    prefab_files = find_files_by_extension(str(assets_dir), '.prefab', ignored_folders=ignored_folders, project_root=input_dir)
    files_to_process = scene_files + prefab_files

    if not files_to_process:
        print("No scene or prefab files found.")
        return
        
    print(f"\n--- Starting Scene/Prefab Parsing ---")
    print(f"Found {len(files_to_process)} files to process.")

    total_files = len(files_to_process)
    for i, file_path_str in enumerate(files_to_process):
        file_path = Path(file_path_str)
        print(f"\n--- Processing file {i+1}/{total_files}: {file_path.relative_to(input_dir)} ---")

        # --- Step 1: Deep Parsing ---
        print("  [1/2] Starting deep parse to extract all GameObjects...")
        gameobject_list = parse_scene_or_prefab(str(file_path))

        relative_path = file_path.relative_to(input_dir)
        asset_output_dir = output_dir / relative_path
        asset_output_dir.mkdir(parents=True, exist_ok=True)

        if gameobject_list:
            print(f"    -> Deep parse complete. Found {len(gameobject_list)} GameObjects.")
            print(f"    -> Saving individual GameObject files to {asset_output_dir}...")
            for go_data in gameobject_list:
                file_id = go_data.get('fileID')
                if file_id:
                    output_json_path = asset_output_dir / f"{file_id}.json"
                    write_json(go_data, output_json_path, indent=args.indent, shrink=args.shrink_json)
            print(f"    -> Finished saving GameObject files.")
        else:
            print(f"    -> Skipped deep parsing for {file_path.name} (no objects found or file was empty).")

        # --- Step 2: Hierarchy Parsing ---
        print("  [2/2] Starting hierarchy parse to identify root objects...")
        try:
            documents = load_unity_yaml(file_path)
            if not documents:
                print("    -> Skipping hierarchy parse (file is empty or could not be read).")
                continue

            print("    -> Converting YAML for hierarchy analysis...")
            raw_object_map = {int(doc.anchor.value): doc for doc in documents if hasattr(doc, 'anchor') and doc.anchor is not None}
            object_map = {file_id: convert_to_plain_python_types(obj) for file_id, obj in raw_object_map.items()}

            print("    -> Analyzing transform hierarchy to find roots...")
            parser = HierarchyParser(object_map)
            root_object_ids = parser.get_root_object_ids()
            
            root_ids_list = [file_id for file_id, _ in root_object_ids]

            if root_ids_list:
                roots_output_path = asset_output_dir / "root_objects.json"
                write_json(root_ids_list, roots_output_path, indent=args.indent, shrink=args.shrink_json)
                print(f"    -> Successfully saved {len(root_ids_list)} root object IDs to {roots_output_path}")
            else:
                print("    -> No root objects found.")

        except Exception as e:
            print(f"    -> Error during hierarchy parsing for {file_path.name}: {e}", file=sys.stderr)

    print("\nScene and prefab parsing complete.")

if __name__ == "__main__":
    main()