import os
from datasets import load_dataset
from PIL import Image


OUTPUT_DIR = "data/images"
NUM_IMAGES = 1000


def _valid_image_count_and_next_index():

    valid = 0
    next_index = 0

    if not os.path.isdir(OUTPUT_DIR):
        return valid, next_index

    for name in sorted(os.listdir(OUTPUT_DIR)):
        base, ext = os.path.splitext(name)
        if ext.lower() not in {".jpg", ".jpeg"}:
            continue
        if not base.isdigit():
            continue

        idx = int(base)
        path = os.path.join(OUTPUT_DIR, name)
        try:
            with Image.open(path) as img:
                img.verify()
            valid += 1
            next_index = max(next_index, idx + 1)
        except Exception:
            # Ignore broken files; they will not be counted as valid samples.
            continue

    return valid, next_index


def download_images():

    os.makedirs(OUTPUT_DIR, exist_ok=True)

    count, next_index = _valid_image_count_and_next_index()

    if count >= NUM_IMAGES:
        print(f"already have {count} valid images in {OUTPUT_DIR}")
        return

    dataset = load_dataset(
        "obvtiger/unsplash-img",
        split="train",
        streaming=True
    )

    for item in dataset:

        try:
            img = item.get("image")
        except Exception:
            continue

        if not isinstance(img, Image.Image):
            continue

        try:
            img = img.convert("RGB")
            img.thumbnail((512, 512))
        except Exception:
            continue

        path = os.path.join(OUTPUT_DIR, f"{next_index:04d}.jpg")

        try:
            img.save(path)
        except Exception:
            continue

        count += 1
        next_index += 1

        print(f"saved {path} ({count}/{NUM_IMAGES})")

        if count >= NUM_IMAGES:
            break


if __name__ == "__main__":
    download_images()