#!/usr/bin/env python3
"""Network stage (run outside Docker): download every map in the frozen list, pinning sha256 on first fetch.
Usage: fetch_maps.py <frozen_csv> <maps_dir> <pins_csv> [--ids id1,id2,...]
- If a pin exists for an image, the downloaded bytes must match it (else status=sha_mismatch).
- Files are stored as <maps_dir>/<image_id>.nii.gz (or .nii); pins_csv rows: image_id,sha256,bytes,status."""
import csv, hashlib, os, subprocess, sys, time
from concurrent.futures import ThreadPoolExecutor
UA = "curl/8.5.0"

def main():
    frozen, maps_dir, pins_path = sys.argv[1:4]
    only = None
    if "--ids" in sys.argv:
        only = {int(x) for x in sys.argv[sys.argv.index("--ids") + 1].split(",")}
    os.makedirs(maps_dir, exist_ok=True)
    rows = [r for r in csv.DictReader(open(frozen)) if only is None or int(r["image_id"]) in only]
    pins = {}
    if os.path.exists(pins_path):
        pins = {r["image_id"]: r for r in csv.DictReader(open(pins_path))}

    def one(r):
        iid = r["image_id"]; ext = ".nii" if r["file_url"].endswith(".nii") else ".nii.gz"
        dest = os.path.join(maps_dir, iid + ext)
        if not os.path.exists(dest):
            ok = False
            for attempt in range(5):
                p = subprocess.run(["curl", "-s", "-f", "-L", "-A", UA, "--max-time", "600", "-o", dest + ".part", r["file_url"]])
                if p.returncode == 0:
                    os.replace(dest + ".part", dest); ok = True; break
                if p.returncode == 22:  # HTTP error (404/410 etc.): record, do not hammer
                    if attempt >= 1: break
                time.sleep(5 * 2 ** attempt)
            if not ok:
                if os.path.exists(dest + ".part"): os.remove(dest + ".part")
                return {"image_id": iid, "sha256": "", "bytes": "", "status": "download_failed"}
        b = open(dest, "rb").read(); sha = hashlib.sha256(b).hexdigest()
        pin = pins.get(iid)
        if pin and pin["sha256"] and pin["sha256"] != sha:
            return {"image_id": iid, "sha256": pin["sha256"], "bytes": pin["bytes"], "status": "sha_mismatch"}
        return {"image_id": iid, "sha256": sha, "bytes": str(len(b)), "status": "ok"}

    with ThreadPoolExecutor(4) as ex:
        res = list(ex.map(one, rows))
    for r in res:
        if r["image_id"] not in pins or pins[r["image_id"]]["status"] != "ok":
            pins[r["image_id"]] = r
    with open(pins_path, "w", newline="") as fh:
        w = csv.DictWriter(fh, fieldnames=["image_id", "sha256", "bytes", "status"], lineterminator="\n"); w.writeheader()
        for k in sorted(pins, key=int): w.writerow(pins[k])
    from collections import Counter
    print(Counter(r["status"] for r in res))

if __name__ == "__main__":
    main()