stack/dev/scripts/fetch_pdfs.py

"""Headless PDF retrieval for skin substitute evidence base.

Downloads PDFs via Unpaywall (legal, ~50%), PMC (free), and SciHub
(fallback) without requiring Zotero's GUI. Attaches downloaded PDFs
to Zotero items in the database.

Usage:
    uv run python dev/scripts/fetch_pdfs.py
    uv run python dev/scripts/fetch_pdfs.py --limit 100
    uv run python dev/scripts/fetch_pdfs.py --source unpaywall
    uv run python dev/scripts/fetch_pdfs.py --source pmc
    uv run python dev/scripts/fetch_pdfs.py --source scihub
    uv run python dev/scripts/fetch_pdfs.py --source s2

    # With SciHub proxy (for blocked networks):
    SCIHUB_PROXY=socks5://localhost:1080 uv run python dev/scripts/fetch_pdfs.py --source scihub
"""

from __future__ import annotations

import argparse
import os
import re
import sqlite3
import time
from pathlib import Path

import httpx

from conf import path as conf_path
from zot.db import FIELD_MAP, TYPE_MAP, generate_key, now_iso

ZOTERO_DB = str(conf_path("db.zotero"))
STORAGE_DIR = Path(conf_path("storage.zotero"))

UNPAYWALL_EMAIL = "dev@fhirworx.io"
USER_AGENT = "stack-pdf-fetcher/1.0 (mailto:{})".format(UNPAYWALL_EMAIL)


# ---------------------------------------------------------------------------
# Source: Unpaywall (legal, no auth, ~50% coverage)
# ---------------------------------------------------------------------------


def fetch_unpaywall(doi: str) -> str | None:
    """Get PDF URL from Unpaywall for a DOI. Returns URL or None."""
    url = f"https://api.unpaywall.org/v2/{doi}"
    try:
        resp = httpx.get(
            url,
            params={"email": UNPAYWALL_EMAIL},
            headers={"User-Agent": USER_AGENT},
            timeout=15,
        )
        if resp.status_code != 200:
            return None
        data = resp.json()
        best = data.get("best_oa_location") or {}
        return best.get("url_for_pdf") or best.get("url")
    except Exception:
        return None


# ---------------------------------------------------------------------------
# Source: PubMed Central (free for OA articles)
# ---------------------------------------------------------------------------


def doi_to_pmcid(dois: list[str]) -> dict[str, str]:
    """Convert DOIs to PMCIDs via NCBI ID Converter. Returns {doi: pmcid}."""
    result = {}
    # API accepts up to 200 IDs per request
    for i in range(0, len(dois), 200):
        batch = dois[i : i + 200]
        try:
            resp = httpx.get(
                "https://www.ncbi.nlm.nih.gov/pmc/utils/idconv/v1.0/",
                params={
                    "ids": ",".join(batch),
                    "format": "json",
                    "tool": "stack-pdf-fetcher",
                    "email": UNPAYWALL_EMAIL,
                },
                timeout=30,
            )
            if resp.status_code == 200:
                for rec in resp.json().get("records", []):
                    doi = rec.get("doi", "")
                    pmcid = rec.get("pmcid", "")
                    if doi and pmcid:
                        result[doi.lower()] = pmcid
        except Exception:
            pass
        time.sleep(0.34)
    return result


def fetch_pmc(pmcid: str) -> str | None:
    """Get PDF URL for a PMC article. Returns URL or None."""
    return f"https://pmc.ncbi.nlm.nih.gov/articles/{pmcid}/pdf/"


# ---------------------------------------------------------------------------
# Source: Semantic Scholar (supplementary, batch endpoint)
# ---------------------------------------------------------------------------


def fetch_s2_batch(dois: list[str]) -> dict[str, str]:
    """Batch lookup OA PDF URLs from Semantic Scholar. Returns {doi: pdf_url}."""
    result = {}
    for i in range(0, len(dois), 500):
        batch = [f"DOI:{d}" for d in dois[i : i + 500]]
        try:
            resp = httpx.post(
                "https://api.semanticscholar.org/graph/v1/paper/batch",
                json={"ids": batch, "fields": "openAccessPdf,externalIds"},
                timeout=30,
            )
            if resp.status_code == 200:
                for paper in resp.json():
                    if paper and paper.get("openAccessPdf"):
                        pdf_url = paper["openAccessPdf"].get("url")
                        ext_ids = paper.get("externalIds", {})
                        doi = ext_ids.get("DOI", "")
                        if pdf_url and doi:
                            result[doi.lower()] = pdf_url
        except Exception:
            pass
        time.sleep(1)
    return result


# ---------------------------------------------------------------------------
# Source: SciHub (fallback)
# ---------------------------------------------------------------------------

