#!/usr/bin/env python3
"""WP44 — harvest arXiv metadata (set=cs) through OAI-PMH (metadataPrefix=arXiv).

Harvests by monthly datestamp windows from --from to --until; each window's raw
XML pages are stored verbatim (immutable raw snapshot) under <out>/raw/<window>/
page_NNNN.xml and parsed into <out>/arxiv_cs.parquet at the end (``--parse-only``
re-parses). Politeness: >= 3 s between requests, honours Retry-After on 503.
Note: OAI ``from``/``until`` filter on the record *datestamp* (last update), so a
window contains every record updated in that window; ``created`` gives the first
submission date. Duplicates across windows are resolved at parse time (latest
datestamp wins).
"""
import argparse, glob, json, os, re, time, datetime as dt, calendar
import requests
from lxml import etree
import pyarrow as pa, pyarrow.parquet as pq

BASE = "https://oaipmh.arxiv.org/oai"
NS = {"oai": "http://www.openarchives.org/OAI/2.0/", "ax": "http://arxiv.org/OAI/arXiv/"}
UA = "WP44-UQO-systematic-mapping/0.1 (academic research; contact simon-pierre.boucher@uqo.ca)"
SCHEMA = pa.schema([
    ("arxiv_id", pa.string()), ("datestamp", pa.string()), ("created", pa.string()), ("updated", pa.string()),
    ("title", pa.string()), ("abstract", pa.string()), ("authors", pa.list_(pa.string())), ("categories", pa.list_(pa.string())),
    ("primary_category", pa.string()), ("comments", pa.string()), ("journal_ref", pa.string()), ("doi", pa.string()),
    ("license", pa.string()), ("report_no", pa.string()), ("msc_class", pa.string()), ("acm_class", pa.string()), ("set_specs", pa.list_(pa.string())),
])


def month_windows(start, end):
    y, m = start.year, start.month
    while (y, m) <= (end.year, end.month):
        last = calendar.monthrange(y, m)[1]
        a = dt.date(y, m, 1); b = min(dt.date(y, m, last), end)
        yield a, b
        m += 1
        if m == 13:
            y, m = y + 1, 1


def fetch(session, params, log):
    for attempt in range(12):
        try:
            r = session.get(BASE, params=params, timeout=180, headers={"User-Agent": UA})
        except requests.RequestException as e:
            log(f"  network error {e}; sleep 30"); time.sleep(30); continue
        if r.status_code == 503:
            wait = int(r.headers.get("Retry-After", "30")); log(f"  503 retry-after {wait}"); time.sleep(wait + 1); continue
        if r.status_code != 200:
            log(f"  HTTP {r.status_code}; sleep 30"); time.sleep(30); continue
        return r.text
    raise RuntimeError("too many failures")


def harvest(a, log):
    s = requests.Session()
    state_path = os.path.join(a.out, "harvest_state.json")
    state = json.load(open(state_path)) if os.path.exists(state_path) else {"windows": {}}
    for w0, w1 in month_windows(dt.date.fromisoformat(a.frm), dt.date.fromisoformat(a.until)):
        key = f"{w0}_{w1}"
        if state["windows"].get(key, {}).get("status") == "done":
            continue
        d = os.path.join(a.out, "raw", key); os.makedirs(d, exist_ok=True)
        params = {"verb": "ListRecords", "metadataPrefix": "arXiv", "set": a.set, "from": str(w0), "until": str(w1)}
        page = 0; nrec = 0; t0 = time.time()
        while True:
            xml = fetch(s, params, log)
            with open(os.path.join(d, f"page_{page:04d}.xml"), "w") as fh:
                fh.write(xml)
            root = etree.fromstring(xml.encode())
            err = root.find("oai:error", NS)
            if err is not None:
                if err.get("code") == "noRecordsMatch":
                    break
                raise RuntimeError(f"OAI error {err.get('code')}: {err.text}")
            nrec += len(root.findall(".//oai:record", NS))
            tok = root.find(".//oai:resumptionToken", NS)
            page += 1
            if tok is None or not (tok.text or "").strip():
                break
            params = {"verb": "ListRecords", "resumptionToken": tok.text.strip()}
            time.sleep(a.delay)
        state["windows"][key] = {"status": "done", "pages": page, "records": nrec, "seconds": round(time.time() - t0),
                                 "when": dt.datetime.now(dt.UTC).isoformat()}
        json.dump(state, open(state_path, "w"), indent=1)
        log(f"window {key}: {page} pages, {nrec} records, {time.time()-t0:.0f}s")
        time.sleep(a.delay)


def parse(a, log):
    recs = {}
    files = sorted(glob.glob(os.path.join(a.out, "raw", "*", "page_*.xml")))
    for f in files:
        root = etree.parse(f).getroot()
        for rec in root.findall(".//oai:record", NS):
            h = rec.find("oai:header", NS)
            if h.get("status") == "deleted":
                continue
            m = rec.find("oai:metadata/ax:arXiv", NS)
            if m is None:
                continue
            g = lambda t: (m.findtext(f"ax:{t}", namespaces=NS) or "").strip() or None
            aid = g("id"); ds = h.findtext("oai:datestamp", namespaces=NS)
            if aid in recs and recs[aid]["datestamp"] >= ds:
                continue
            cats = (g("categories") or "").split()
            authors = []
            for au in m.findall("ax:authors/ax:author", NS):
                kn = au.findtext("ax:keyname", namespaces=NS) or ""; fn = au.findtext("ax:forenames", namespaces=NS) or ""
                authors.append((fn + " " + kn).strip())
            recs[aid] = dict(arxiv_id=aid, datestamp=ds, created=g("created"), updated=g("updated"),
                             title=re.sub(r"\s+", " ", g("title") or ""), abstract=re.sub(r"\s+", " ", g("abstract") or ""),
                             authors=authors, categories=cats, primary_category=cats[0] if cats else None,
                             comments=g("comments"), journal_ref=g("journal-ref"), doi=(g("doi") or "").lower() or None, license=g("license"),
                             report_no=g("report-no"), msc_class=g("msc-class"), acm_class=g("acm-class"),
                             set_specs=[x.text for x in h.findall("oai:setSpec", NS)])
    cols = {n: [r[n] for r in recs.values()] for n in SCHEMA.names}
    pq.write_table(pa.Table.from_pydict(cols, schema=SCHEMA), os.path.join(a.out, "arxiv_cs.parquet"), compression="zstd")
    rep = {"files": len(files), "records": len(recs), "parsed_at": dt.datetime.now(dt.UTC).isoformat(), "set": a.set, "from": a.frm, "until": a.until}
    json.dump(rep, open(os.path.join(a.out, "arxiv_parse_report.json"), "w"), indent=1); log(json.dumps(rep))


def main():
    ap = argparse.ArgumentParser(); ap.add_argument("--out", required=True); ap.add_argument("--set", default="cs")
    ap.add_argument("--from", dest="frm", default="2018-01-01"); ap.add_argument("--until", default="2026-11-02")
    ap.add_argument("--delay", type=float, default=3.5); ap.add_argument("--parse-only", action="store_true")
    a = ap.parse_args(); os.makedirs(a.out, exist_ok=True)
    log = lambda s: print(f"[{dt.datetime.now():%H:%M:%S}] {s}", flush=True)
    if not a.parse_only:
        harvest(a, log)
    parse(a, log)


if __name__ == "__main__":
    main()