#!/usr/bin/env python3.13
"""Download a Stokastisk release bundle from Hugging Face (CDN-backed `resolve/` URLs) - standard library only.

    python3.13 fetch_release.py stokastisk-0.1.0-kaggle-t4-sm75 [DEST]      # prints the bundle directory

 * ONE file (the release's `<release>-kaggle-input.zip`, ~37 MB) as a single stream by default (workers=1; `workers=N` opts into N parallel range requests, which stalled against the Hub CDN in testing); falls back
   to a single stream when the server does not support ranges, and to `file://` for offline tests;
 * SHA256 verified against the `.sha256` sidecar BEFORE extraction; path-traversal-safe extraction; the extracted directory must carry
   the release name and `manifests/RELEASE-MANIFEST.json` must say the same name;
 * nothing is installed here: run `<bundle>/scripts/kaggle/02_install.py <bundle>` afterwards (it verifies every file in SHA256SUMS first).
Default destination is /tmp (NOT /kaggle/working, which has a 20 GiB quota)."""
import hashlib, json, os, pathlib, shutil, sys, time, urllib.request, zipfile
from concurrent.futures import ThreadPoolExecutor

HF_REPO = "waqasm86/stokastisk"
HF_BASE = f"https://huggingface.co/{HF_REPO}/resolve"


def _open(url, headers=None):
    return urllib.request.urlopen(urllib.request.Request(url, headers=headers or {"User-Agent": "stokastisk-fetch/1"}), timeout=60)


def _range_ok(url):
    if not url.startswith("http"):
        return None
    with _open(url, {"Range": "bytes=0-0", "User-Agent": "stokastisk-fetch/1"}) as r:
        cr = r.headers.get("Content-Range")                      # "bytes 0-0/37080766"
        return int(cr.rsplit("/", 1)[1]) if r.status == 206 and cr else None


def _download(url, dest, workers=1, retries=3):
    """Download to `<dest>.part`, check free disk space and the advertised length, then replace `dest` atomically; the partial file is removed on any failure."""
    dest = pathlib.Path(dest); part = dest.with_name(dest.name + ".part")
    part.unlink(missing_ok=True)
    try:
        total = _range_ok(url)
        if total and shutil.disk_usage(dest.parent).free < total + (64 << 20):
            raise RuntimeError(f"not enough free disk space in {dest.parent} for {total / 1e6:.0f} MB")
        _download_to(url, part, workers, retries)
        if total and part.stat().st_size != total:
            raise RuntimeError(f"incomplete download: {part.stat().st_size} of {total} bytes")
        os.replace(part, dest)
    except BaseException:
        part.unlink(missing_ok=True)
        raise


def _download_to(url, dest, workers=1, retries=3):
    dest = pathlib.Path(dest); total = _range_ok(url)
    t0 = time.time()
    if total and total > (4 << 20):
        chunk = -(-total // workers); parts = [(i * chunk, min(total, (i + 1) * chunk) - 1) for i in range(workers) if i * chunk < total]
        with open(dest, "wb") as f:
            f.truncate(total)

        def part(rng):
            a, b = rng
            for attempt in range(retries):
                try:
                    with _open(url, {"Range": f"bytes={a}-{b}", "User-Agent": "stokastisk-fetch/1"}) as r, open(dest, "r+b") as f:
                        f.seek(a); shutil.copyfileobj(r, f, 1 << 20)
                    return
                except Exception:
                    if attempt == retries - 1: raise
                    time.sleep(2 ** attempt)
        with ThreadPoolExecutor(len(parts)) as ex:
            list(ex.map(part, parts))
        mode = f"{len(parts)} parallel ranges"
    else:
        for attempt in range(retries):
            try:
                with _open(url) as r, open(dest, "wb") as f:
                    shutil.copyfileobj(r, f, 1 << 20)
                break
            except Exception:
                if attempt == retries - 1: raise
                time.sleep(2 ** attempt)
        mode = "single stream"
    n = dest.stat().st_size
    print(f"downloaded {n / 1e6:.1f} MB in {time.time() - t0:.1f}s ({n / 1e6 / max(time.time() - t0, 1e-6):.1f} MB/s, {mode})")


def _sha256(path):
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for c in iter(lambda: f.read(1 << 20), b""):
            h.update(c)
    return h.hexdigest()


def fetch_release(release, dest="/tmp/stokastisk-release", base=None, revision="main", workers=1):
    """Return the verified, extracted bundle directory (…/<release>). `base` overrides the Hub URL prefix (tests use file://)."""
    base = (base or HF_BASE).rstrip("/")
    root = f"{base}/{revision}/releases/{release}" if base.startswith("http") else f"{base}/releases/{release}"
    dest = pathlib.Path(dest); dest.mkdir(parents=True, exist_ok=True)
    zname = f"{release}-kaggle-input.zip"; zpath = dest / zname
    if zpath.exists(): zpath.unlink()
    with _open(f"{root}/{zname}.sha256") as r:
        expected = r.read().decode().split()[0]
    _download(f"{root}/{zname}", zpath, workers=workers)
    got = _sha256(zpath)
    if got != expected:
        zpath.unlink(); raise RuntimeError(f"SHA256 mismatch for {zname}: expected {expected}, got {got}")
    out = dest / release
    if out.exists(): shutil.rmtree(out)
    with zipfile.ZipFile(zpath) as z:
        assert z.testzip() is None
        for n in z.namelist():
            if n.startswith("/") or ".." in pathlib.PurePosixPath(n).parts or not n.startswith(release + "/"):
                raise RuntimeError(f"unsafe or foreign path in archive: {n}")
        z.extractall(dest)
    m = json.loads((out / "manifests" / "RELEASE-MANIFEST.json").read_text())
    if m.get("name") != release:
        raise RuntimeError(f"manifest name {m.get('name')!r} != {release!r}")
    print(f"release {release}: SHA256 {got[:16]}... verified, extracted to {out}")
    return out


if __name__ == "__main__":
    if len(sys.argv) < 2:
        raise SystemExit(__doc__)
    print(fetch_release(sys.argv[1], *(sys.argv[2:3] or ["/tmp/stokastisk-release"])))