SCIHUB_MIRRORS = [
    "https://sci-hub.se",
    "https://sci-hub.st",
    "https://sci-hub.ru",
    "https://sci-hub.ren",
    "https://sci-hub.ee",
    "https://sci-hub.wf",
]

# Proxy for SciHub (set SCIHUB_PROXY env var, e.g. socks5://localhost:1080)
SCIHUB_PROXY = os.environ.get("SCIHUB_PROXY", "")

# Cloudflare Worker proxy (preferred — bypasses ISP blocking via CF edge)
CF_WORKER_URL = os.environ.get(
    "PDF_PROXY_URL", "https://pdf-proxy.lite7889.workers.dev"
)
CF_WORKER_SECRET = os.environ.get("PDF_PROXY_SECRET", "")


def fetch_via_worker(doi: str, dest: Path) -> bool:
    """Fetch PDF via Cloudflare Worker proxy. Returns True on success.

    The Worker fetches from SciHub on Cloudflare's edge network,
    bypassing ISP-level DNS/IP blocking.
    """
    if not CF_WORKER_SECRET:
        return False
    try:
        resp = httpx.get(
            f"{CF_WORKER_URL}/pdf",
            params={"doi": doi},
            headers={
                "X-Proxy-Key": CF_WORKER_SECRET,
                "User-Agent": USER_AGENT,
            },
            timeout=45,
        )
        if resp.status_code != 200:
            return False
        if b"%PDF" not in resp.content[:1024]:
            return False
        dest.parent.mkdir(parents=True, exist_ok=True)
        dest.write_bytes(resp.content)
        return True
    except Exception:
        return False


def fetch_scihub(doi: str) -> str | None:
    """Get PDF URL from SciHub directly. Returns direct PDF URL or None.

    Set SCIHUB_PROXY env var for networks where SciHub is blocked.
    Prefer fetch_via_worker() which uses Cloudflare edge.
    """
    transport = httpx.HTTPTransport(proxy=SCIHUB_PROXY) if SCIHUB_PROXY else None
    client_kwargs = {"transport": transport} if transport else {}

    for mirror in SCIHUB_MIRRORS:
        try:
            with httpx.Client(
                timeout=15, follow_redirects=True, **client_kwargs
            ) as client:
                resp = client.get(
                    f"{mirror}/{doi}",
                    headers={"User-Agent": "Mozilla/5.0"},
                )
            if resp.status_code != 200:
                continue
            # Parse HTML for #pdf iframe/embed src
            m = re.search(r'id="pdf"[^>]*src="([^"]+)"', resp.text)
            if not m:
                m = re.search(r'<iframe[^>]*src="([^"]*\.pdf[^"]*)"', resp.text)
            if not m:
                m = re.search(r'<embed[^>]*src="([^"]*\.pdf[^"]*)"', resp.text)
            if m:
                pdf_url = m.group(1)
                if pdf_url.startswith("//"):
                    pdf_url = "https:" + pdf_url
                elif pdf_url.startswith("/"):
                    pdf_url = mirror + pdf_url
                return pdf_url
        except Exception:
            continue
    return None


# ---------------------------------------------------------------------------
# Download PDF
# ---------------------------------------------------------------------------


def download_pdf(url: str, dest: Path) -> bool:
    """Download a PDF file. Returns True on success."""
    try:
        resp = httpx.get(
            url,
            headers={"User-Agent": USER_AGENT},
            timeout=30,
            follow_redirects=True,
        )
        if resp.status_code != 200:
            return False
        content_type = resp.headers.get("content-type", "")
        # Verify it's actually a PDF
        if b"%PDF" not in resp.content[:1024] and "pdf" not in content_type:
            return False
        dest.parent.mkdir(parents=True, exist_ok=True)
        dest.write_bytes(resp.content)
        return True
    except Exception:
        return False


# ---------------------------------------------------------------------------
# Zotero attachment
# ---------------------------------------------------------------------------


def attach_pdf_to_item(
    con: sqlite3.Connection, item_id: int, pdf_path: Path, title: str = ""
) -> bool:
    """Create a Zotero attachment record for a downloaded PDF."""
    key = generate_key()
    while con.execute("SELECT 1 FROM items WHERE key = ?", (key,)).fetchone():
        key = generate_key()

    now = now_iso()

    # Create attachment item (itemTypeID=14 = attachment)
    cur = con.execute(
        """INSERT INTO items
           (itemTypeID, dateAdded, dateModified, clientDateModified,
            libraryID, key, version, synced)
           VALUES (?, ?, ?, ?, 1, ?, 0, 0)""",
        (TYPE_MAP["attachment"], now, now, now, key),
    )
    att_item_id = cur.lastrowid

    # Create attachment record
    # linkMode: 0 = imported file, 1 = imported URL, 2 = linked file
    con.execute(
        """INSERT INTO itemAttachments
           (itemID, parentItemID, linkMode, contentType, path, storageModTime)
           VALUES (?, ?, 0, 'application/pdf', ?, ?)""",
        (
            att_item_id,
            item_id,
            f"storage:{key}/{pdf_path.name}",
            int(time.time() * 1000),
        ),
    )

    # Move PDF to Zotero storage
    storage_dir = STORAGE_DIR / key
    storage_dir.mkdir(parents=True, exist_ok=True)
    dest = storage_dir / pdf_path.name
    if pdf_path != dest:
        import shutil

        shutil.copy2(pdf_path, dest)

    con.commit()
    return True


