#!/usr/bin/env python3
"""Resumable NeuroVault metadata crawler (metadata only; no map files).
Stage A: page all collections. Stage B: page images of every collection that has a publication DOI.
Each API page is saved as raw JSON (gzipped); existing files are skipped on restart."""
import json, os, subprocess, sys, time, gzip, glob
from concurrent.futures import ThreadPoolExecutor
BASE = "https://neurovault.org/api"
UA = "curl/8.5.0"
OUT = os.path.dirname(os.path.abspath(__file__))
LIM = 1000

def fetch(url, dest):
    if os.path.exists(dest):
        return json.load(gzip.open(dest))
    for attempt in range(8):
        try:
            r = subprocess.run(["curl", "-s", "-f", "-A", UA, "--max-time", "300", url],
                               capture_output=True, timeout=320)
            if r.returncode == 0:
                d = json.loads(r.stdout)
                tmp = dest + ".tmp"
                with gzip.open(tmp, "wb") as f: f.write(r.stdout)
                os.replace(tmp, dest)
                return d
            err = f"curl rc={r.returncode}"
        except Exception as e:
            err = repr(e)
        wait = min(300, 10 * 2 ** attempt)
        print(f"retry {attempt} {url} {err}; sleep {wait}", flush=True)
        time.sleep(wait)
    raise RuntimeError("failed " + url)

def stage_a():
    os.makedirs(f"{OUT}/collections", exist_ok=True)
    d = fetch(f"{BASE}/collections/?format=json&limit={LIM}&offset=0", f"{OUT}/collections/p000000.json.gz")
    n = d["count"]
    offs = list(range(LIM, n, LIM))
    with ThreadPoolExecutor(3) as ex:
        list(ex.map(lambda o: fetch(f"{BASE}/collections/?format=json&limit={LIM}&offset={o}",
                                    f"{OUT}/collections/p{o:06d}.json.gz"), offs))
    cols = []
    for f in sorted(glob.glob(f"{OUT}/collections/p*.json.gz")):
        cols += json.load(gzip.open(f))["results"]
    print("collections", len(cols), "of", n, flush=True)
    return cols

def crawl_collection(cid):
    d = fetch(f"{BASE}/collections/{cid}/images/?format=json&limit={LIM}&offset=0",
              f"{OUT}/images/c{cid}_p000000.json.gz")
    for o in range(LIM, d["count"], LIM):
        fetch(f"{BASE}/collections/{cid}/images/?format=json&limit={LIM}&offset={o}",
              f"{OUT}/images/c{cid}_p{o:06d}.json.gz")
    return cid

def main():
    cols = stage_a()
    os.makedirs(f"{OUT}/images", exist_ok=True)
    doi = [c for c in cols if (c.get("DOI") or "").strip() and (c.get("number_of_images") or 0) > 0]
    doi.sort(key=lambda c: c["id"])
    print("DOI collections with images", len(doi), "images", sum(c["number_of_images"] for c in doi), flush=True)
    done = 0
    with ThreadPoolExecutor(4) as ex:
        for cid in ex.map(crawl_collection, [c["id"] for c in doi]):
            done += 1
            if done % 50 == 0: print(time.strftime("%H:%M:%S"), "done", done, "/", len(doi), flush=True)
    open(f"{OUT}/CRAWL_DONE", "w").write(time.strftime("%Y-%m-%dT%H:%M:%S%z") + "\n")
    print("CRAWL_DONE", flush=True)

if __name__ == "__main__":
    main()