wiki-kb/wikikb/scraper.py

"""Mirror a wiki into `corpus/`.

Where pages come from is decided by a :mod:`.sources` Source — a MediaWiki API when
one exists, a sitemap + HTML extraction when it does not. This module owns the parts
that are the same either way: politeness (User-Agent, ~1 req/s, exponential backoff),
noise filtering, resumability, and writing the corpus + manifest.
"""

from __future__ import annotations

import argparse
import json
import re
import sys
import time
from collections.abc import Callable

import httpx

from . import config, variant
from .html2md import build_document
from .sources import HTML, MediaWikiSource, SitemapHtmlSource, Source

_UNSAFE_CHARS = re.compile(r'[/\\:*?"<>|\x00-\x1f]')


class Throttle:
    """Single-flight rate limiter maintaining a minimum request interval."""

    def __init__(
        self,
        interval: float = config.REQUEST_INTERVAL_S,
        *,
        sleep: Callable[[float], None] = time.sleep,
        clock: Callable[[], float] = time.monotonic,
    ) -> None:
        self._interval = interval
        self._sleep = sleep
        self._clock = clock
        self._last = 0.0

    def wait(self) -> None:
        now = self._clock()
        delay = self._interval - (now - self._last)
        if delay > 0:
            self._sleep(delay)
        self._last = self._clock()


def is_core_title(title: str) -> bool:
    """True for content pages worth mirroring; False for noise titles.

    Whether a "/" in a title means noise is per-wiki and configured, never assumed —
    see ``config.SUBPAGE_POLICY``.
    """
    if "/" in title:
        if title.rsplit("/", 1)[-1].lower() in config.NOISE_SUBPAGE_TAILS:
            return False
        if config.SUBPAGE_POLICY == "reject":
            return False
    lowered = title.lower()
    return not any(lowered.startswith(prefix) for prefix in config.MOD_TITLE_PREFIXES)


def sanitize_filename(title: str) -> str:
    """Map a page title to a readable, filesystem-safe ``.md`` filename."""
    name = _UNSAFE_CHARS.sub("_", title).replace(" ", "_").strip("._")
    return f"{name or 'untitled'}.md"


def _api_get(
    client: httpx.Client,
    params: dict[str, object],
    *,
    throttle: Throttle,
    sleep: Callable[[float], None] = time.sleep,
) -> dict:
    """GET the API with maxlag, retrying on 429/503 and maxlag errors."""
    query = {
        **params,
        "format": "json",
        "formatversion": "2",
        "maxlag": config.MAXLAG,
    }
    for attempt in range(config.MAX_RETRIES):
        throttle.wait()
        response = client.get(config.API_URL, params=query)
        if response.status_code in (429, 503):
            sleep(config.BACKOFF_BASE_S**attempt)
            continue
        response.raise_for_status()
        try:
            data = response.json()
        except ValueError:
            # Wikis behind Cloudflare (Fandom, many wiki.gg installs) serve a bot
            # challenge as a 200 with an HTML body, so raise_for_status() passes it
            # straight through. Back off and retry rather than letting a challenge
            # page be converted to markdown and written into the corpus.
            sleep(config.BACKOFF_BASE_S**attempt)
            continue
        error = data.get("error")
        if error and error.get("code") == "maxlag":
            sleep(config.BACKOFF_BASE_S**attempt)
            continue
        return data
    raise RuntimeError(f"API request failed after {config.MAX_RETRIES} retries: {params}")


def build_source(
    client: httpx.Client,
    *,
    throttle: Throttle,
    kind: str | None = None,
) -> Source:
    """Construct the page source for the configured wiki.

    Defaults to the MediaWiki API, which is preferred whenever it exists: it gives a
    real page list, expanded templates, categories and revision ids. ``kind="html"``
    selects the sitemap fallback for wikis with no API.
    """
    if (kind or config.SOURCE_KIND) == HTML:
        return SitemapHtmlSource(
            lambda url: _get(client, url, throttle=throttle),
            sitemap_url=config.SITEMAP_URL,
        )
    return MediaWikiSource(lambda params: _api_get(client, params, throttle=throttle))


