"""
RBSP Cross-Encoder Reranking Pipeline.

Two-stage retrieval: a fast first stage (BM25 / LSH / resonance) retrieves
candidates, then a cross-encoder reranker scores each (query, document) pair
for precise relevance.  This is the standard pattern used by all SOTA search
systems.

Components:
    - RerankerConfig:         Configuration for the reranking pipeline.
    - CrossEncoder:           Abstract base for scoring (query, document) pairs.
    - DummyCrossEncoder:      Zero-dep test encoder using BM25-like heuristic.
    - CohereReranker:         Optional Cohere Rerank API provider.
    - TransformerReranker:    Optional HuggingFace cross-encoder provider.
    - RerankerPipeline:       Orchestrates first-stage results through a CrossEncoder.

Zero required external dependencies.
"""

from __future__ import annotations

import heapq
import math
import re
from dataclasses import dataclass
from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple

Ranking = List[Tuple[str, float]]


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

_SIMPLE_TOKEN_RE = re.compile(r'[a-zA-Z0-9_]+')


def _simple_tokenize(text: str) -> list[str]:
    """Lowercase tokenize text into words.

    No stemming -- keeps the encoder independent of rbsp.core.stemmer.
    """
    return [m.group().lower() for m in _SIMPLE_TOKEN_RE.finditer(text)]


def _minmax_normalize(scores: list[float]) -> list[float]:
    """Min-max normalize a list of scores to [0, 1]."""
    if not scores:
        return []
    lo = min(scores)
    hi = max(scores)
    span = hi - lo
    if span == 0.0:
        return [0.0] * len(scores)
    return [(s - lo) / span for s in scores]


# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------


@dataclass(slots=True)
class RerankerConfig:
    """Configuration for the cross-encoder reranking pipeline.

    Attributes:
        top_n:           Number of first-stage candidates to rerank (default 50).
                         Higher = better recall, slower.  Should be >= final top_k.
        batch_size:      Number of (query, doc) pairs per scoring call (default 32).
                         Only relevant for batched encoders (Cohere, Transformers).
        score_threshold: Minimum reranker score to keep a result (default 0.0).
                         Results below this are discarded after reranking.
        normalize:       Whether to min-max normalize reranker scores to [0, 1]
                         (default True).  Useful for downstream fusion.
    """

    top_n: int = 50
    batch_size: int = 32
    score_threshold: float = 0.0
    normalize: bool = True


# ---------------------------------------------------------------------------
# Abstract cross-encoder
# ---------------------------------------------------------------------------


class CrossEncoder:
    """Abstract base class for cross-encoder scoring.

    A cross-encoder takes a (query, document_text) pair and produces a
    relevance score.  Higher scores = more relevant.

    Subclasses MUST implement :meth:`score_batch`.  The default
    :meth:`score` delegates to ``score_batch`` with a single-element list.
    """

    def score(self, query: str, document: str) -> float:
        """Score a single (query, document) pair.

        Args:
            query:    The search query string.
            document: The document text.

        Returns:
            Relevance score (higher = more relevant).
        """
        results = self.score_batch(query, [document])
        return results[0]

    def score_batch(self, query: str, documents: list[str]) -> list[float]:
        """Score a batch of documents against a single query.

        Args:
            query:     The search query string.
            documents: List of document texts.

        Returns:
            List of relevance scores, one per document, same order.
        """
        raise NotImplementedError


# ---------------------------------------------------------------------------
# Dummy cross-encoder (zero-dep test encoder)
# ---------------------------------------------------------------------------


class DummyCrossEncoder(CrossEncoder):
    """A zero-dependency cross-encoder for testing and baselines.

    Computes a BM25-like relevance score using simple term frequency and
    inverse document frequency heuristics.  No external models needed.

    The score formula per query term present in the document::

        score += log(1 + tf_saturated)

    where ``tf_saturated = tf * (k1 + 1) / (tf + k1 * (1 - 0.75 + 0.75 *
    doc_len / avg_doc_len))``, tf = count of term in document, doc_len =
    number of tokens in document, avg_doc_len is estimated as 256 (a
    reasonable default for code files).

    Args:
        avg_doc_len: Assumed average document length (default 256).
        k1:          BM25-like saturation parameter (default 1.2).
    """

    def __init__(self, avg_doc_len: float = 256.0, k1: float = 1.2) -> None:
        self._avg_doc_len = avg_doc_len if avg_doc_len > 0 else 256.0
        self._k1 = k1

    def score_batch(self, query: str, documents: list[str]) -> list[float]:
        """Score documents using BM25-like term matching.

        Args:
            query:     The search query string.
            documents: List of document texts.

        Returns:
            List of BM25-like relevance scores.
        """
        q_terms = _simple_tokenize(query)
        if not q_terms:
            return [0.0] * len(documents)

        scores: list[float] = []
        for doc in documents:
            doc_tokens = _simple_tokenize(doc)
            doc_len = len(doc_tokens)

            # Build term freq map for this document
            tf_map: dict[str, int] = {}
            for t in doc_tokens:
                tf_map[t] = tf_map.get(t, 0) + 1

            total = 0.0
            for term in q_terms:
                tf = tf_map.get(term, 0)
                if tf <= 0:
                    continue
                # BM25-style TF saturation
                numerator = tf * (self._k1 + 1.0)
                denominator = tf + self._k1 * (
                    1.0 - 0.75 + 0.75 * doc_len / self._avg_doc_len
                )
                total += math.log(1.0 + numerator / denominator)
            scores.append(total)

        return scores

    def __repr__(self) -> str:
        return (
            f"DummyCrossEncoder(avg_doc_len={self._avg_doc_len}, "
            f"k1={self._k1})"
        )


