"""Download arXiv PDFs for the master paper list.

Usage: python download.py <papers.json> <pdf_dir>
Idempotent: skips files that already exist with a plausible size.
"""
import json
import os
import sys
import time
import urllib.request

UA = 'Mozilla/5.0 (compatible; paperbot/1.0)'
MIN_BYTES = 20000


def fetch(url, timeout=180):
    req = urllib.request.Request(url, headers={'User-Agent': UA})
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return r.read()


def main():
    papers = json.load(open(sys.argv[1], encoding='utf-8'))
    outdir = sys.argv[2]
    os.makedirs(outdir, exist_ok=True)
    ok = 0
    fail = []
    for p in papers:
        dest = os.path.join(outdir, p['id'] + '.pdf')
        if os.path.exists(dest) and os.path.getsize(dest) > MIN_BYTES:
            print('SKIP', p['id'], os.path.getsize(dest))
            ok += 1
            continue
        got = False
        for attempt in range(3):
            try:
                data = fetch(p['pdf_url'])
                if data[:4] == b'%PDF' and len(data) > MIN_BYTES:
                    with open(dest, 'wb') as f:
                        f.write(data)
                    print('OK  ', p['id'], len(data), p['title'][:60])
                    ok += 1
                    got = True
                    break
                print('  bad payload', p['id'], len(data), data[:20])
            except Exception as e:  # noqa: BLE001
                print('  retry', attempt, p['id'], str(e)[:100])
            time.sleep(3)
        if not got:
            fail.append(p['id'])
            print('FAIL', p['id'])
        time.sleep(1)
    print('=== downloaded/verified %d/%d, failures: %s' % (ok, len(papers), fail))


if __name__ == '__main__':
    main()