import os
import string
import random
import argparse

def generate_unique_code(existing_codes, length=10):
    """Generate a unique alphanumeric code."""
    chars = string.ascii_letters + string.digits
    while True:
        code = ''.join(random.choice(chars) for _ in range(length))
        if code not in existing_codes:
            return code

def get_existing_codes(directory):
    """Get all existing alphanumeric codes in the directory."""
    existing_codes = set()
    for filename in os.listdir(directory):
        parts = filename.split('_')
        if len(parts) > 1:
            existing_codes.add(parts[-1].split('.')[0])
    return existing_codes

def rename_files_with_unique_code(directory, search_string):
    """Rename all files containing the search string by appending a unique alphanumeric code."""
    existing_codes = get_existing_codes(directory)
    matched_files = [filename for filename in os.listdir(directory) if search_string in filename]

    if not matched_files:
        print(f"No files found containing the string '{search_string}'")
        return

    unique_code = generate_unique_code(existing_codes)

    for filename in matched_files:
        base_filename = filename.split('.')[0]
        new_file_name = f"{base_filename}_{unique_code}.pkl.gz"

        old_file_path = os.path.join(directory, filename)
        new_file_path = os.path.join(directory, new_file_name)

        os.rename(old_file_path, new_file_path)
        print(f"Renamed '{filename}' to '{new_file_name}'")

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Rename files containing a specific string by appending a unique alphanumeric code.")
    parser.add_argument('search_string', type=str, help="The string to search for in filenames.")
    parser.add_argument('--directory', type=str, default='histos', help="The directory containing the files (default: 'histos').")
    args = parser.parse_args()

    # Ensure the directory exists
    if not os.path.exists(args.directory):
        print(f"Directory '{args.directory}' does not exist.")
    else:
        # Rename the files
        rename_files_with_unique_code(args.directory, args.search_string)