import random
import shutil
from pathlib import Path
from typing import List


def list_numeric_class_dirs(root: Path) -> List[Path]:
    """Return class directories whose names are numeric (e.g., 0..999), sorted by int value."""
    class_dirs = [p for p in root.iterdir() if p.is_dir() and p.name.isdigit()]
    class_dirs.sort(key=lambda p: int(p.name))
    return class_dirs


def list_image_files(directory: Path) -> List[Path]:
    """List image files in a directory. Accept common jpeg/jpg extensions."""
    valid_exts = {".jpeg", ".jpg", ".JPEG", ".JPG"}
    return [p for p in directory.iterdir() if p.is_file() and p.suffix in valid_exts]


def copy_random_one_per_class(source_root: Path, dest_dir: Path, seed: int = 42) -> int:
    """Copy one random image from each numeric-named subdirectory to dest_dir.

    Filenames follow convention:
    IN-1k_v2_{original_name_without_extension}_class{parent_dir}.jpeg
    Returns count of images copied.
    """
    if seed is not None:
        random.seed(seed)

    dest_dir.mkdir(parents=True, exist_ok=True)

    class_dirs = list_numeric_class_dirs(source_root)
    if not class_dirs:
        raise FileNotFoundError(f"No numeric class directories found under: {source_root}")

    copied = 0
    for class_dir in class_dirs:
        images = list_image_files(class_dir)
        if not images:
            # Skip empty classes but report
            print(f"Warning: No image files found in {class_dir}")
            continue
        chosen = random.choice(images)

        new_name = f"IN-1k_v2_{chosen.stem}_class{class_dir.name}.jpeg"
        dest_path = dest_dir / new_name

        shutil.copy2(chosen, dest_path)
        copied += 1

    return copied


def main() -> None:
    repo_root = Path(__file__).resolve().parent.parent
    in1k_root = repo_root / "IN-1k_v2"
    source_root = in1k_root / "imagenetv2-top-images-format-val"
    dest_dir = in1k_root / "Dataset"

    print(f"Source root: {source_root}")
    print(f"Destination: {dest_dir}")

    if not source_root.exists():
        raise FileNotFoundError(f"Expected source directory not found: {source_root}")

    count = copy_random_one_per_class(source_root=source_root, dest_dir=dest_dir)
    print(f"Copied {count} images to {dest_dir}")


if __name__ == "__main__":
    main()


