"""Fetch preview cache on another machine if the stand's public IP is rate limited."""
from concurrent.futures import ThreadPoolExecutor
import json
from pathlib import Path
import time
import urllib.parse
import urllib.request

cache = Path(__file__).parent / 'preview-cache'
cache.mkdir(exist_ok=True)


def fetch(url):
    for attempt in range(6):
        try:
            with urllib.request.urlopen(url, timeout=30) as response:
                return response.read()
        except Exception:
            if attempt == 5:
                raise
            time.sleep(min(30, 3 * (attempt + 1)))


def metadata(args):
    split, offset = args
    path = cache / f'rows-{split}-{offset}.json'
    if not path.exists():
        query = urllib.parse.urlencode(dict(dataset='uoft-cs/cifar10', config='plain_text',
                                           split=split, offset=offset, length=100))
        path.write_bytes(fetch('https://datasets-server.huggingface.co/rows?' + query))
        time.sleep(0.5)
    return [(split, r) for r in json.loads(path.read_text())['rows'] if r['row']['label'] in (3, 5)]


def image(args):
    split, row = args
    path = cache / f'{split}-{row["row_idx"]}.jpg'
    if not path.exists():
        path.write_bytes(fetch(row['row']['img']['src']))


tasks = [('train', i) for i in range(0, 6000, 100)] + [('test', i) for i in range(0, 3000, 100)]
with ThreadPoolExecutor(max_workers=2) as pool:
    images = [item for batch in pool.map(metadata, tasks) for item in batch]
print('metadata ready', len(images), flush=True)
with ThreadPoolExecutor(max_workers=4) as pool:
    for i, _ in enumerate(pool.map(image, images)):
        if i % 200 == 0:
            print('images', i, flush=True)
print('Cache ready', flush=True)