"""Download released source files with retries and streamed checksum verification."""
import argparse
import hashlib
from pathlib import Path
import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry

p=argparse.ArgumentParser()
p.add_argument('dataset',choices=['tahoe','sciplex'])
p.add_argument('--output',required=True)
p.add_argument('--shards',type=int,default=2)
p.add_argument('--indices',type=int,nargs='+',help='Explicit official Tahoe shard indices')
a=p.parse_args()
out=Path(a.output);out.mkdir(parents=True,exist_ok=True)
s=requests.Session()
s.mount('https://',HTTPAdapter(max_retries=Retry(total=5,backoff_factor=2,status_forcelist=[429,500,502,503,504],respect_retry_after_header=True)))
base='https://huggingface.co/datasets/tahoebio/Tahoe-100M/resolve/main/'
files=[(base+'metadata/'+n+'.parquet',n+'.parquet',None) for n in ['sample_metadata','gene_metadata','drug_metadata']]
if a.dataset=='tahoe':
    files += [(base+f'data/train-{i:05d}-of-03388.parquet',f'train-{i:05d}-of-03388.parquet',None) for i in (a.indices if a.indices is not None else range(a.shards))]
else:
    files += [('https://zenodo.org/records/13350497/files/SrivatsanTrapnell2020_sciplex3.h5ad?download=1','sciplex3.h5ad','c9e70629505d98c7ca1a837f62b14e89')]
for url,name,expected in files:
    path=out/name
    if path.exists():
        if expected and hashlib.file_digest(path.open('rb'),'md5').hexdigest()!=expected:
            raise ValueError(f'Checksum mismatch for {name}')
        print(f'Using {path}',flush=True);continue
    temp=path.with_suffix(path.suffix+'.part')
    with s.get(url,stream=True,timeout=(30,180)) as r:
        r.raise_for_status()
        with temp.open('wb') as f:
            for block in r.iter_content(8*1024*1024):f.write(block)
    if expected:
        with temp.open('rb') as f:digest=hashlib.file_digest(f,'md5').hexdigest()
        if digest!=expected:raise ValueError(f'Checksum mismatch for {name}')
    temp.replace(path)
    print(f'Downloaded {path} ({path.stat().st_size:,} bytes)',flush=True)