"""Scrape arXiv search HTML for a query and dump id/title/authors.

Usage: python scrape_search.py <query> <pages> <out.json>
"""
import urllib.request
import urllib.parse
import re
import json
import sys
import time
import html as htmlmod

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


def get(url, timeout=40):
    req = urllib.request.Request(url, headers={'User-Agent': UA})
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return r.read().decode('utf-8', 'ignore')


def strip(s):
    s = re.sub(r'<[^>]+>', ' ', s)
    s = htmlmod.unescape(s)
    return ' '.join(s.split())


def parse_page(t):
    out = []
    # each result block
    for m in re.finditer(r'<li class="arxiv-result">(.*?)</li>\s*(?=<li class="arxiv-result">|</ol>)', t, re.S):
        blk = m.group(1)
        aid = re.search(r'arxiv\.org/abs/([^"\'<]+)', blk)
        title = re.search(r'<p class="title is-5 mathjax">(.*?)</p>', blk, re.S)
        authors = re.search(r'<p class="authors">(.*?)</p>', blk, re.S)
        dates = re.search(r'<p class="is-size-7">(.*?)</p>', blk, re.S)
        if not aid:
            continue
        out.append({
            'id': aid.group(1),
            'title': strip(title.group(1)) if title else '',
            'authors': strip(authors.group(1)) if authors else '',
            'dates': strip(dates.group(1)) if dates else '',
        })
    return out


def main():
    q = sys.argv[1]
    pages = int(sys.argv[2]) if len(sys.argv) > 2 else 1
    out_path = sys.argv[3] if len(sys.argv) > 3 else 'search.json'
    size = 200
    res = []
    for p in range(pages):
        url = ('https://arxiv.org/search/?searchtype=all&query=' + urllib.parse.quote(q)
               + '&start=' + str(p * size) + '&size=' + str(size))
        try:
            t = get(url)
        except Exception as e:  # noqa: BLE001
            print('page fail', p, e, file=sys.stderr)
            break
        rows = parse_page(t)
        total = re.search(r'of\s+([\d,]+)\s+results', t)
        print('page', p, 'rows', len(rows), 'total', total.group(1) if total else '?', file=sys.stderr)
        res.extend(rows)
        if len(rows) < size:
            break
        time.sleep(2)
    seen = {}
    for r in res:
        seen.setdefault(r['id'].split('v')[0], r)
    with open(out_path, 'w', encoding='utf-8') as f:
        json.dump(list(seen.values()), f, ensure_ascii=False, indent=1)
    print('saved', len(seen), file=sys.stderr)


if __name__ == '__main__':
    main()