# ---------------------------------------------------------------------------
# Main pipeline
# ---------------------------------------------------------------------------


def get_items_needing_pdfs(con: sqlite3.Connection) -> list[dict]:
    """Get skin-subs items with DOI but no PDF attachment."""
    rows = con.execute(
        """
        SELECT DISTINCT i.itemID,
               (SELECT idv.value FROM itemData id
                JOIN itemDataValues idv ON id.valueID = idv.valueID
                WHERE id.itemID = i.itemID AND id.fieldID = ?) AS doi
        FROM items i
        JOIN itemTags it ON i.itemID = it.itemID
        JOIN tags t ON it.tagID = t.tagID
        WHERE t.name = 'module:skin-subs'
          AND i.itemTypeID = ?
          AND i.itemID NOT IN (
              SELECT DISTINCT parentItemID FROM itemAttachments
              WHERE parentItemID IS NOT NULL
                AND contentType = 'application/pdf'
          )
    """,
        (FIELD_MAP["DOI"], TYPE_MAP["journalArticle"]),
    ).fetchall()
    return [{"item_id": r[0], "doi": r[1]} for r in rows if r[1]]


def main() -> None:
    parser = argparse.ArgumentParser(description="Headless PDF retrieval")
    parser.add_argument(
        "--limit", type=int, default=0, help="Max items to process (0=all)"
    )
    parser.add_argument(
        "--source",
        choices=["all", "unpaywall", "pmc", "s2", "worker", "scihub"],
        default="all",
        help="Which source to use",
    )
    parser.add_argument(
        "--proxy", help="SOCKS/HTTP proxy for SciHub (e.g. socks5://localhost:1080)"
    )
    args = parser.parse_args()

    global SCIHUB_PROXY
    if args.proxy:
        SCIHUB_PROXY = args.proxy

    print("=" * 60)
    print("Headless PDF Retrieval")
    if SCIHUB_PROXY:
        print(f"  SciHub proxy: {SCIHUB_PROXY}")
    print("=" * 60)

    con = sqlite3.connect(ZOTERO_DB, timeout=10)
    con.row_factory = sqlite3.Row

    items = get_items_needing_pdfs(con)
    if args.limit:
        items = items[: args.limit]
    print(f"\nItems needing PDFs: {len(items)}")

    tmp_dir = Path("/tmp/pdf_downloads")
    tmp_dir.mkdir(exist_ok=True)

    stats = {"unpaywall": 0, "pmc": 0, "scihub": 0, "failed": 0, "total": 0}

    # Phase 1: Unpaywall (fast, legal, no CAPTCHA)
    if args.source in ("all", "unpaywall"):
        print("\n--- Phase 1: Unpaywall ---")
        for i, item in enumerate(items):
            if not item["doi"]:
                continue
            stats["total"] += 1

            pdf_url = fetch_unpaywall(item["doi"])
            if pdf_url:
                dest = tmp_dir / f"{item['item_id']}.pdf"
                if download_pdf(pdf_url, dest):
                    attach_pdf_to_item(con, item["item_id"], dest, item["doi"])
                    stats["unpaywall"] += 1
                    item["_done"] = True

            if (i + 1) % 50 == 0:
                print(f"  {i + 1}/{len(items)}  unpaywall={stats['unpaywall']}")
            time.sleep(0.1)  # Unpaywall is generous but be polite

        print(f"  Unpaywall: {stats['unpaywall']} PDFs")

    # Phase 2: PMC (batch convert DOIs to PMCIDs, then download)
    remaining = [it for it in items if not it.get("_done")]
    if args.source in ("all", "pmc") and remaining:
        print(f"\n--- Phase 2: PubMed Central ({len(remaining)} remaining) ---")
        dois = [it["doi"] for it in remaining if it["doi"]]
        pmcid_map = doi_to_pmcid(dois)
        print(f"  DOIs with PMCIDs: {len(pmcid_map)}")

        for item in remaining:
            doi = (item["doi"] or "").lower()
            if doi in pmcid_map:
                pdf_url = fetch_pmc(pmcid_map[doi])
                dest = tmp_dir / f"{item['item_id']}.pdf"
                if download_pdf(pdf_url, dest):
                    attach_pdf_to_item(con, item["item_id"], dest, item["doi"])
                    stats["pmc"] += 1
                    item["_done"] = True
                time.sleep(0.34)

        print(f"  PMC: {stats['pmc']} PDFs")

    # Phase 3: Semantic Scholar (batch, catches some repos Unpaywall misses)
    remaining = [it for it in items if not it.get("_done")]
    if args.source in ("all", "s2") and remaining:
        print(f"\n--- Phase 3: Semantic Scholar ({len(remaining)} remaining) ---")
        dois = [it["doi"] for it in remaining if it["doi"]]
        s2_map = fetch_s2_batch(dois)
        print(f"  S2 OA PDFs found: {len(s2_map)}")

        for item in remaining:
            doi = (item["doi"] or "").lower()
            if doi in s2_map:
                dest = tmp_dir / f"{item['item_id']}.pdf"
                if download_pdf(s2_map[doi], dest):
                    attach_pdf_to_item(con, item["item_id"], dest, item["doi"])
                    stats["s2"] = stats.get("s2", 0) + 1
                    item["_done"] = True
                time.sleep(0.2)

        print(f"  Semantic Scholar: {stats.get('s2', 0)} PDFs")

    # Phase 4: Cloudflare Worker → SciHub (bypasses ISP blocking)
    remaining = [it for it in items if not it.get("_done")]
    if args.source in ("all", "worker", "scihub") and remaining and CF_WORKER_SECRET:
        print(f"\n--- Phase 4: CF Worker → SciHub ({len(remaining)} remaining) ---")
        consecutive_failures = 0
        for i, item in enumerate(remaining):
            if consecutive_failures >= 20:
                print(f"  Stopping: {consecutive_failures} consecutive failures")
                break
            dest = tmp_dir / f"{item['item_id']}.pdf"
            if fetch_via_worker(item["doi"], dest):
                attach_pdf_to_item(con, item["item_id"], dest, item["doi"])
                stats["worker"] = stats.get("worker", 0) + 1
                item["_done"] = True
                consecutive_failures = 0
            else:
                consecutive_failures += 1
            if (i + 1) % 25 == 0:
                print(
                    f"  {i + 1}/{len(remaining)}  worker={stats.get('worker', 0)}  "
                    f"fails={consecutive_failures}"
                )
            time.sleep(1.5)  # Be polite to Worker + SciHub
        print(f"  CF Worker: {stats.get('worker', 0)} PDFs")

    # Phase 5: SciHub direct (fallback if Worker unavailable)
    remaining = [it for it in items if not it.get("_done")]
    if args.source in ("scihub",) and remaining and not CF_WORKER_SECRET:
        print(f"\n--- Phase 5: SciHub direct ({len(remaining)} remaining) ---")
        if not SCIHUB_PROXY:
            print("  WARNING: No proxy set. SciHub may be blocked from this network.")
            print("  Set SCIHUB_PROXY or use --proxy socks5://host:port")
        consecutive_failures = 0
        for i, item in enumerate(remaining):
            if consecutive_failures >= 10:
                print(
                    f"  Stopping: {consecutive_failures} consecutive failures (CAPTCHA?)"
                )
                break

            pdf_url = fetch_scihub(item["doi"])
            if pdf_url:
                dest = tmp_dir / f"{item['item_id']}.pdf"
                if download_pdf(pdf_url, dest):
                    attach_pdf_to_item(con, item["item_id"], dest, item["doi"])
                    stats["scihub"] += 1
                    item["_done"] = True
                    consecutive_failures = 0
                else:
                    consecutive_failures += 1
            else:
                consecutive_failures += 1

            if (i + 1) % 25 == 0:
                print(
                    f"  {i + 1}/{len(remaining)}  scihub={stats['scihub']}  "
                    f"fails={consecutive_failures}"
                )
            time.sleep(1)  # Be polite to SciHub

        print(f"  SciHub: {stats['scihub']} PDFs")

    stats["failed"] = len([it for it in items if not it.get("_done")])
    total_found = (
        stats["unpaywall"]
        + stats["pmc"]
        + stats.get("s2", 0)
        + stats.get("worker", 0)
        + stats.get("scihub", 0)
    )

    print(f"\n{'=' * 60}")
    print(f"Results: {total_found} PDFs downloaded")
    print(f"  Unpaywall: {stats['unpaywall']}")
    print(f"  PMC:       {stats['pmc']}")
    print(f"  S2:        {stats.get('s2', 0)}")
    print(f"  CF Worker: {stats.get('worker', 0)}")
    print(f"  SciHub:    {stats.get('scihub', 0)}")
    print(f"  Failed:    {stats['failed']}")
    print(f"  Coverage:  {total_found * 100 // max(len(items), 1)}%")

    con.close()


if __name__ == "__main__":
    main()