"""Download PDFs by arXiv id list.

Usage: python download_ids.py <ids.txt> <out_dir>
"""
import os
import sys
import time
import urllib.request

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


def main():
    ids = [l.strip() for l in open(sys.argv[1], encoding='utf-8') if l.strip()]
    outdir = sys.argv[2]
    os.makedirs(outdir, exist_ok=True)
    fail = []
    for aid in ids:
        dest = os.path.join(outdir, aid + '.pdf')
        if os.path.exists(dest) and os.path.getsize(dest) > MIN:
            print('SKIP', aid, os.path.getsize(dest))
            continue
        done = False
        for attempt in range(4):
            try:
                req = urllib.request.Request('https://arxiv.org/pdf/' + aid,
                                             headers={'User-Agent': UA})
                data = urllib.request.urlopen(req, timeout=180).read()
                if data[:4] == b'%PDF' and len(data) > MIN:
                    open(dest, 'wb').write(data)
                    print('OK  ', aid, len(data))
                    done = True
                    break
                print('  bad', aid, len(data))
            except Exception as e:  # noqa: BLE001
                print('  retry', attempt, aid, str(e)[:90])
            time.sleep(3)
        if not done:
            fail.append(aid)
    print('failures:', fail)


if __name__ == '__main__':
    main()