"""
reranker.py
-----------
Re-ranks retrieved chunks using a cross-encoder model for better relevance ordering.

Cross-encoders (unlike bi-encoders) compute similarity by processing the query
and document together, providing more accurate ranking than vector similarity alone.

Uses: cross-encoder/ms-marco-MiniLM-L-12-v2 (fast, accurate, lightweight)
"""

import logging
from typing import List, Tuple

from langchain.schema import Document
from sentence_transformers import CrossEncoder

from app.config import RERANKER_MODEL_NAME, RERANKER_TOP_K

logger = logging.getLogger(__name__)


class Reranker:
    """
    Re-ranks retrieved document chunks using a cross-encoder model.

    Args:
        model_name: HuggingFace model ID for the cross-encoder (default from config).
        top_k:      Number of results to return after re-ranking.
    """

    def __init__(
        self,
        model_name: str = RERANKER_MODEL_NAME,
        top_k: int = RERANKER_TOP_K,
    ) -> None:
        self.model_name = model_name
        self.top_k = top_k
        self._model: CrossEncoder | None = None

    def _load_model(self) -> CrossEncoder:
        """
        Lazily load the cross-encoder model on first use.
        """
        if self._model is None:
            logger.info(f"Loading cross-encoder model: {self.model_name}")
            self._model = CrossEncoder(self.model_name)
        return self._model

    def rerank(
        self,
        query: str,
        documents: List[Tuple[Document, float]],
        top_k: int | None = None,
    ) -> List[Tuple[Document, float]]:
        """
        Re-rank retrieved documents by relevance to the query.

        Uses cross-encoder scores instead of raw vector similarity.
        Lower cross-encoder scores are better (0=irrelevant, 1=highly relevant in most cases,
        but the scale depends on the model).

        Args:
            query:     User's natural-language question.
            documents: List of (Document, vector_score) tuples from initial retrieval.
            top_k:     Number of results to return (falls back to self.top_k).

        Returns:
            List of (Document, cross_encoder_score) tuples, sorted by relevance.
        """
        if not documents:
            logger.debug("No documents to rerank.")
            return []

        k = top_k or self.top_k
        model = self._load_model()

        # Extract document texts and original scores
        doc_list = [doc for doc, _ in documents]
        original_scores = {doc.page_content: score for doc, score in documents}

        # Prepare pairs for cross-encoder: (query, document_text)
        pairs = [[query, doc.page_content] for doc in doc_list]

        try:
            # Get cross-encoder scores
            # Returns a list of scores; higher is better for most models
            ce_scores = model.predict(pairs)

            # Combine documents with their cross-encoder scores
            ranked = list(zip(doc_list, ce_scores))

            # Sort by cross-encoder score (descending)
            ranked.sort(key=lambda x: x[1], reverse=True)

            # Return top-k
            result = ranked[:k]
            logger.debug(
                f"Reranked {len(documents)} documents → top {len(result)} by cross-encoder"
            )
            return result

        except Exception as exc:
            logger.error(f"Cross-encoder reranking failed: {exc}")
            # Fall back to original vector scores if reranking fails
            return documents[:k]

    def rerank_and_get_documents(
        self,
        query: str,
        documents: List[Tuple[Document, float]],
        top_k: int | None = None,
    ) -> List[Document]:
        """
        Convenience wrapper — returns only Document objects (no scores).
        """
        return [doc for doc, _ in self.rerank(query, documents, top_k=top_k)]