"""Shared helpers for the two OpenAlex snapshot fetchers.

Sources and institutions both pull from the OpenAlex S3 bucket, translate
``s3://`` URLs to the public HTTPS gateway, and defend-in-depth against
a compromised manifest by allowlisting the ``s3://openalex/`` prefix.
This file owns those three shared symbols so ``openalex.py`` and
``institutions.py`` don't duplicate them (and can't drift). It also
owns the per-partition streaming helper they both use to iterate
records with consistent malformed-line suppression and tmp-file
lifecycle.
"""

from __future__ import annotations

import gzip
import json
from pathlib import Path
from typing import Callable, Iterator, Tuple

from loguru import logger

# Public OpenAlex snapshot — CC0, no auth, no rate limits.
# Manifest format documented at:
#   https://docs.openalex.org/download-all-data/snapshot-data-format
# Each entry in ``manifest["entries"]`` has ``url`` (s3://...) and
# ``meta.content_length`` / ``meta.record_count``. We translate s3:// to
# the public HTTPS gateway so we don't need boto3.
OPENALEX_S3_BASE = "https://openalex.s3.amazonaws.com"

# Only fetch parts hosted under the OpenAlex public S3 bucket — defense
# in depth on top of safe_get's private-IP block. A compromised or
# malformed manifest could otherwise list arbitrary attacker-controlled
# URLs.
OPENALEX_MANIFEST_ALLOWED_PREFIX = "s3://openalex/"


def s3_to_https(s3_url: str) -> str:
    """Translate ``s3://openalex/...`` to the public HTTPS gateway."""
    return s3_url.replace(
        OPENALEX_MANIFEST_ALLOWED_PREFIX, OPENALEX_S3_BASE + "/", 1
    )


def validate_manifest_entries(entries: list[dict], label: str) -> None:
    """Refuse to fetch if any manifest entry escapes the S3 allowlist.

    Defense-in-depth: a compromised or tampered manifest could list
    URLs outside the OpenAlex bucket. Refusing the whole fetch rather
    than fetching some-and-not-others keeps failure modes simple.
    """
    for entry in entries:
        raw = entry.get("url", "")
        if not raw.startswith(OPENALEX_MANIFEST_ALLOWED_PREFIX):
            raise ValueError(
                f"{label} manifest contains disallowed URL "
                f"(must start with {OPENALEX_MANIFEST_ALLOWED_PREFIX!r}): "
                f"{raw!r}"
            )


# Per-partition retry budget. The default ``safe_get_with_retries``
# budget (3 retries, 1-2-4 s backoff = ~7 s total) is sized for small
# request bodies and trips on a sustained mid-stream S3 hiccup: every
# retry of a ~5–10 MB partition that lands inside the same bad window
# fails the same way, exhausts the budget in seconds, and aborts the
# whole 30-partition pull. The release-gate workflow saw this twice in
# a row on 2026-04-26.
#
# 5 retries with 2-5-10-20-40 s backoff rides out a ~75 s S3 blip
# instead, while still bounding total wall-clock per partition at
# roughly ``timeout * 6 + 77 s`` — well inside the 45 min job timeout
# even if every partition needed all retries.
_PARTITION_MAX_RETRIES = 5
_PARTITION_BACKOFF_SECONDS = (2, 5, 10, 20, 40)


def iter_partitions(
    entries: list[dict],
    data_dir: Path,
    *,
    file_prefix: str,
    label: str,
    safe_get: Callable,
    timeout: int = 120,
    max_retries: int = _PARTITION_MAX_RETRIES,
    backoff_times: tuple = _PARTITION_BACKOFF_SECONDS,
) -> Iterator[Tuple[int, int, list[dict]]]:
    """Download each partition, yielding ``(idx, total_parts, records)``.

    Shared between ``openalex.py`` and ``institutions.py`` so the
    tmp-file lifecycle and malformed-JSON suppression (first-10
    warnings + one "further suppressed" notice) are defined once.

    The caller iterates ``records`` for per-record work and is
    responsible for per-partition progress logging and ``progress_cb``
    invocations — those need caller-specific state (running record
    count, schema-drift counters) that doesn't belong in the helper.

    Args:
        entries: ``manifest["entries"]`` — each dict has ``url``
            starting with ``s3://openalex/``.
        data_dir: Directory used for the transient ``.<prefix>_part_<n>.gz``
            files. Cleaned up even on exception.
        file_prefix: Leaf prefix for tmp files
            (e.g. ``openalex_sources`` / ``openalex_institutions``).
        label: Human-readable label used in log messages
            (e.g. ``"OpenAlex sources"`` / ``"Institutions"``).
        safe_get: Dependency-injected HTTP getter (lets the caller
            pick ``safe_get_with_retries`` without forcing a global
            import at module load). Must accept ``consume_body=True``
            so body-stream transients (``ChunkedEncodingError``,
            ``ReadTimeout``) raised during ``resp.content`` are
            retried inside the wrapper, not propagated to abort the
            whole multi-partition pull.
        timeout: Per-partition HTTP timeout (seconds).
        max_retries: Per-partition retry budget. Defaults higher than
            ``safe_get_with_retries``' generic 3 because partition
            bodies are MB-sized and a mid-stream IncompleteRead aborts
            the whole multi-partition pull on exhaustion.
        backoff_times: Per-attempt sleep schedule. Defaults to a
            longer schedule than the generic ``safe_get_with_retries``
            (1, 2, 4) so we ride out a sustained S3 blip instead of
            burning all retries inside the same bad window.
    """
    malformed_total = 0
    total_parts = len(entries)

    for idx, entry in enumerate(entries):
        part_url = s3_to_https(entry["url"])
        tmp_part = data_dir / f".{file_prefix}_part_{idx}.gz"
        records: list[dict] = []

        try:
            # consume_body=True: an OpenAlex S3 partition is ~10 MB
            # gzipped. A mid-stream ChunkedEncodingError /
            # IncompleteRead would otherwise abort the whole 30+
            # partition pull. With consume_body, safe_get_with_retries
            # reads resp.content inside its retry loop and retries
            # body-stream transients the same way it retries
            # header-stage failures.
            resp = safe_get(
                part_url,
                timeout=timeout,
                consume_body=True,
                max_retries=max_retries,
                backoff_times=backoff_times,
            )
            resp.raise_for_status()
            tmp_part.write_bytes(resp.content)

            with gzip.open(tmp_part, "rt", encoding="utf-8") as fh:
                for line in fh:
                    line = line.strip()
                    if not line:
                        continue
                    try:
                        rec = json.loads(line)
                    except (json.JSONDecodeError, ValueError):
                        malformed_total += 1
                        if malformed_total <= 10:
                            logger.warning(
                                f"{label} partition {idx}: skipping "
                                f"malformed JSON line"
                            )
                        elif malformed_total == 11:
                            logger.warning(
                                f"{label} partition {idx}: further "
                                "malformed lines suppressed"
                            )
                        continue
                    records.append(rec)
        finally:
            tmp_part.unlink(missing_ok=True)

        yield idx, total_parts, records

    if malformed_total:
        logger.warning(
            f"{label}: {malformed_total:,} malformed lines skipped across "
            f"{total_parts} partitions"
        )
