scripts/fetch_corpus.py

"""Download canonical RL papers from arXiv into data/papers/.

Idempotent: papers already on disk (with valid PDF content) are skipped.
Truncated/corrupt PDFs are auto-deleted and retried.

Uses `requests` with streaming + chunked writes — the `arxiv` library's
built-in downloader was truncating large PDFs at power-of-2 byte boundaries
on Windows.

Usage:
    uv run python scripts/fetch_corpus.py
"""
from __future__ import annotations

import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

import pymupdf
import requests
from rich.console import Console
from rich.table import Table

from researchpath.corpus import CANONICAL_RL_PAPERS

console = Console()
DATA_DIR = Path(__file__).resolve().parents[1] / "data" / "papers"

USER_AGENT = "ResearchPath/0.1.0 (mailto:chetanchowdary01@gmail.com)"
MIN_PDF_BYTES = 50_000
INTER_REQUEST_DELAY = 12  # arXiv asks for >=3s; PDF endpoint rate-limits harder
MAX_ATTEMPTS = 3
RETRY_BACKOFF_SEC = 30
DOWNLOAD_TIMEOUT = 180


def validate_pdf(path: Path) -> tuple[bool, str]:
    if not path.exists():
        return False, "file missing"
    if path.stat().st_size < MIN_PDF_BYTES:
        return False, f"too small ({path.stat().st_size} bytes)"
    try:
        with pymupdf.open(path) as doc:
            if doc.page_count < 1:
                return False, "0 pages"
            _ = doc[0].get_text()
        return True, "ok"
    except Exception as e:
        return False, f"parse error: {e}"


def download_pdf(arxiv_id: str, dest: Path) -> tuple[bool, str]:
    """Stream a single PDF from arXiv. Returns (success, message)."""
    url = f"https://arxiv.org/pdf/{arxiv_id}.pdf"
    headers = {"User-Agent": USER_AGENT}
    try:
        with requests.get(url, headers=headers, stream=True, timeout=DOWNLOAD_TIMEOUT) as r:
            if r.status_code == 429:
                return False, "rate limited (429)"
            r.raise_for_status()
            expected = int(r.headers.get("Content-Length", "0") or "0")
            written = 0
            with open(dest, "wb") as f:
                for chunk in r.iter_content(chunk_size=64 * 1024):
                    if chunk:
                        f.write(chunk)
                        written += len(chunk)
            if expected and written < expected:
                dest.unlink(missing_ok=True)
                return False, f"truncated: {written}/{expected} bytes"
        return True, f"ok ({written} bytes)"
    except requests.exceptions.HTTPError as e:
        return False, f"http {e.response.status_code}"
    except Exception as e:
        return False, f"network: {e}"


def fetch_with_retry(arxiv_id: str, dest: Path) -> tuple[bool, str]:
    for attempt in range(1, MAX_ATTEMPTS + 1):
        ok, msg = download_pdf(arxiv_id, dest)
        if not ok:
            if attempt < MAX_ATTEMPTS:
                console.print(f"  [dim]attempt {attempt}: {msg}; backing off {RETRY_BACKOFF_SEC}s[/dim]")
                time.sleep(RETRY_BACKOFF_SEC)
                continue
            return False, msg

        valid, vmsg = validate_pdf(dest)
        if valid:
            return True, msg
        dest.unlink(missing_ok=True)
        if attempt < MAX_ATTEMPTS:
            console.print(f"  [dim]attempt {attempt}: invalid pdf ({vmsg}); backing off {RETRY_BACKOFF_SEC}s[/dim]")
            time.sleep(RETRY_BACKOFF_SEC)
    return False, "exhausted retries"


def main() -> int:
    DATA_DIR.mkdir(parents=True, exist_ok=True)

    n_ok = 0
    n_skip = 0
    n_fail = 0
    failures: list[tuple[str, str]] = []

    pending = []
    for paper in CANONICAL_RL_PAPERS:
        target = DATA_DIR / f"{paper.arxiv_id}.pdf"
        if target.exists():
            ok, reason = validate_pdf(target)
            if ok:
                console.print(f"[yellow]SKIP[/yellow]  {paper.arxiv_id}  ({paper.tag})")
                n_skip += 1
                continue
            console.print(f"[magenta]REDO[/magenta]  {paper.arxiv_id}  ({paper.tag}) — {reason}")
            target.unlink()
        pending.append(paper)

    if not pending:
        console.print("\n[bold green]All papers already present and valid.[/bold green]")
        return 0

    console.print(f"\n[bold]Need to download {len(pending)} paper(s). Spacing requests by {INTER_REQUEST_DELAY}s.[/bold]\n")

    for i, paper in enumerate(pending):
        target = DATA_DIR / f"{paper.arxiv_id}.pdf"
        console.print(f"[cyan]FETCH[/cyan] {paper.arxiv_id}  {paper.tag} - {paper.title[:60]}")
        ok, msg = fetch_with_retry(paper.arxiv_id, target)
        if ok:
            console.print(f"  [green]ok[/green] {msg}")
            n_ok += 1
        else:
            console.print(f"[red]FAIL[/red]  {paper.arxiv_id}: {msg}")
            failures.append((paper.arxiv_id, msg))
            n_fail += 1
        # Spacing between successful papers (within retry, the backoff is already 30s)
        if i + 1 < len(pending):
            time.sleep(INTER_REQUEST_DELAY)

    console.print()
    table = Table(title="Corpus fetch summary")
    table.add_column("Outcome", style="bold")
    table.add_column("Count", justify="right")
    table.add_row("[green]Downloaded[/green]", str(n_ok))
    table.add_row("[yellow]Skipped (already valid)[/yellow]", str(n_skip))
    table.add_row("[red]Failed[/red]", str(n_fail))
    console.print(table)

    if failures:
        console.print("\n[bold red]Failures (re-run later to retry):[/bold red]")
        for arxiv_id, err in failures:
            console.print(f"  {arxiv_id}: {err}")

    return 0 if n_fail == 0 else 1


if __name__ == "__main__":
    sys.exit(main())