"""Extract clean text + a digest for each paper from arXiv HTML.

Usage: python extract_text.py <papers.json> <text_dir> <digest_dir>
"""
import json
import os
import re
import sys
import time
import html as htmlmod
import urllib.request

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


def fetch(url, timeout=90):
    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 html_to_text(t):
    t = re.sub(r'<script.*?</script>', ' ', t, flags=re.S | re.I)
    t = re.sub(r'<style.*?</style>', ' ', t, flags=re.S | re.I)
    t = re.sub(r'<!--.*?-->', ' ', t, flags=re.S)
    # headings
    t = re.sub(r'<h([1-6])[^>]*>(.*?)</h\1>', lambda m: '\n\n## ' + re.sub(r'<[^>]+>', ' ', m.group(2)) + '\n', t, flags=re.S | re.I)
    # block elements -> newline
    t = re.sub(r'</(p|div|li|tr|table|section|figure|figcaption|blockquote)>', '\n', t, flags=re.I)
    t = re.sub(r'<br\s*/?>', '\n', t, flags=re.I)
    t = re.sub(r'<[^>]+>', ' ', t)
    t = htmlmod.unescape(t)
    t = t.replace('\u00a0', ' ')
    lines = []
    for ln in t.split('\n'):
        ln = ' '.join(ln.split())
        if ln:
            lines.append(ln)
    return '\n'.join(lines)


def main():
    papers = json.load(open(sys.argv[1], encoding='utf-8'))
    tdir, ddir = sys.argv[2], sys.argv[3]
    os.makedirs(tdir, exist_ok=True)
    os.makedirs(ddir, exist_ok=True)
    report = []
    for p in papers:
        aid = p['id']
        tpath = os.path.join(tdir, aid + '.txt')
        dpath = os.path.join(ddir, aid + '.txt')
        if os.path.exists(tpath) and os.path.getsize(tpath) > 5000:
            report.append((aid, 'cached', os.path.getsize(tpath)))
            continue
        txt = None
        for suffix in ('', 'v1', 'v2', 'v3'):
            try:
                h = fetch('https://arxiv.org/html/%s%s' % (aid, suffix))
                if 'No HTML' in h[:5000] or len(h) < 5000:
                    continue
                txt = html_to_text(h)
                if len(txt) > 5000:
                    break
                txt = None
            except Exception:  # noqa: BLE001
                time.sleep(2)
        if not txt:
            report.append((aid, 'NO-HTML', 0))
            continue
        with open(tpath, 'w', encoding='utf-8') as f:
            f.write(txt)
        # digest: everything before References, capped
        m = re.search(r'\n## .{0,40}References', txt, re.I)
        body = txt[:m.start()] if m else txt
        head = body[:60000]
        tail = body[-14000:] if len(body) > 74000 else body[60000:]
        headings = [l for l in body.split('\n') if l.startswith('## ')]
        digest = '\n'.join([
            'TITLE: ' + p['title'],
            'ARXIV: ' + aid,
            'AUTHORS: ' + p['authors'][:1500],
            'DATE: ' + p['dateline'],
            '',
            '=== SECTION HEADINGS ===',
            '\n'.join(headings[:120]),
            '',
            '=== BODY (head) ===',
            head,
            '',
            '=== BODY (tail) ===',
            tail,
        ])
        with open(dpath, 'w', encoding='utf-8') as f:
            f.write(digest)
        report.append((aid, 'ok', len(txt)))
        print('ok  ', aid, len(txt), 'digest', len(digest))
        time.sleep(1)
    for r in report:
        print(r)
    bad = [r[0] for r in report if r[1] == 'NO-HTML']
    print('NO-HTML:', bad)


if __name__ == '__main__':
    main()