"""Download a versioned scPerturb count matrix and verify its source checksum."""

from pathlib import Path
import argparse, hashlib, os
import requests

FILES = {
    "norman": ("NormanWeissman2019_filtered.h5ad", "c870e6967d91c017d9da827bab183cd6"),
    "replogle_k562": (
        "ReplogleWeissman2022_K562_essential.h5ad",
        "d8cba17576d1a8afc0f7d71b79cad0f7",
    ),
}


def md5(path):
    h = hashlib.md5()
    with open(path, "rb") as f:
        for block in iter(lambda: f.read(8 * 1024**2), b""):
            h.update(block)
    return h.hexdigest()


def download(dataset, output):
    name, digest = FILES[dataset]
    out = Path(output)
    out.mkdir(parents=True, exist_ok=True)
    dest = out / name
    if dest.exists():
        if md5(dest) != digest:
            raise ValueError(f"Checksum mismatch: {dest}")
        return dest
    part = dest.with_suffix(".h5ad.partial")
    offset = part.stat().st_size if part.exists() else 0
    url = f"https://zenodo.org/records/10044268/files/{name}?download=1"
    with requests.get(
        url,
        stream=True,
        headers={"Range": f"bytes={offset}-"} if offset else {},
        timeout=(30, 120),
    ) as response:
        response.raise_for_status()
        append = offset > 0 and response.status_code == 206
        with open(part, "ab" if append else "wb") as f:
            for chunk in response.iter_content(8 * 1024**2):
                f.write(chunk)
    if md5(part) != digest:
        raise ValueError(f"Source checksum mismatch: {part}")
    os.replace(part, dest)
    return dest


if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("dataset", choices=FILES)
    p.add_argument("--output", default="data/raw")
    a = p.parse_args()
    print(download(a.dataset, a.output))