def _get(
    client: httpx.Client,
    url: str,
    *,
    throttle: Throttle,
    sleep: Callable[[float], None] = time.sleep,
) -> httpx.Response:
    """Throttled GET with the same backoff policy as the API path.

    404 is returned to the caller rather than retried: for the sitemap source a
    missing page is a normal outcome, not a transient failure.
    """
    for attempt in range(config.MAX_RETRIES):
        throttle.wait()
        response = client.get(url, follow_redirects=True)
        if response.status_code in (429, 503):
            sleep(config.BACKOFF_BASE_S**attempt)
            continue
        if response.status_code == 404:
            return response
        response.raise_for_status()
        return response
    raise RuntimeError(f"GET failed after {config.MAX_RETRIES} retries: {url}")


def load_manifest() -> dict[str, dict]:
    if config.MANIFEST_PATH.exists():
        return json.loads(config.MANIFEST_PATH.read_text(encoding="utf-8"))
    return {}


def save_manifest(manifest: dict[str, dict]) -> None:
    config.MANIFEST_PATH.write_text(
        json.dumps(manifest, indent=2, sort_keys=True, ensure_ascii=False),
        encoding="utf-8",
    )


def scrape(
    *,
    limit: int | None = None,
    refresh: bool = False,
    titles: list[str] | None = None,
    client: httpx.Client | None = None,
    throttle: Throttle | None = None,
    source: Source | None = None,
    log: Callable[[str], None] = print,
) -> dict[str, int]:
    """Mirror the wiki. Returns counts of {written, skipped, missing}."""
    owns_client = client is None
    client = client or httpx.Client(
        headers={"User-Agent": config.USER_AGENT},
        timeout=config.REQUEST_TIMEOUT_S,
    )
    throttle = throttle or Throttle()
    source = source or build_source(client, throttle=throttle)
    selected = (
        (t for t in titles if is_core_title(t))
        if titles is not None
        else source.iter_titles(limit=limit)
    )
    config.CORPUS_DIR.mkdir(parents=True, exist_ok=True)
    manifest = load_manifest()
    written = skipped = missing = 0
    try:
        for title in selected:
            if titles is None and not is_core_title(title):
                continue
            filename = sanitize_filename(title)
            target = config.CORPUS_DIR / filename
            if not refresh and title in manifest and target.exists():
                skipped += 1
                continue
            page = source.fetch_page(title)
            if page is None:
                missing += 1
                log(f"MISSING {title}")
                continue
            page_variant = variant.classify(page.title, page.categories)
            document = build_document(
                variant=page_variant,
                title=page.title,
                url=page.url,
                revid=page.revid,
                categories=page.categories,
                source=page.source,
                html=page.html,
            )
            target.write_text(document, encoding="utf-8")
            manifest[page.title] = {
                "revid": page.revid,
                "filename": filename,
                "url": page.url,
                "categories": page.categories,
                "source": page.source,
                "variant": page_variant,
            }
            written += 1
            log(f"[{written:>5}] {page.title}")
            if written % 25 == 0:
                save_manifest(manifest)
    finally:
        save_manifest(manifest)
        if owns_client:
            client.close()
    return {"written": written, "skipped": skipped, "missing": missing}


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Mirror the CK3 wiki into corpus/.")
    parser.add_argument("--limit", type=int, default=None, help="max pages to fetch")
    parser.add_argument("--refresh", action="store_true", help="re-fetch pages already mirrored")
    parser.add_argument(
        "--titles", default=None, help="comma-separated titles to mirror instead of the whole wiki"
    )
    args = parser.parse_args(argv)
    named = [t.strip() for t in args.titles.split(",") if t.strip()] if args.titles else None
    result = scrape(limit=args.limit, refresh=args.refresh, titles=named)
    print(
        f"\nDone: {result['written']} written, "
        f"{result['skipped']} skipped, {result['missing']} missing."
    )
    return 0


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