#!/usr/bin/env python3
"""Bounded, rate-limited, count-reconciling pagination for this skill's APIs.

Six of the ten databases here paginate differently -- absolute record offsets,
opaque cursors, continuation tokens, 1-based pages -- and each reports totals its
own way. Re-deriving the walk per query is how records get silently dropped. The
worst case is bioRxiv: `cursor` is an absolute offset, `/details/` returns 30 per
page but `/pubs/` returns 100, and an out-of-step cursor returns **HTTP 200**, so
stepping by 100 skips records 30-99 of every hundred and looks successful.

Every walk here:

- steps by the page size the response actually reported, never an assumed one
- stops on this API's real terminator (Europe PMC echoes your cursor back rather
  than sending null; bioRxiv just returns an empty collection)
- reconciles retrieved against the expected total and **exits 4 on a shortfall**
- refuses to exceed --max-records / --max-calls, and says so rather than
  truncating quietly

    python3 paginate.py --api biorxiv --query 2024-01-01/2024-01-03
    python3 paginate.py --api europepmc --query 'SRC:"PPR" AND "organoid"' --max-records 200
    python3 paginate.py --api openalex --query 'filter=publication_year:2024' --dry-run

Needs network access. No credentials required for bioRxiv, medRxiv, Europe PMC,
Crossref, or OpenAlex; NCBI_API_KEY and S2_API_KEY raise limits where relevant.
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable

sys.path.insert(0, str(Path(__file__).resolve().parent))

from _common import Reconciliation, emit, fail, redact_url  # noqa: E402

USER_AGENT = "paper-lookup-skill/2.0 (+https://agentskills.io)"
DEFAULT_MAX_RECORDS = 1000
DEFAULT_MAX_CALLS = 50
REQUEST_TIMEOUT = 60


@dataclass
class Page:
    """One response, normalized."""

    records: list[Any]
    total: int | None = None
    #: The next cursor/token/offset, or None when this API says it is done.
    next_state: Any = None
    #: Anything the caller must be told that is not a record.
    notes: list[str] | None = None


@dataclass
class Api:
    name: str
    #: Seconds to wait between requests. Serialized: never parallelize one host.
    delay: float
    build_url: Callable[[str, Any, int], str]
    parse: Callable[[Any, Any], Page]
    initial_state: Any = 0
    note: str = ""


def fetch(url: str, *, headers: dict[str, str] | None = None) -> Any:
    request = urllib.request.Request(url, headers={"User-Agent": USER_AGENT, **(headers or {})})
    try:
        with urllib.request.urlopen(request, timeout=REQUEST_TIMEOUT) as response:
            body = response.read().decode("utf-8", errors="replace")
    except urllib.error.HTTPError as error:
        detail = error.read().decode("utf-8", errors="replace")[:400]
        raise RuntimeError(f"HTTP {error.code} from {url}: {detail}") from error
    except urllib.error.URLError as error:
        raise RuntimeError(f"could not reach {url}: {error.reason}") from error
    try:
        return json.loads(body)
    except json.JSONDecodeError as error:
        raise RuntimeError(f"response from {url} was not JSON: {error}; first 200 bytes: {body[:200]}")


# --- bioRxiv / medRxiv ------------------------------------------------------
#
# `cursor` is an absolute record offset. The page size is 30 on /details/ and 100
# on /pubs/, and an out-of-step cursor is accepted with HTTP 200 -- so the step
# comes from the response's own `count`, never from a constant.


def _rxiv_url(server: str) -> Callable[[str, Any, int], str]:
    def build(query: str, state: Any, _limit: int) -> str:
        endpoint = "pubs" if query.startswith("pubs:") else "details"
        interval = query[5:] if query.startswith("pubs:") else query
        return f"https://api.biorxiv.org/{endpoint}/{server}/{interval}/{int(state)}/json"

    return build


def _rxiv_parse(payload: Any, state: Any) -> Page:
    if not isinstance(payload, dict):
        raise RuntimeError(f"expected a JSON object, got {type(payload).__name__}")

    messages = payload.get("messages") or [{}]
    message = messages[0] if isinstance(messages[0], dict) else {}
    status = message.get("status")
    records = payload.get("collection") or []
    notes: list[str] = []

    if status and status != "ok":
        # "no articles found" arrives with HTTP 200 and an empty collection, which
        # is indistinguishable from a genuine no-match unless status is read.
        notes.append(f"server status: {status!r} (HTTP 200 with an empty collection)")
        return Page(records=[], total=0, next_state=None, notes=notes)

    total = message.get("total")
    total = int(total) if total is not None and str(total).isdigit() else None

    new_papers = message.get("count_new_papers")
    if new_papers is not None:
        notes.append(
            f"count_new_papers={new_papers} counts distinct first-posting preprints while "
            f"total={total} counts every version record; deduplicate by DOI to compare against "
            "count_new_papers"
        )

    reported = message.get("count")
    step = int(reported) if isinstance(reported, int) and reported > 0 else len(records)
    if not records:
        return Page(records=[], total=total, next_state=None, notes=notes)

    if step != len(records):
        notes.append(f"response reported count={step} but returned {len(records)} records")
        step = len(records)

    next_state = int(state) + step
    if total is not None and next_state >= total:
        next_state = None
    return Page(records=records, total=total, next_state=next_state, notes=notes)


# --- Europe PMC ------------------------------------------------------------
#
# cursorMark. At exhaustion it returns an empty result list and echoes back the
# cursor you sent, rather than a null -- so detecting the end costs one extra
# empty request.


def _europepmc_url(query: str, state: Any, limit: int) -> str:
    params = {
        "query": query,
        "format": "json",
        "pageSize": str(min(limit, 1000)),
        "cursorMark": str(state),
        "resultType": "lite",
    }
    return "https://www.ebi.ac.uk/europepmc/webservices/rest/search?" + urllib.parse.urlencode(params)


def _europepmc_parse(payload: Any, state: Any) -> Page:
    if not isinstance(payload, dict):
        raise RuntimeError(f"expected a JSON object, got {type(payload).__name__}")
    # Europe PMC reports errors with HTTP 200 and an errCode in the body.
    if "errCode" in payload:
        raise RuntimeError(
            f"Europe PMC errCode {payload['errCode']}: {payload.get('errMsg', 'no message')}"
        )

    total = payload.get("hitCount")
    records = (payload.get("resultList") or {}).get("result") or []
    next_cursor = payload.get("nextCursorMark")
    notes: list[str] = []

    echoed = (payload.get("request") or {}).get("queryString")
    if echoed:
        notes.append(f"query as parsed by Europe PMC: {echoed!r}")

    if not records or next_cursor in (None, state):
        next_cursor = None
    return Page(
        records=records,
        total=int(total) if isinstance(total, int) else None,
        next_state=next_cursor,
        notes=notes,
    )


# --- OpenAlex --------------------------------------------------------------


def _openalex_url(query: str, state: Any, limit: int) -> str:
    # `query` is a raw parameter string, e.g. `search=crispr` or
    # `filter=publication_year:2024`, so both forms work without a second flag.
    base = "https://api.openalex.org/works?"
    params = {"per-page": str(min(limit, 200)), "cursor": str(state)}
    mail = os.environ.get("OPENALEX_EMAIL")
    if mail:
        params["mailto"] = mail
    key = os.environ.get("OPENALEX_API_KEY")
    if key:
        params["api_key"] = key
    return base + query + "&" + urllib.parse.urlencode(params)


def _openalex_parse(payload: Any, _state: Any) -> Page:
    if not isinstance(payload, dict):
        raise RuntimeError(f"expected a JSON object, got {type(payload).__name__}")
    meta = payload.get("meta") or {}
    records = payload.get("results") or []
    notes = []
    if meta.get("cost_usd") is not None:
        notes.append(f"OpenAlex reported cost_usd={meta['cost_usd']} for this call")
    next_cursor = meta.get("next_cursor")
    if not records:
        next_cursor = None
    return Page(
        records=records,
        total=meta.get("count") if isinstance(meta.get("count"), int) else None,
        next_state=next_cursor,
        notes=notes,
    )


# --- Crossref --------------------------------------------------------------


def _crossref_url(query: str, state: Any, limit: int) -> str:
    params = {"rows": str(min(limit, 1000)), "cursor": str(state)}
    mail = os.environ.get("CROSSREF_MAILTO")
    if mail:
        params["mailto"] = mail
    return "https://api.crossref.org/works?" + query + "&" + urllib.parse.urlencode(params)


def _crossref_parse(payload: Any, _state: Any) -> Page:
    if not isinstance(payload, dict):
        raise RuntimeError(f"expected a JSON object, got {type(payload).__name__}")
    message = payload.get("message") or {}
    records = message.get("items") or []
    next_cursor = message.get("next-cursor")
    if not records:
        next_cursor = None
    total = message.get("total-results")
    notes = ["Crossref cursors expire after 5 minutes; a long walk must keep moving"]
    return Page(
        records=records,
        total=int(total) if isinstance(total, int) else None,
        next_state=next_cursor,
        notes=notes,
    )


APIS: dict[str, Api] = {
    "biorxiv": Api(
        name="biorxiv",
        delay=1.0,
        build_url=_rxiv_url("biorxiv"),
        parse=_rxiv_parse,
        initial_state=0,
        note=(
            "query is an interval (2024-01-01/2024-01-03), Nd, N, or a DOI. "
            "Prefix with 'pubs:' to walk /pubs/ instead of /details/."
        ),
    ),
    "medrxiv": Api(
        name="medrxiv",
        delay=1.0,
        build_url=_rxiv_url("medrxiv"),
        parse=_rxiv_parse,
        initial_state=0,
        note="same as biorxiv; always via api.biorxiv.org, never api.medrxiv.org",
    ),
    "europepmc": Api(
        name="europepmc",
        delay=0.5,
        build_url=_europepmc_url,
        parse=_europepmc_parse,
        initial_state="*",
        note="query is Europe PMC query syntax, e.g. 'SRC:\"PPR\" AND \"organoid\"'",
    ),
    "openalex": Api(
        name="openalex",
        delay=0.2,
        build_url=_openalex_url,
        parse=_openalex_parse,
        initial_state="*",
        note="query is a raw parameter string, e.g. 'search=crispr' or 'filter=publication_year:2024'",
    ),
    "crossref": Api(
        name="crossref",
        delay=0.3,
        build_url=_crossref_url,
        parse=_crossref_parse,
        initial_state="*",
        note="query is a raw parameter string, e.g. 'query.bibliographic=attention+is+all+you+need'",
    ),
}


def walk(
    api: Api,
    query: str,
    *,
    page_size: int,
    max_records: int,
    max_calls: int,
    verbose: bool,
) -> tuple[list[Any], Reconciliation, list[str]]:
    records: list[Any] = []
    reconciliation = Reconciliation()
    urls: list[str] = []
    state = api.initial_state
    seen_notes: set[str] = set()

    while True:
        if reconciliation.pages >= max_calls:
            reconciliation.stopped_at_limit = True
            reconciliation.note(
                f"stopped at the --max-calls limit of {max_calls}; the walk is INCOMPLETE"
            )
            break
        if len(records) >= max_records:
            reconciliation.stopped_at_limit = True
            reconciliation.note(
                f"stopped at the --max-records limit of {max_records}; the walk is INCOMPLETE"
            )
            break

        remaining = max_records - len(records)
        url = api.build_url(query, state, min(page_size, remaining))
        # Record and log the redacted form only. OpenAlex and Crossref authenticate
        # by query string, so the fetched URL carries the credential and this
        # provenance list is printed to the user.
        safe_url = redact_url(url)
        urls.append(safe_url)
        if verbose:
            sys.stderr.write(f"  page {reconciliation.pages + 1}: {safe_url}\n")

        try:
            page = api.parse(fetch(url), state)
        except RuntimeError as error:
            fail(str(error))

        reconciliation.pages += 1
        if page.total is not None and reconciliation.expected is None:
            reconciliation.expected = page.total
        for note in page.notes or []:
            if note not in seen_notes:
                seen_notes.add(note)
                reconciliation.note(note)

        records.extend(page.records)
        if page.next_state is None:
            break
        state = page.next_state
        time.sleep(api.delay)

    # Trim only after the walk, so the reported page count stays truthful.
    if len(records) > max_records:
        reconciliation.note(
            f"last page overshot --max-records; kept the first {max_records} of {len(records)}"
        )
        records = records[:max_records]

    reconciliation.retrieved = len(records)
    return records, reconciliation, urls


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description=(
            "Paginate one of this skill's APIs with the correct step, the correct stop "
            "condition, and count reconciliation. Exits 4 on a reconciliation shortfall."
        ),
        epilog="python3 %(prog)s --api biorxiv --query 2024-01-01/2024-01-03",
    )
    # Not `required=True`: --list-apis is the flag you reach for when you do not yet
    # know what to pass for either of these.
    parser.add_argument("--api", choices=sorted(APIS), help="which API to walk")
    parser.add_argument("--query", help="see --list-apis for the per-API format")
    parser.add_argument("--page-size", type=int, default=100, help="requested page size (default 100)")
    parser.add_argument(
        "--max-records",
        type=int,
        default=DEFAULT_MAX_RECORDS,
        help=f"stop after this many records (default {DEFAULT_MAX_RECORDS})",
    )
    parser.add_argument(
        "--max-calls",
        type=int,
        default=DEFAULT_MAX_CALLS,
        help=f"stop after this many requests (default {DEFAULT_MAX_CALLS})",
    )
    parser.add_argument("-o", "--output", help="write JSON here instead of stdout")
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="print the first URL that would be requested and exit without fetching",
    )
    parser.add_argument("--list-apis", action="store_true", help="describe each API's query format")
    parser.add_argument("-v", "--verbose", action="store_true", help="log each URL to stderr")
    return parser


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)

    if args.list_apis:
        emit(
            {
                name: {"delay_seconds": api.delay, "query_format": api.note}
                for name, api in sorted(APIS.items())
            },
            args.output,
        )
        return 0

    if not args.api or not args.query:
        fail("--api and --query are both required (use --list-apis to see the query format)", code=2)
    if args.page_size < 1:
        fail("--page-size must be at least 1")
    if args.max_records < 1:
        fail("--max-records must be at least 1")
    if args.max_calls < 1:
        fail("--max-calls must be at least 1")

    api = APIS[args.api]

    if args.dry_run:
        emit(
            {
                "api": api.name,
                "first_url": redact_url(
                    api.build_url(args.query, api.initial_state, args.page_size)
                ),
                "delay_seconds": api.delay,
                "query_format": api.note,
            },
            args.output,
        )
        return 0

    records, reconciliation, urls = walk(
        api,
        args.query,
        page_size=args.page_size,
        max_records=args.max_records,
        max_calls=args.max_calls,
        verbose=args.verbose,
    )

    emit(
        {
            "api": api.name,
            "query": args.query,
            "provenance": {"urls": urls, "delay_seconds": api.delay},
            "reconciliation": reconciliation.as_dict(),
            "records": records,
        },
        args.output,
    )

    if not reconciliation.ok:
        # Exit 4 is reserved for the unexplained case: the walk terminated on its
        # own and still came up short, which means records went missing. A bound
        # the caller set is not a failure and exits 0 with the partiality recorded.
        fail(
            f"reconciliation failed: the walk ended on its own but retrieved "
            f"{reconciliation.retrieved} of {reconciliation.expected}. Records are missing -- "
            "say so before drawing any conclusion from this result.",
            code=4,
        )
    if reconciliation.stopped_at_limit:
        sys.stderr.write(
            f"note: stopped at a caller-set bound with {reconciliation.retrieved}"
            f"{f' of {reconciliation.expected}' if reconciliation.expected is not None else ''} "
            "records. Report this result as partial.\n"
        )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())