"""OpenAlex Institutions data source.

Downloads the **bulk snapshot** of ~120K institutions from the public
OpenAlex S3 gateway (``openalex.s3.amazonaws.com``) and writes a
compact gzipped JSON snapshot used by the institution-scoring tier of
the journal filter. Each institution carries its ROR ID, country,
h-index, works count, and 2-year mean citedness.

Why bulk snapshot instead of the REST API
-----------------------------------------

The previous implementation cursor-paginated ``/api/institutions`` at
200/page with 100 ms polite sleeps — ~550 requests + ~55 s of sleep +
actual transfer. Wall-clock was 5–10 minutes per "Download Data"
click, which dominated the user-facing latency. The S3 dump is the
documented bulk path, CC0, no auth, no rate limits, and finishes
much faster.
"""

from __future__ import annotations

import gzip
import json
import time
from pathlib import Path

from loguru import logger

from ._openalex_common import (
    OPENALEX_S3_BASE,
    iter_partitions,
    validate_manifest_entries,
)
from ..scoring import normalize_name
from .base import DataSource

_OPENALEX_INSTITUTIONS_MANIFEST = (
    f"{OPENALEX_S3_BASE}/data/jsonl/institutions/manifest.json"
)

# Safety floor — OpenAlex has ~120K institutions. Refuse to overwrite
# good data if the fetch produced far fewer records (likely schema
# change, empty partitions, or broken manifest).
_MIN_INSTITUTIONS = 50_000


class InstitutionSource(DataSource):
    key = "institutions"  # gitleaks:allow
    name = "OpenAlex Institutions"
    url = "https://openalex.org"
    dataset_url = (
        "https://docs.openalex.org/download-all-data/openalex-snapshot"
    )
    license = "CC0 1.0"
    license_url = "https://creativecommons.org/publicdomain/zero/1.0/"
    description = (
        "~120K research institutions with ROR ID, country, h-index, "
        "works count, and citation counts"
    )
    filename = "openalex_institutions.json.gz"
    count_label = "institutions"
    auto_download = False  # large; user opts in via dashboard
    required = False
    approx_size_mb = 10.0  # final compact output, NOT the raw snapshot

    def fetch(self, data_dir: Path, progress_cb=None) -> int:
        from ...security.safe_requests import (
            safe_get_with_retries as safe_get,
        )

        # 1. Fetch the manifest.
        logger.info(
            f"Fetching OpenAlex institutions manifest: "
            f"{_OPENALEX_INSTITUTIONS_MANIFEST}"
        )
        # consume_body=True: small JSON but serial bottleneck — a body
        # transient here aborts the whole 10-min institutions pull.
        manifest_resp = safe_get(
            _OPENALEX_INSTITUTIONS_MANIFEST, timeout=30, consume_body=True
        )
        manifest_resp.raise_for_status()
        manifest = manifest_resp.json()

        # OpenAlex's 2026-06 "standard-format" snapshot renamed the
        # manifest's part list from ``entries`` to ``files`` (see
        # ``openalex.py``). Entry shape (url + meta) is unchanged.
        entries = manifest.get("files", [])

        # Validate every part URL before fetching any part — SSRF
        # defense in depth. If any entry points outside the OpenAlex
        # bucket we refuse the whole fetch rather than partially trust.
        validate_manifest_entries(entries, "Institutions")

        total_records = sum(
            e.get("meta", {}).get("record_count", 0) for e in entries
        )
        total_bytes = sum(
            e.get("meta", {}).get("content_length", 0) for e in entries
        )
        logger.info(
            f"OpenAlex institutions snapshot: {len(entries)} parts, "
            f"{total_records:,} records, "
            f"{total_bytes / 1024 / 1024:.0f} MB compressed"
        )

        # 2. Stream-process each part. The ``iter_partitions`` helper
        #    in ``_openalex_common`` owns the tmp-file lifecycle and
        #    first-10 malformed-JSON suppression; we focus on the
        #    compact-format + secondary-index extraction. Compact
        #    format matches the journal sources snapshot; ROR and
        #    name indexes keep the runtime lookup path O(1).
        institutions: dict = {}
        ror_index: dict = {}
        name_index: dict = {}
        start = time.time()

        for idx, total_parts, records in iter_partitions(
            entries,
            data_dir,
            file_prefix="openalex_institutions",
            label="Institutions",
            safe_get=safe_get,
        ):
            for inst in records:
                inst_id = (inst.get("id") or "").split("/")[-1]
                if not inst_id:
                    continue

                stats = inst.get("summary_stats") or {}
                ror = (inst.get("ror") or "").rstrip("/").split("/")[-1]
                compact = {
                    "n": inst.get("display_name", ""),
                    "c": inst.get("country_code", ""),
                    "t": inst.get("type", ""),
                    "h": stats.get("h_index"),
                    "if": stats.get("2yr_mean_citedness"),
                    "w": inst.get("works_count"),
                    "cb": inst.get("cited_by_count"),
                    "r": ror or None,
                }
                institutions[inst_id] = compact

                if ror:
                    ror_index[ror] = inst_id

                primary = normalize_name(inst.get("display_name", "") or "")
                if primary:
                    name_index[primary] = inst_id
                for alt in inst.get("display_name_alternatives") or []:
                    alt_lower = normalize_name(alt or "")
                    if alt_lower and alt_lower not in name_index:
                        name_index[alt_lower] = inst_id

            if (idx + 1) % 5 == 0 or idx == total_parts - 1:
                elapsed = time.time() - start
                logger.info(
                    f"OpenAlex institutions: processed "
                    f"{idx + 1}/{total_parts} parts "
                    f"({len(institutions):,} records, {elapsed:.0f}s)"
                )
            # Per-partition UI ping so the dashboard bar moves
            # frequently enough to feel live (vs every 5th partition).
            if progress_cb is not None:
                try:
                    progress_cb(
                        idx + 1,
                        total_parts,
                        f"{len(institutions):,} records",
                    )
                except Exception:
                    logger.debug(
                        "institutions progress_cb raised; continuing",
                        exc_info=True,
                    )

        if len(institutions) < _MIN_INSTITUTIONS:
            raise RuntimeError(
                f"OpenAlex institutions: suspiciously few records "
                f"({len(institutions):,} < {_MIN_INSTITUTIONS:,}); "
                "refusing to overwrite existing data"
            )

        payload = {
            "i": institutions,  # institution ID → compact record
            "r": ror_index,  # ROR ID → institution ID
            "nm": name_index,  # lower-cased name → institution ID
        }

        output = data_dir / self.filename
        tmp = data_dir / f"{self.filename}.tmp"
        with gzip.open(tmp, "wt", encoding="utf-8") as f:
            json.dump(payload, f)
        tmp.rename(output)

        elapsed = time.time() - start
        logger.info(
            f"OpenAlex institutions: saved {len(institutions):,} "
            f"institutions in {elapsed:.0f}s"
        )
        return len(institutions)
