"""Parallel HTTP ranges from the official CIFAR host; validate full archive afterwards."""
from concurrent.futures import ThreadPoolExecutor, as_completed
import hashlib
from pathlib import Path
import time
import urllib.request

URL = 'https://cave.cs.toronto.edu/kriz/cifar-10-binary.tar.gz'
SIZE = 170052171
MD5 = 'c32a1d4ab5d03f1284b67883e8d87530'
root = Path('/home/int/run/cats-dogs/data')
parts = root / 'parts'
parts.mkdir(exist_ok=True)
chunk = (SIZE + 31) // 32


def fetch(i):
    lo, hi = i * chunk, min(SIZE, (i + 1) * chunk) - 1
    path = parts / str(i)
    if path.exists() and path.stat().st_size == hi - lo + 1:
        return i
    for attempt in range(3):
        try:
            request = urllib.request.Request(URL, headers={'Range': f'bytes={lo}-{hi}'})
            with urllib.request.urlopen(request, timeout=180) as response:
                assert response.status == 206
                assert response.headers['Content-Range'].startswith(f'bytes {lo}-{hi}/')
                data = response.read()
            assert len(data) == hi - lo + 1
            path.write_bytes(data)
            return i
        except Exception:
            if attempt == 2:
                raise
            time.sleep(1)


with ThreadPoolExecutor(max_workers=32) as pool:
    for future in as_completed([pool.submit(fetch, i) for i in range(32)]):
        print('part', future.result(), 'ready', flush=True)
archive = root / 'cifar-10-binary.tar.gz'
temporary = root / 'cifar-download.tmp'
with temporary.open('wb') as out:
    for i in range(32):
        out.write((parts / str(i)).read_bytes())
assert hashlib.md5(temporary.read_bytes()).hexdigest() == MD5
temporary.replace(archive)
print('Checksum OK', flush=True)