#!/usr/bin/env python3
"""Build the text-to-image retrieval set for the embedding-model benchmark.

Samples N images from MS-COCO val2017 (seeded, so the sample is reproducible),
takes the first human caption of each image as its query, and downloads the
images. Output: <out>/images/<id>.jpg and <out>/queries.json with
[{"image_id", "file", "query", "flickr_url", "license"}].

COCO annotations are CC BY 4.0; images carry the Flickr licence recorded per
item in queries.json and are not redistributed by this repository.

Usage:
  python build_dataset.py --annotations captions_val2017.json --out data --n 1000
"""
import argparse
import json
import os
import random
import time
from concurrent.futures import ThreadPoolExecutor

import requests

IMG_BASE = "http://images.cocodataset.org/val2017/"


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--annotations", required=True)
    ap.add_argument("--out", default="data")
    ap.add_argument("--n", type=int, default=1000)
    ap.add_argument("--seed", type=int, default=20260910)
    a = ap.parse_args()

    coco = json.load(open(a.annotations))
    first_caption: dict[int, str] = {}
    for ann in sorted(coco["annotations"], key=lambda x: x["id"]):
        first_caption.setdefault(ann["image_id"], ann["caption"].strip())
    images = sorted(coco["images"], key=lambda x: x["id"])
    random.Random(a.seed).shuffle(images)
    chosen = images[: a.n]

    img_dir = os.path.join(a.out, "images")
    os.makedirs(img_dir, exist_ok=True)

    def fetch(im: dict) -> dict:
        path = os.path.join(img_dir, im["file_name"])
        if not os.path.exists(path):
            for attempt in range(3):
                r = requests.get(IMG_BASE + im["file_name"], timeout=60)
                if r.ok:
                    open(path, "wb").write(r.content)
                    break
                time.sleep(2 * (attempt + 1))
            else:
                raise RuntimeError(f"download failed: {im['file_name']}")
        return {
            "image_id": im["id"],
            "file": im["file_name"],
            "query": first_caption[im["id"]],
            "flickr_url": im.get("flickr_url"),
            "license": im.get("license"),
        }

    with ThreadPoolExecutor(8) as ex:
        rows = list(ex.map(fetch, chosen))
    json.dump(
        {
            "source": "MS-COCO val2017 captions (CC BY 4.0 annotations)",
            "seed": a.seed,
            "n": len(rows),
            "items": rows,
        },
        open(os.path.join(a.out, "queries.json"), "w"),
        indent=1,
    )
    print(f"{len(rows)} images -> {a.out}")


if __name__ == "__main__":
    main()