"""
Download a stratified sample of Open Images V7.

Open Images is published by Google under CC BY 2.0 (commercial use OK).
Each downloaded image is recorded in the manifest with its source URL and
license string for audit trail.

Usage
-----
    python scripts/dataset/fetch_open_images.py \
        --out data/raw/real \
        --count 50000

Notes
-----
We use the FiftyOne library because it has first-class Open Images support
and handles the (large, non-trivial) metadata files for us. FiftyOne is
Apache 2.0 licensed.

This script is idempotent: re-running with the same --out resumes where it
left off and skips already-downloaded files.
"""
from __future__ import annotations

import argparse
import csv
import hashlib
import time
from pathlib import Path


def _load_with_retry(foz, split: str, count: int, max_attempts: int = 20):
    """Call foz.load_zoo_dataset with retry on transient network errors.

    Open Images is hosted on S3; a single dropped TCP connection bubbles up
    as botocore.EndpointConnectionError and kills the whole download even
    when 99% of the work is done. Each successful image is cached locally,
    so retrying the call resumes from where it left off — we just need to
    keep retrying until the network cooperates for one full pass.

    Caught broadly to also handle ConnectionError, OSError, and the
    botocore wrappers without a hard import on botocore.
    """
    for attempt in range(1, max_attempts + 1):
        try:
            return foz.load_zoo_dataset(
                "open-images-v7",
                split=split,
                max_samples=count,
                shuffle=True,
                # We only need image data, not labels — bbox/segmentation/etc.
                # add tens of GB of metadata we don't use.
                label_types=[],
            )
        except Exception as exc:  # noqa: BLE001 — intentional broad catch
            cls = type(exc).__name__
            transient = any(
                token in cls
                for token in (
                    "Connection", "Endpoint", "Timeout", "SSL", "Proxy",
                )
            ) or isinstance(exc, OSError)
            if not transient or attempt >= max_attempts:
                raise
            wait = min(60, 5 * attempt)
            print(
                f"\n[retry {attempt}/{max_attempts}] {cls}: {exc}\n"
                f"  Sleeping {wait}s before retrying — already-downloaded "
                "images are cached, so the next attempt resumes."
            )
            time.sleep(wait)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--out", type=Path, required=True, help="Output directory")
    parser.add_argument("--count", type=int, default=50_000, help="Number of images")
    parser.add_argument(
        "--split",
        choices=["train", "validation", "test"],
        default="train",
        help="Open Images split to draw from",
    )
    parser.add_argument(
        "--max-attempts",
        type=int,
        default=20,
        help="Max retry attempts on network errors before giving up",
    )
    args = parser.parse_args()

    args.out.mkdir(parents=True, exist_ok=True)
    manifest_path = args.out.parent / "real_manifest.csv"

    print(f"Downloading {args.count} images from Open Images {args.split} split...")
    print("Importing FiftyOne (this is slow the first time)...")

    # Lazy import — fiftyone has heavy native deps.
    import fiftyone.zoo as foz

    dataset = _load_with_retry(foz, args.split, args.count, args.max_attempts)

    rows: list[dict] = []
    for sample in dataset:
        src = Path(sample.filepath)
        dst = args.out / src.name
        if not dst.exists():
            dst.write_bytes(src.read_bytes())

        with dst.open("rb") as fh:
            sha = hashlib.sha256(fh.read()).hexdigest()

        rows.append({
            # as_posix() so the manifest is portable to Linux GPU boxes —
            # str() on Windows produces backslashes that break Path parsing
            # on POSIX.
            "path": dst.relative_to(args.out.parent.parent).as_posix(),
            "class": "authentic",
            "source": "open_images_v7",
            "license": "CC-BY-2.0",
            "license_url": "https://creativecommons.org/licenses/by/2.0/",
            "sha256": sha,
        })

    # Write manifest fragment for the real half.
    with manifest_path.open("w", newline="") as fh:
        writer = csv.DictWriter(
            fh,
            fieldnames=["path", "class", "source", "license", "license_url", "sha256"],
        )
        writer.writeheader()
        writer.writerows(rows)

    print(f"Done. {len(rows)} images. Manifest fragment: {manifest_path}")


if __name__ == "__main__":
    main()