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

import argparse
import sys
import json
import shutil
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 replicate_directory_structure, find_files_by_extension, create_guid_to_path_map
from json_utils import write_json
from config_utils import load_config
from scene_processor import UnitySceneProcessor

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.
    """
    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

    print("--> Finding all .meta files...")
    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 to process.")

    asset_type_map = {
        '.prefab': 'prefabs', '.unity': 'scenes', '.mat': 'materials',
        '.cs': 'scripts', '.png': 'textures', '.jpg': 'textures',
        '.jpeg': 'textures', '.asset': 'scriptable_objects',
    }
    
    guid_maps = {value: {} for value in asset_type_map.values()}
    guid_maps['others'] = {}

    print("--> Parsing .meta files and mapping GUIDs to asset paths...")
    total_files = len(meta_files)
    for i, meta_file_path_str in enumerate(meta_files):
        if (i + 1) % 250 == 0:
            print(f"    ...processed {i+1}/{total_files} .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:
            asset_ext = asset_file_path.suffix.lower()
            asset_type = asset_type_map.get(asset_ext, 'others')
            # Use the full path from the project root for the guid_map value
            full_asset_path = input_dir / asset_file_path
            guid_maps[asset_type][guid] = full_asset_path.as_posix()
    print(f"    ...finished processing all {total_files} .meta files.")

    mappers_dir = output_dir / "GuidMappers"
    try:
        mappers_dir.mkdir(parents=True, exist_ok=True)
        print(f"--> Writing GUID mapper files to {mappers_dir}...")
        for asset_type, guid_map in guid_maps.items():
            if guid_map:
                output_path = mappers_dir / f"{asset_type}.json"
                # For the output JSON, we still want the project-relative path
                relative_guid_map = {g: Path(p).relative_to(input_dir).as_posix() for g, p in guid_map.items()}
                write_json(relative_guid_map, output_path, indent=indent, shrink=shrink)
                print(f"    -> Created {asset_type}.json with {len(guid_map)} entries.")
        print(f"--> Successfully created all GUID mappers.")
    except OSError as e:
        print(f"Error: Could not create GUID mapper directory or files. {e}", file=sys.stderr)
    
    # Return the map with full paths for the processor
    full_path_guid_map = {}
    for asset_type in guid_maps:
        full_path_guid_map.update(guid_maps[asset_type])
    return full_path_guid_map


def main():
    """
    Main function to run the mid-level data extraction process.
    This script generates a virtual file structure and detailed GUID mappers.
    """
    parser = argparse.ArgumentParser(
        description="Generates a virtual representation of the project's structure and GUID maps."
    )
    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."
    )
    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 Directories ---
    mid_level_output_dir = output_dir / "MidLevel"
    output_assets_dir = mid_level_output_dir / "Assets"
    try:
        output_assets_dir.mkdir(parents=True, exist_ok=True)
        print(f"Output will be saved to: {mid_level_output_dir}")
    except OSError as e:
        print(f"Error: Could not create output directory '{mid_level_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 Mid-Level Extraction ---")

    # --- Task 1: Replicate 'Assets' directory structure ---
    print("\n[1/3] Replicating 'Assets' directory structure...")
    replicate_directory_structure(str(assets_dir), str(output_assets_dir), ignored_folders=ignored_folders, project_root=input_dir)
    print("--> Directory structure replication complete.")

    # --- Task 2: Generate GUID Map and Mappers ---
    print("\n[2/3] Generating GUID Mappers...")
    guid_map = generate_guid_mappers(
        input_dir, 
        mid_level_output_dir, 
        indent=indent_level, 
        shrink=shrink_json, 
        ignored_folders=ignored_folders
    )
    print("--> GUID Mapper generation complete.")

    # --- Task 3: Orchestrate Scene and Prefab Parsing for Hierarchy ---
    print("\n[3/3] Parsing Scene and Prefab Hierarchies...")
    
    print("--> Finding scene and prefab files...")
    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
    
    print(f"--> Found {len(files_to_process)} total scene/prefab 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)
        
        relative_path = file_path.relative_to(assets_dir)
        output_json_path = (output_assets_dir / relative_path).with_suffix('.json')
        
        try:
            print(f"\n--- Processing {file_path.name} ({i+1}/{total_files}) ---")
            
            # Use the sophisticated processor for building the visual tree
            processor = UnitySceneProcessor(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"    -> Pass 1/6: Building relationship maps and creating basic nodes...")
            processor.process_first_pass()

            print(f"    -> Pass 2/6: Building hierarchy relationships...")
            processor.process_second_pass()

            print(f"    -> Pass 3/6: Verifying and fixing parent-child relationships...")
            processor.verification_pass()

            print(f"    -> Pass 4/6: Extracting components...")
            processor.process_third_pass()

            print(f"    -> Pass 5/6: Merging prefab data...")
            processor.merge_prefab_data_pass()
            
            print(f"    -> Pass 6/6: Assembling final hierarchy...")
            hierarchy = processor.get_hierarchy()
            
            output_json_path.parent.mkdir(parents=True, exist_ok=True)
            write_json(hierarchy, output_json_path, indent=indent_level, shrink=shrink_json)
            print(f"--> Successfully processed hierarchy for {file_path.name} -> {output_json_path}")

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

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

if __name__ == "__main__":
    main()