terra-llm/Assets/LLM/source/extract_mid_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 replicate_directory_structure, find_files_by_extension
from virtual_tree_builder import build_virtual_tree

def copy_scripts(assets_dir, output_assets_dir):
    """
    Copies all C# scripts (.cs) to the target directory.
    """
    print("\n--- Starting Task 3: 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
        
        # Ensure the destination directory exists
        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 mid-level data extraction process.
    """
    parser = argparse.ArgumentParser(
        description="Generates a virtual representation of the project's structure."
    )
    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()

    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)

    # --- Task 1: Create output folder and replicate structure ---
    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 assets_dir.is_dir():
        print(f"\n--- Replicating 'Assets' directory structure ---")
        replicate_directory_structure(str(assets_dir), str(output_assets_dir))
        print("Directory structure replication complete.")
    else:
        print(f"Warning: 'Assets' directory not found in '{input_dir}'. Skipping structure replication.", file=sys.stderr)
        return # Exit if there's no Assets directory

    # --- Task 3: Script Handling ---
    copy_scripts(assets_dir, output_assets_dir)

    # --- Task 4: Orchestration ---
    print("\n--- Starting Task 4: Orchestration ---")
    
    # Find all scene and prefab files
    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"Found {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: {file_path.name}")
        
        # Generate the virtual tree
        virtual_tree = build_virtual_tree(str(file_path))
        
        if virtual_tree:
            # Construct the output path
            relative_path = file_path.relative_to(assets_dir)
            output_path = output_assets_dir / relative_path
            output_path = output_path.with_suffix('.json')
            
            # Create parent directories if they don't exist
            output_path.parent.mkdir(parents=True, exist_ok=True)
            
            # Save the JSON output
            try:
                with open(output_path, 'w', encoding='utf-8') as f:
                    json.dump(virtual_tree, f, indent=4)
                print(f"Successfully created: {output_path}")
            except IOError as e:
                print(f"Error writing to {output_path}: {e}", file=sys.stderr)
        else:
            print(f"Skipped {file_path.name} as virtual tree generation failed.")

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

if __name__ == "__main__":
    main()