# ---------------------------------------------------------------------------
# Cohere Reranker (optional)
# ---------------------------------------------------------------------------


class CohereReranker(CrossEncoder):
    """Cross-encoder using the Cohere Rerank API.

    Requires the ``cohere`` package (``pip install cohere``).
    Import is deferred to first use so the module loads without it.

    Args:
        api_key: Cohere API key.
        model:   Model name (default ``"rerank-english-v3.0"``).
        top_n:   Number of results to return from Cohere (default None = all).
    """

    def __init__(
        self,
        api_key: str,
        model: str = "rerank-english-v3.0",
        top_n: Optional[int] = None,
    ) -> None:
        self._api_key = api_key
        self._model = model
        self._top_n = top_n
        self._client: Any = None

    def _get_client(self) -> Any:
        """Lazily initialize the Cohere client.

        Raises:
            ImportError: If the ``cohere`` package is not installed.
        """
        if self._client is None:
            try:
                import cohere  # type: ignore[import-untyped]
            except ImportError as exc:
                raise ImportError(
                    "CohereReranker requires the 'cohere' package. "
                    "Install it with: pip install cohere"
                ) from exc
            self._client = cohere.ClientV2(api_key=self._api_key)
        return self._client

    def score_batch(self, query: str, documents: list[str]) -> list[float]:
        """Score documents via Cohere Rerank API.

        Returns scores in the SAME ORDER as the input documents list
        (not Cohere's relevance-sorted order).

        Args:
            query:     The search query string.
            documents: List of document texts.

        Returns:
            List of relevance scores from the Cohere API.
        """
        if not documents:
            return []

        client = self._get_client()
        response = client.rerank(
            query=query,
            documents=documents,
            model=self._model,
            top_n=self._top_n or len(documents),
        )

        # Cohere returns results sorted by relevance; map back to input order
        scores = [0.0] * len(documents)
        for result in response.results:
            scores[result.index] = result.relevance_score
        return scores

    def __repr__(self) -> str:
        return f"CohereReranker(model={self._model!r})"


# ---------------------------------------------------------------------------
# Transformer Reranker (optional)
# ---------------------------------------------------------------------------


class TransformerReranker(CrossEncoder):
    """Cross-encoder using a HuggingFace sentence-transformers model.

    Requires ``sentence-transformers`` (``pip install sentence-transformers``).
    Import is deferred to first use.

    Args:
        model_name: HuggingFace model ID
                    (default ``"cross-encoder/ms-marco-MiniLM-L-6-v2"``).
        device:     PyTorch device string (default None = auto).
    """

    def __init__(
        self,
        model_name: str = "cross-encoder/ms-marco-MiniLM-L-6-v2",
        device: Optional[str] = None,
    ) -> None:
        self._model_name = model_name
        self._device = device
        self._model: Any = None

    def _get_model(self) -> Any:
        """Lazily initialize the sentence-transformers model.

        Raises:
            ImportError: If ``sentence-transformers`` is not installed.
        """
        if self._model is None:
            try:
                from sentence_transformers import CrossEncoder as STCrossEncoder  # type: ignore[import-untyped]
            except ImportError as exc:
                raise ImportError(
                    "TransformerReranker requires 'sentence-transformers'. "
                    "Install it with: pip install sentence-transformers"
                ) from exc
            self._model = STCrossEncoder(self._model_name, device=self._device)
        return self._model

    def score_batch(self, query: str, documents: list[str]) -> list[float]:
        """Score documents using a local cross-encoder model.

        Args:
            query:     The search query string.
            documents: List of document texts.

        Returns:
            List of model relevance scores.
        """
        if not documents:
            return []

        model = self._get_model()
        pairs = [[query, doc] for doc in documents]
        predictions = model.predict(pairs)
        return [float(s) for s in predictions]

    def __repr__(self) -> str:
        return f"TransformerReranker(model_name={self._model_name!r})"


# ---------------------------------------------------------------------------
# Reranking pipeline
# ---------------------------------------------------------------------------


