import argparse
import os
from pathlib import Path
from typing import Union


def validate_dataset(data_path: Union[str, Path]) -> bool:
    """
    Validate that each image in the dataset directory has a corresponding text file and vice versa.

    Args:
        data_path: Path to the dataset directory.

    Returns:
        True if the dataset is valid, False otherwise.
    """
    data_dir = Path(data_path)
    image_extensions = {'.png', '.jpg', '.jpeg', '.webp'}

    images = []
    texts = []

    for file in data_dir.iterdir():
        if file.suffix.lower() in image_extensions:
            images.append(file.stem)
        elif file.suffix == '.txt':
            texts.append(file.stem)

    images.sort()
    texts.sort()

    print(f"Found {len(images)} images and {len(texts)} text files")

    missing_texts = set(images) - set(texts)
    missing_images = set(texts) - set(images)

    if missing_texts:
        print(f"Missing text files for: {missing_texts}")
    if missing_images:
        print(f"Missing images for: {missing_images}")

    if not missing_texts and not missing_images:
        print("Dataset structure is valid!")

    return len(images) == len(texts) and not missing_texts and not missing_images

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Validate image-text dataset structure.")
    parser.add_argument(
        "--path",
        type=str,
        required=True,
        help="Path to the dataset directory to validate."
    )
    args = parser.parse_args()
    validate_dataset(args.path)