#!/usr/bin/env python3
"""
arxiv_oaipmh_harvest.py — Harvest arXiv math papers via OAI-PMH into arxiv DB.

Usage:
  python3 4-Infrastructure/shim/arxiv_oaipmh_harvest.py --set=math:math:NT
  python3 4-Infrastructure/shim/arxiv_oaipmh_harvest.py --set=math:math:NT --set=math:math:CO
  python3 4-Infrastructure/shim/arxiv_oaipmh_harvest.py --all-math
  python3 4-Infrastructure/shim/arxiv_oaipmh_harvest.py --append     # just append, no DROP
"""
from __future__ import annotations

import argparse
import subprocess
import sys
import time
import xml.etree.ElementTree as ET
from pathlib import Path
from urllib.request import urlopen, Request

OAI_BASE = "https://oaipmh.arxiv.org/oai"
NEON_HOST = "neon-64gb"
CONTAINER = "arxiv-pg"
DB = "arxiv"
BATCH_INSERT = 100

def ssh_psql(sql: str, timeout: int = 300) -> str:
    result = subprocess.run([
        "ssh", NEON_HOST,
        f"podman exec -i {CONTAINER} psql -U postgres -d {DB} -t -A"
    ], input=sql, capture_output=True, text=True, timeout=timeout)
    return result.stdout.strip()

def setup_table(append: bool = False):
    if not append:
        ssh_psql("DROP TABLE IF EXISTS arxiv_papers;")
    ssh_psql("""
        CREATE TABLE IF NOT EXISTS arxiv_papers (
            paper_id TEXT PRIMARY KEY,
            title TEXT NOT NULL DEFAULT '',
            categories TEXT NOT NULL DEFAULT '',
            authors TEXT NOT NULL DEFAULT '',
            abstract TEXT NOT NULL DEFAULT ''
        );
    """)
    print("Table arxiv_papers ready", file=sys.stderr)

def fetch_records(set_spec: str, resumption_token: str | None = None) -> tuple[list[dict], str | None]:
    if resumption_token:
        url = f"{OAI_BASE}?verb=ListRecords&resumptionToken={resumption_token}"
    else:
        url = f"{OAI_BASE}?verb=ListRecords&metadataPrefix=arXivRaw&set={set_spec}"

    req = Request(url, headers={"User-Agent": "ResearchStack/1.0 arxiv-harvester"})
    with urlopen(req, timeout=120) as resp:
        data = resp.read()

    root = ET.fromstring(data)
    ns = {
        "oai": "http://www.openarchives.org/OAI/2.0/",
        "arxiv": "http://arxiv.org/OAI/arXivRaw/",
    }

    records = []
    error = root.find(".//oai:error", ns)
    if error is not None:
        print(f"  OAI error: {error.text}", file=sys.stderr)
        return records, None

    for rec in root.findall(".//oai:record", ns):
        header = rec.find("oai:header", ns)
        if header is not None and header.get("status", "") == "deleted":
            continue

        metadata = rec.find("oai:metadata/arxiv:arXivRaw", ns)
        if metadata is None:
            continue

        paper_id = metadata.findtext("arxiv:id", "", ns).strip()
        if not paper_id:
            continue

        records.append({
            "paper_id": paper_id,
            "title": metadata.findtext("arxiv:title", "", ns).strip(),
            "categories": metadata.findtext("arxiv:categories", "", ns).strip(),
            "authors": metadata.findtext("arxiv:authors", "", ns).strip(),
            "abstract": metadata.findtext("arxiv:abstract", "", ns).strip(),
        })

    token_el = root.find(".//oai:resumptionToken", ns)
    next_token = token_el.text.strip() if token_el is not None and token_el.text else None

    return records, next_token

def escape_sql(s: str) -> str:
    return s.replace("'", "''").replace("\\", "\\\\")

def insert_batch(records: list[dict]):
    if not records:
        return 0
    values = []
    for r in records:
        pid = escape_sql(r["paper_id"])
        title = escape_sql(r["title"][:1000])
        cats = escape_sql(r["categories"][:500])
        authors = escape_sql(r["authors"][:500])
        abstract = escape_sql(r["abstract"][:5000])
        values.append(f"('{pid}','{title}','{cats}','{authors}','{abstract}')")

    sql = (f"INSERT INTO arxiv_papers (paper_id, title, categories, authors, abstract) VALUES "
           + ','.join(values)
           + " ON CONFLICT (paper_id) DO UPDATE SET "
           + "title=EXCLUDED.title, categories=EXCLUDED.categories, "
           + "authors=EXCLUDED.authors, abstract=EXCLUDED.abstract")
    ssh_psql(sql, timeout=120)
    return len(records)

def harvest(set_spec: str):
    print(f"Harvesting set: {set_spec}", file=sys.stderr)
    token = None
    total = 0
    batch = []
    t0 = time.time()
    page = 0

    while True:
        records, token = fetch_records(set_spec, token)
        batch.extend(records)
        total += len(records)
        page += 1

        if len(batch) >= BATCH_INSERT or (not records and token is None):
            insert_batch(batch)
            batch.clear()

        elapsed = time.time() - t0
        rate = total / elapsed if elapsed > 0 else 0
        if page % 10 == 0:
            print(f"  {total} records, {elapsed:.0f}s ({rate:.0f}/s), more={'yes' if token else 'no'}", file=sys.stderr)

        if token is None:
            break

        time.sleep(0.5)

    if batch:
        insert_batch(batch)

    elapsed = time.time() - t0
    print(f"  Done: {total} records in {elapsed:.0f}s", file=sys.stderr)

def main():
    ap = argparse.ArgumentParser(description="Harvest arXiv papers via OAI-PMH")
    ap.add_argument("--set", action="append", default=[], help="OAI set spec")
    ap.add_argument("--all-math", action="store_true", help="Harvest all math")
    ap.add_argument("--append", action="store_true", help="Don't DROP table first")
    args = ap.parse_args()

    sets = args.set
    if args.all_math:
        sets = ["math:math"]
    if not sets:
        sets = ["math:math:NT"]

    setup_table(append=args.append)
    for s in sets:
        harvest(s)

    count = ssh_psql("SELECT COUNT(*) FROM arxiv_papers")
    print(f"\nTotal rows: {count}", file=sys.stderr)

if __name__ == "__main__":
    main()