class RerankerPipeline:
    """Two-stage reranking pipeline.

    Takes first-stage retrieval results and re-scores them through a
    CrossEncoder for higher-precision ranking.

    The pipeline:
        1. Takes the top ``config.top_n`` candidates from the first stage.
        2. Retrieves document text via ``doc_fn(doc_id) -> str``.
        3. Scores (query, doc_text) pairs in batches through the encoder.
        4. Optionally filters by ``config.score_threshold``.
        5. Optionally normalizes scores to [0, 1].
        6. Returns results sorted by descending reranker score.

    Args:
        encoder: A CrossEncoder instance.
        config:  RerankerConfig (optional, defaults apply).
    """

    def __init__(
        self,
        encoder: CrossEncoder,
        config: Optional[RerankerConfig] = None,
    ) -> None:
        self._encoder = encoder
        self._config = config or RerankerConfig()

    @property
    def config(self) -> RerankerConfig:
        """Return the pipeline configuration."""
        return self._config

    @property
    def encoder(self) -> CrossEncoder:
        """Return the underlying cross-encoder."""
        return self._encoder

    def __repr__(self) -> str:
        return (
            f"RerankerPipeline(encoder={self._encoder.__class__.__name__}, "
            f"config={self._config!r})"
        )

    def rerank(
        self,
        query: str,
        first_stage: Ranking,
        doc_fn: Callable[[str], str],
        *,
        top_k: Optional[int] = None,
    ) -> Ranking:
        """Rerank first-stage results through the cross-encoder.

        Args:
            query:       The search query string.
            first_stage: First-stage ranking as ``[(doc_id, score), ...]``,
                         assumed to be sorted by descending first-stage score.
            doc_fn:      Callable that returns document text given a doc_id.
                         This allows the pipeline to work with any storage
                         backend (engine index, file system, database, etc.).
            top_k:       Number of final results to return.  Defaults to
                         ``len(first_stage)`` if not specified.

        Returns:
            Reranked list of ``(doc_id, reranker_score)`` sorted by
            descending score.
        """
        cfg = self._config

        # Step 1: take top_n candidates from first stage
        candidates = first_stage[: cfg.top_n]
        if not candidates:
            return []

        # Step 2: retrieve document texts
        doc_ids = [doc_id for doc_id, _ in candidates]
        doc_texts = [doc_fn(doc_id) for doc_id in doc_ids]

        # Step 3: score in batches
        all_scores: list[float] = []
        for i in range(0, len(doc_texts), cfg.batch_size):
            batch = doc_texts[i : i + cfg.batch_size]
            batch_scores = self._encoder.score_batch(query, batch)
            all_scores.extend(batch_scores)

        # Step 4: pair doc_ids with scores
        scored: list[tuple[str, float]] = list(zip(doc_ids, all_scores))

        # Step 5: filter by threshold
        if cfg.score_threshold > 0.0:
            scored = [(did, s) for did, s in scored if s >= cfg.score_threshold]

        # Step 6: normalize to [0, 1] if configured
        if cfg.normalize and scored:
            raw = [s for _, s in scored]
            normed = _minmax_normalize(raw)
            scored = [(scored[i][0], normed[i]) for i in range(len(scored))]

        # Step 7: top-k by descending score
        k = top_k if top_k is not None else len(scored)
        return heapq.nlargest(k, scored, key=lambda x: x[1])

    def rerank_with_fusion(
        self,
        query: str,
        first_stage: Ranking,
        doc_fn: Callable[[str], str],
        *,
        alpha: float = 0.7,
        top_k: Optional[int] = None,
    ) -> Ranking:
        """Rerank and fuse first-stage + reranker scores via linear interpolation.

        ``fused_score = alpha * reranker_score + (1 - alpha) * first_stage_score``

        Both sets of scores are min-max normalized before fusion.

        Args:
            query:       The search query.
            first_stage: First-stage ranking.
            doc_fn:      Document text retrieval function.
            alpha:       Weight for reranker score (default 0.7).
            top_k:       Number of final results.

        Returns:
            Fused ranking sorted by descending score.
        """
        cfg = self._config
        candidates = first_stage[: cfg.top_n]
        if not candidates:
            return []

        doc_ids = [doc_id for doc_id, _ in candidates]
        doc_texts = [doc_fn(doc_id) for doc_id in doc_ids]

        # Score in batches
        rerank_scores: list[float] = []
        for i in range(0, len(doc_texts), cfg.batch_size):
            batch = doc_texts[i : i + cfg.batch_size]
            rerank_scores.extend(self._encoder.score_batch(query, batch))

        # Normalize first-stage scores
        fs_scores = [s for _, s in candidates]
        fs_norm = _minmax_normalize(fs_scores)

        # Normalize reranker scores
        rr_norm = _minmax_normalize(rerank_scores)

        # Fuse
        fused: list[tuple[str, float]] = []
        for i, doc_id in enumerate(doc_ids):
            fused_score = alpha * rr_norm[i] + (1.0 - alpha) * fs_norm[i]
            if fused_score >= cfg.score_threshold:
                fused.append((doc_id, fused_score))

        k = top_k if top_k is not None else len(fused)
        return heapq.nlargest(k, fused, key=lambda x: x[1])
