"""Fetch specific figure images from ArxivCap on HuggingFace via DuckDB range requests."""
import os
import re
from collections import defaultdict
from pathlib import Path

HF_SRC = "hf://datasets/MMInstruction/ArxivCap/data"


def yymm_of(arxiv_id):
    s = str(arxiv_id)
    if "/" in s:
        return re.sub(r"\D", "", s.split("/", 1)[1])[:4]
    return s[:4]


def safe_name(arxiv_id, fig_index):
    return re.sub(r"[^0-9a-zA-Z.]+", "_", str(arxiv_id)) + f"_f{fig_index}.jpg"


def fetch_images(needed, image_dir, hf_src=HF_SRC):
    """needed: set of (arxiv_id, fig_index 1-based). Skips files already on disk."""
    import duckdb
    image_dir = Path(image_dir)
    image_dir.mkdir(parents=True, exist_ok=True)
    todo = {(a, f) for (a, f) in needed
            if not (image_dir / safe_name(a, f)).exists()}
    if not todo:
        print("  all images already cached")
        return
    by_chunk = defaultdict(set)
    for a, f in todo:
        by_chunk[yymm_of(a)].add((a, f))
    print(f"  fetching {len(todo)} figures from {len(by_chunk)} ArxivCap chunks")

    con = duckdb.connect()
    for setting in ["SET http_retries = 8", "SET http_retry_wait_ms = 2000",
                    "SET threads = 2"]:
        try:
            con.execute(setting)
        except Exception:
            pass
    token = os.environ.get("HF_TOKEN", "")
    if token:
        con.execute(f"CREATE SECRET hf (TYPE HUGGINGFACE, TOKEN '{token}')")
    else:
        print("  WARNING: no HF_TOKEN; anonymous HF requests rate-limit fast")

    for yymm, pairs in sorted(by_chunk.items()):
        ids = sorted({a for a, _ in pairs})
        idlist = ",".join("'" + i.replace("'", "''") + "'" for i in ids)
        src = f"{hf_src}/arXiv_src_{yymm}_*.parquet"
        try:
            tab = con.execute(
                f"SELECT arxiv_id, caption_images FROM read_parquet('{src}') "
                f"WHERE arxiv_id IN ({idlist})").fetch_arrow_table()
        except Exception as e:
            print(f"  chunk {yymm}: fetch failed ({str(e)[:120]}), skipping")
            continue
        got = 0
        for row in tab.to_pylist():
            a = row["arxiv_id"]
            caps = row["caption_images"] or []
            for (aa, f) in pairs:
                if aa != a or f > len(caps):
                    continue
                entry = caps[f - 1] or {}
                cil = entry.get("cil_pairs") or []
                img = (cil[0] or {}).get("image") if cil else None
                data = img.get("bytes") if isinstance(img, dict) else img
                if not data:
                    continue
                (image_dir / safe_name(a, f)).write_bytes(data)
                got += 1
        print(f"  chunk {yymm}: {got} figures saved")


if __name__ == "__main__":
    import argparse
    import csv
    ap = argparse.ArgumentParser()
    ap.add_argument("ids_csv", help="csv with columns arxiv_id,fig_index")
    ap.add_argument("--image-dir", default="vision_images")
    args = ap.parse_args()
    needed = {(r["arxiv_id"], int(r["fig_index"]))
              for r in csv.DictReader(open(args.ids_csv))}
    fetch_images(needed, Path(args.image_dir))