arctic_shift.py

"""Arctic Shift API client — primary Reddit archive source.

Arctic Shift (arctic-shift.photon-reddit.com) archives Reddit data up to
the present day, making it the preferred source over PullPush for all
subreddit pulls. PullPush is used only as a fallback when Arctic Shift
is unavailable or rate-limited.

Pagination uses the `before` cursor: results are returned DESC by
created_utc, and each page sets before=(last item's created_utc - 1).
The `after` param filters out items older than since_ts.
"""
from __future__ import annotations

import logging
import time
from typing import Iterator, Optional

import requests

ARCTIC_SHIFT_BASE = "https://arctic-shift.photon-reddit.com/api"
DEFAULT_TIMEOUT = 30
USER_AGENT = "pakka-corpus/0.1 (research)"

log = logging.getLogger(__name__)


class ArcticShiftError(Exception):
    pass


class ArcticShiftClient:
    def __init__(
        self,
        delay_seconds: float = 1.0,
        max_retries: int = 3,
        backoff_base: int = 2,
    ):
        self.delay = delay_seconds
        self.max_retries = max_retries
        self.backoff_base = backoff_base
        self.session = requests.Session()
        self.session.headers.update({"User-Agent": USER_AGENT})

    def _get(self, path: str, params: dict) -> dict:
        url = f"{ARCTIC_SHIFT_BASE}{path}"
        last_err = None
        for attempt in range(self.max_retries):
            try:
                r = self.session.get(url, params=params, timeout=DEFAULT_TIMEOUT)
                if r.status_code == 200:
                    try:
                        return r.json()
                    except ValueError as e:
                        raise ArcticShiftError(f"Bad JSON: {e}") from e
                if r.status_code == 429 or r.status_code >= 500:
                    wait = self.backoff_base**attempt
                    log.warning(
                        "ArcticShift %s on %s; backoff %ss (attempt %d)",
                        r.status_code, path, wait, attempt + 1,
                    )
                    time.sleep(wait)
                    last_err = f"HTTP {r.status_code}"
                    continue
                raise ArcticShiftError(f"HTTP {r.status_code}: {r.text[:200]}")
            except requests.RequestException as e:
                wait = self.backoff_base**attempt
                log.warning(
                    "ArcticShift request error on %s: %s; backoff %ss", path, e, wait
                )
                time.sleep(wait)
                last_err = str(e)
        raise ArcticShiftError(f"Exhausted retries on {path}: {last_err}")

    # ----------------------------------------------------------- bulk fetch

    def fetch_posts_by_ids(self, ids: list[str]) -> list[dict]:
        """Fetch up to 100 specific posts by ID list."""
        params = {"ids": ",".join(ids), "limit": len(ids)}
        data = self._get("/posts/search", params)
        return data.get("data", []) or []

    # ---------------------------------------------------------------- search

    def search_submissions(
        self,
        subreddit: Optional[str] = None,
        author: Optional[str] = None,
        after: Optional[int] = None,
        before: Optional[int] = None,
        limit: int = 100,
    ) -> list[dict]:
        params: dict = {"limit": limit}
        if subreddit:
            params["subreddit"] = subreddit
        if author:
            params["author"] = author
        if after:
            params["after"] = after
        if before:
            params["before"] = before
        data = self._get("/posts/search", params)
        return data.get("data", []) or []

    def search_comments(
        self,
        subreddit: Optional[str] = None,
        link_id: Optional[str] = None,
        author: Optional[str] = None,
        after: Optional[int] = None,
        before: Optional[int] = None,
        limit: int = 100,
    ) -> list[dict]:
        params: dict = {"limit": limit}
        if subreddit:
            params["subreddit"] = subreddit
        if link_id:
            params["link_id"] = link_id.lstrip("t3_") if link_id.startswith("t3_") else link_id
        if author:
            params["author"] = author
        if after:
            params["after"] = after
        if before:
            params["before"] = before
        data = self._get("/comments/search", params)
        return data.get("data", []) or []

    # -------------------------------------------------------------- pagination

    def paginate_submissions(
        self,
        subreddit: str,
        since_ts: Optional[int] = None,
        max_items: int = 5000,
    ) -> Iterator[dict]:
        before: Optional[int] = None
        yielded = 0
        seen: set[str] = set()

        while yielded < max_items:
            batch = self.search_submissions(
                subreddit=subreddit,
                after=since_ts,
                before=before,
                limit=100,
            )
            if not batch:
                return

            new_in_batch = 0
            for item in batch:
                iid = item.get("id", "")
                if iid in seen:
                    continue
                seen.add(iid)
                new_in_batch += 1
                created = int(float(item.get("created_utc", 0)))
                if since_ts is not None and created < since_ts:
                    return
                yield item
                yielded += 1
                if yielded >= max_items:
                    return

            if new_in_batch == 0:
                return

            oldest_ts = int(float(batch[-1].get("created_utc", 0)))
            before = oldest_ts - 1
            time.sleep(self.delay)

    def paginate_comments_for_subreddit(
        self,
        subreddit: str,
        since_ts: Optional[int] = None,
        max_items: int = 10000,
    ) -> Iterator[dict]:
        before: Optional[int] = None
        yielded = 0
        seen: set[str] = set()

        while yielded < max_items:
            batch = self.search_comments(
                subreddit=subreddit,
                after=since_ts,
                before=before,
                limit=100,
            )
            if not batch:
                return

            new_in_batch = 0
            for item in batch:
                iid = item.get("id", "")
                if iid in seen:
                    continue
                seen.add(iid)
                new_in_batch += 1
                created = int(float(item.get("created_utc", 0)))
                if since_ts is not None and created < since_ts:
                    return
                yield item
                yielded += 1
                if yielded >= max_items:
                    return

            if new_in_batch == 0:
                return

            oldest_ts = int(float(batch[-1].get("created_utc", 0)))
            before = oldest_ts - 1
            time.sleep(self.delay)

    def paginate_comments_for_post(
        self,
        post_id: str,
        max_items: int = 20000,
    ) -> Iterator[dict]:
        before: Optional[int] = None
        yielded = 0
        seen: set[str] = set()
        limit = 100

        while yielded < max_items:
            batch = self.search_comments(
                link_id=post_id,
                before=before,
                limit=limit,
            )
            if not batch:
                return

            new_in_batch = 0
            for item in batch:
                iid = item.get("id", "")
                if iid in seen:
                    continue
                seen.add(iid)
                new_in_batch += 1
                yield item
                yielded += 1
                if yielded >= max_items:
                    return

            if new_in_batch == 0:
                return

            # Fewer items than requested → no more pages, skip the extra round-trip
            if len(batch) < limit:
                return

            oldest_ts = int(float(batch[-1].get("created_utc", 0)))
            before = oldest_ts - 1
            time.sleep(self.delay)