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

import argparse
import sys
import json
import shutil
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 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 copy_scripts(assets_dir, output_assets_dir):
    """
    Copies all C# scripts (.cs) to the target directory.
    """
    print("\n--- Starting Script Handling ---")
    cs_files = find_files_by_extension(str(assets_dir), '.cs')
    print(f"Found {len(cs_files)} C# script files to copy.")

    for script_path_str in cs_files:
        script_path = Path(script_path_str)
        relative_path = script_path.relative_to(assets_dir)
        destination_path = output_assets_dir / relative_path
        
        destination_path.parent.mkdir(parents=True, exist_ok=True)
        
        try:
            shutil.copy(script_path, destination_path)
        except IOError as e:
            print(f"Error copying {script_path} to {destination_path}: {e}", file=sys.stderr)
    
    print("Script copying complete.")

def main():
    """
    Main function to run the low-level data extraction process.
    """
    parser = argparse.ArgumentParser(
        description="Generates a deeply detailed, per-GameObject breakdown of scenes and prefabs."
    )
    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 output folder will be saved."
    )
    parser.add_argument(
        "--indent",
        type=int,
        default=None,
        help="Indentation level for JSON output. Defaults to None (compact)."
    )
    args = parser.parse_args()

    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)

    # Create the main output folder, named "LowLevel"
    low_level_output_dir = output_dir / "LowLevel"
    output_assets_dir = low_level_output_dir / "Assets"
    try:
        output_assets_dir.mkdir(parents=True, exist_ok=True)
        print(f"Output will be saved to: {low_level_output_dir}")
    except OSError as e:
        print(f"Error: Could not create output directory '{low_level_output_dir}'. {e}", file=sys.stderr)
        sys.exit(1)

    # --- Orchestration ---
    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

    # --- Task 1: Copy C# Scripts ---
    copy_scripts(assets_dir, output_assets_dir)

    # --- Task 2: Process Scenes and Prefabs ---
    scene_files = find_files_by_extension(str(assets_dir), '.unity')
    prefab_files = find_files_by_extension(str(assets_dir), '.prefab')
    files_to_process = scene_files + prefab_files
    
    print(f"\nFound {len(files_to_process)} scene/prefab files to process.")

    for file_path_str in files_to_process:
        file_path = Path(file_path_str)
        print(f"\nProcessing Scene/Prefab: {file_path.name}")

        # --- Deep Parsing for Individual GameObjects ---
        gameobject_list = parse_scene_or_prefab(str(file_path))

        # Create the output subdirectory for this asset
        relative_path = file_path.relative_to(input_dir)
        asset_output_dir = low_level_output_dir / relative_path
        try:
            asset_output_dir.mkdir(parents=True, exist_ok=True)
        except OSError as e:
            print(f"Error creating directory {asset_output_dir}: {e}", file=sys.stderr)
            continue

        if gameobject_list:
            print(f"Saving {len(gameobject_list)} GameObjects 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"
                    try:
                        write_json(go_data, output_json_path, indent=args.indent)
                    except Exception as e:
                        print(f"Error writing to {output_json_path}: {e}", file=sys.stderr)
        else:
            print(f"Skipped deep parsing for {file_path.name}.")

        # --- Hierarchy Parsing for Root Object Identification ---
        try:
            documents = load_unity_yaml(file_path)
            if not documents:
                print(f"Could not load YAML from {file_path.name} for hierarchy parsing.")
                continue

            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()}

            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)
                print(f"Successfully saved root object list to {roots_output_path}")

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

    # --- Task 3: Process .asset files ---
    asset_files = find_files_by_extension(str(assets_dir), '.asset')
    print(f"\nFound {len(asset_files)} .asset files to process.")

    for file_path_str in asset_files:
        file_path = Path(file_path_str)
        print(f"\nProcessing Asset: {file_path.name}")

        relative_path = file_path.relative_to(input_dir)
        asset_output_dir = low_level_output_dir / relative_path
        try:
            asset_output_dir.mkdir(parents=True, exist_ok=True)
        except OSError as e:
            print(f"Error creating directory {asset_output_dir}: {e}", file=sys.stderr)
            continue

        try:
            documents = load_unity_yaml(file_path)
            if not documents:
                print(f"Skipped {file_path.name} as it's empty or could not be parsed.")
                continue

            print(f"Saving {len(documents)} objects from {file_path.name} to {asset_output_dir}")
            for doc in documents:
                if not hasattr(doc, 'anchor') or doc.anchor is None:
                    continue
                
                file_id = int(doc.anchor.value)
                obj_data = convert_to_plain_python_types(doc)
                
                final_obj_data = {}
                for key, value in obj_data.items():
                    if isinstance(value, dict):
                        new_value = value.copy()
                        new_value['fileID'] = file_id
                        final_obj_data[key] = new_value
                    else:
                        final_obj_data[key] = value

                output_json_path = asset_output_dir / f"{file_id}.json"
                write_json(final_obj_data, output_json_path, indent=args.indent)

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


    print("\nLow-level extraction complete.")

if __name__ == "__main__":
    main()