"""On-demand figure fetch from ArxivCap on HuggingFace, with a local cache.

Each figure is fetched at most once: the JPEG and a sidecar json holding the
paper title and full caption are written to the cache dir, keyed by
(arxiv_id, fig_index). Misses are grouped by ArxivCap month chunk so one
query's candidates cost one DuckDB range read per chunk.
"""
import json
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_stem(arxiv_id, fig_index):
    return re.sub(r"[^0-9a-zA-Z.]+", "_", str(arxiv_id)) + f"_f{fig_index}"


def _cached(cache, a, f):
    stem = safe_stem(a, f)
    img, meta = cache / f"{stem}.jpg", cache / f"{stem}.json"
    if img.exists() and meta.exists():
        d = json.loads(meta.read_text())
        d["image_path"] = str(img)
        return d
    return None


def fetch_figures(needed, cache_dir="figure_cache", log=None):
    """needed: iterable of (arxiv_id, fig_index 1-based).
    Returns {(a, f): {"image_path", "title", "caption"}} for what was found."""
    cache = Path(cache_dir)
    cache.mkdir(parents=True, exist_ok=True)
    out, todo = {}, set()
    for a, f in needed:
        hit = _cached(cache, a, f)
        if hit:
            out[(a, f)] = hit
        else:
            todo.add((a, f))
    if not todo:
        return out

    import duckdb
    con = duckdb.connect()
    for s in ["SET http_retries = 6", "SET http_retry_wait_ms = 1500",
              "SET threads = 2"]:
        try:
            con.execute(s)
        except Exception:
            pass
    token = os.environ.get("HF_TOKEN", "")
    if token:
        con.execute(f"CREATE SECRET hf (TYPE HUGGINGFACE, TOKEN '{token}')")

    by_chunk = defaultdict(set)
    for a, f in todo:
        by_chunk[yymm_of(a)].add((a, f))
    for yymm, pairs in sorted(by_chunk.items()):
        if log:
            log(f"  fetching {len(pairs)} figure(s) from chunk {yymm}...")
        ids = sorted({a for a, _ in pairs})
        idlist = ",".join("'" + i.replace("'", "''") + "'" for i in ids)
        src = f"{HF_SRC}/arXiv_src_{yymm}_*.parquet"
        try:
            rows = con.execute(
                f"SELECT arxiv_id, title, caption_images "
                f"FROM read_parquet('{src}') "
                f"WHERE arxiv_id IN ({idlist})").fetch_arrow_table().to_pylist()
        except Exception as e:
            if log:
                log(f"  chunk {yymm} failed: {str(e)[:80]}")
            continue
        for row in rows:
            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
                stem = safe_stem(a, f)
                (cache / f"{stem}.jpg").write_bytes(data)
                d = {"title": row.get("title") or "",
                     "caption": entry.get("caption") or ""}
                (cache / f"{stem}.json").write_text(json.dumps(d))
                d["image_path"] = str(cache / f"{stem}.jpg")
                out[(a, f)] = d
    return out