"""Fetch open-source document images with dense, varied layouts for the model card gallery.

Pulls sample pages from the CDLA-Permissive-1.0-licensed `creative-graphic-design/PubLayNet`
dataset on the Hugging Face Hub (a re-hosting of PubLayNet: Zhong et al., 2019,
https://arxiv.org/abs/1908.07836 — scientific articles from PubMed Central Open Access).
Scores a scan window of rows by their COCO-annotation count (a proxy for layout density —
how many text/title/list/table/figure regions are on the page) and downloads the densest ones,
so the gallery shows the model handling as many layout elements as possible per image.

Uses only the HF `datasets-server` `/rows` REST API (plain paginated JSON + signed image URLs)
— no `datasets`/torch/pyarrow dependency needed.

Usage:
    python fetch_example_images.py --count 5 --scan 3000 --out-dir examples/inputs
"""

from __future__ import annotations

import argparse
import json
import urllib.request
from pathlib import Path

DATASET = "creative-graphic-design/PubLayNet"
ROWS_URL = "https://datasets-server.huggingface.co/rows"
PAGE = 100


def scan(scan_rows: int, split: str) -> list[dict]:
    candidates = []
    for offset in range(0, scan_rows, PAGE):
        url = (
            f"{ROWS_URL}?dataset={DATASET.replace('/', '%2F')}&config=default"
            f"&split={split}&offset={offset}&length={PAGE}"
        )
        try:
            with urllib.request.urlopen(url, timeout=30) as r:
                payload = json.load(r)
        except Exception as exc:
            print(f"  offset {offset}: skipped ({exc})")
            continue
        for item in payload["rows"]:
            row = item["row"]
            candidates.append(
                {
                    "file_name": row["file_name"],
                    "width": row["width"],
                    "height": row["height"],
                    "n_boxes": len(row["annotations"]["bbox"]),
                    "n_classes": len(set(row["annotations"]["category_id"])),
                    "src": row["image"]["src"],
                }
            )
        print(f"  offset {offset}: scanned ({len(candidates)} candidates so far)")
    return candidates


def main() -> int:
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--count", type=int, default=5, help="how many images to download")
    p.add_argument("--scan", type=int, default=2000, help="how many dataset rows to scan for density")
    p.add_argument("--split", default="train")
    p.add_argument("--out-dir", type=Path, default=Path("examples/inputs"))
    args = p.parse_args()

    args.out_dir.mkdir(parents=True, exist_ok=True)
    print(f"Scanning {args.scan} rows of {DATASET} ({args.split}) for the densest layouts...")
    candidates = scan(args.scan, args.split)
    if not candidates:
        raise SystemExit("no candidates found — dataset-server may be unreachable")

    candidates.sort(key=lambda c: (-c["n_boxes"], -c["n_classes"]))
    chosen = candidates[: args.count]

    manifest = []
    print("\nDownloading:")
    for c in chosen:
        dest = args.out_dir / c["file_name"]
        urllib.request.urlretrieve(c["src"], dest)
        print(f"  {c['n_boxes']:3d} boxes, {c['n_classes']} classes -> {dest.name} ({c['width']}x{c['height']})")
        manifest.append({k: v for k, v in c.items() if k != "src"})

    (args.out_dir / "SOURCE.json").write_text(
        json.dumps(
            {
                "dataset": DATASET,
                "dataset_url": f"https://huggingface.co/datasets/{DATASET}",
                "license": "CDLA-Permissive-1.0",
                "citation": (
                    "Zhong, X., Tang, J., & Yepes, A. J. (2019). "
                    "PubLayNet: largest dataset ever for document layout analysis. "
                    "arXiv:1908.07836"
                ),
                "note": "Pages sourced from PubMed Central open-access scientific articles.",
                "images": manifest,
            },
            indent=2,
        )
    )
    print(f"\n{len(chosen)} image(s) -> {args.out_dir}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())