import os
import re

def modify_synthetic_imports(file_path):
    """Modify imports of synthetic_dataset_generator to add src."""
    try:
        with open(file_path, 'r') as file:
            content = file.read()
        
        # Replace both import patterns to add src.
        modified_content = re.sub(
            r'from src.synthetic_dataset_generator',
            'from src.synthetic_dataset_generator',
            content
        )
        modified_content = re.sub(
            r'import src.synthetic_dataset_generator',
            'import src.synthetic_dataset_generator',
            modified_content
        )

        # Only write if changes were made
        if modified_content != content:
            with open(file_path, 'w') as file:
                file.write(modified_content)
            print(f"Modified imports in: {file_path}")

    except Exception as e:
        print(f"Error processing {file_path}: {str(e)}")

def process_directory(start_path):
    """Recursively process all Python files in directory"""
    for root, _, files in os.walk(start_path):
        for file in files:
            if file.endswith('.py'):
                file_path = os.path.join(root, file)
                modify_synthetic_imports(file_path)

if __name__ == "__main__":
    import sys
    if len(sys.argv) != 2:
        print("Usage: python script.py <directory_path>")
        sys.exit(1)

    directory_path = sys.argv[1]
    if not os.path.isdir(directory_path):
        print(f"Error: {directory_path} is not a valid directory")
        sys.exit(1)

    process_directory(directory_path)
    print("Processing complete!")