import os
from PIL import Image


def validate_images(directory, recursive=True, extensions=None):
    if extensions is None:
        extensions = ('.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.gif')  # Add more extensions if needed

    error_log = []

    for root, dirs, files in os.walk(directory):
        if not recursive and root != directory:
            break

        for file in files:
            file_path = os.path.join(root, file)

            if file_path.endswith(extensions):
                try:
                    img = Image.open(file_path)
                    img.verify()  # Optionally verify that the image is intact (slower)
                    img.close()
                    print(f"Validated: {file_path}")
                except Exception as e:
                    error_log.append((file_path, str(e)))
                    print(f"Error loading image: {file_path} - {str(e)}")

    return error_log


if __name__ == "__main__":
    directory_to_validate = "../dataset/train2014"
    recursive_search = True  # Set to False to only check the specified directory (no subdirectories)

    error_log = validate_images(directory_to_validate, recursive=recursive_search)
    print("\nValidation complete.")

    if error_log:
        print("\nErrors encountered:")
        for file_path, error_msg in error_log:
            print(f"{file_path}: {error_msg}")
    else:
        print("No errors found.")