"""Author RAG Chatbot SaaS — Cross-Encoder Re-Ranker.

Uses cross-encoder/ms-marco-MiniLM-L-6-v2 (free, local) to re-rank
retrieved chunks by relevance to the original query.
Significantly improves precision over cosine similarity alone.
RULE: Keep top N chunks above minimum score threshold.

BUG-6 fix: reranker.predict() is CPU-bound sync — wrapped in asyncio.to_thread().
"""

import asyncio

import structlog

from app.config import get_settings
from app.services.vector_store import RetrievedChunk

logger = structlog.get_logger(__name__)
cfg = get_settings()

_reranker = None


async def get_reranker():
    """Lazily load and cache the cross-encoder re-ranker.

    Returns:
        Loaded CrossEncoder model.
    """
    global _reranker
    if _reranker is None:
        from sentence_transformers import CrossEncoder
        logger.info("Loading cross-encoder re-ranker (first load)...")
        _reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")
        logger.info("Cross-encoder re-ranker loaded successfully")
    return _reranker


async def rerank_chunks(
    query: str,
    chunks: list[RetrievedChunk],
    top_n: int | None = None,
    min_score: float | None = None,
) -> list[RetrievedChunk]:
    """Re-rank retrieved chunks using cross-encoder scoring.

    Args:
        query: The original (non-rewritten) user query.
        chunks: List of RetrievedChunk from the retriever.
        top_n: Maximum chunks to keep after re-ranking.
        min_score: Minimum cross-encoder score to keep a chunk.

    Returns:
        Re-ranked and filtered list of chunks (best first).
    """
    top_n = top_n or cfg.RAG_RERANK_TOP_N
    min_score = min_score or cfg.RAG_RERANK_MIN_SCORE

    if not chunks:
        return []

    try:
        reranker = await get_reranker()

        # Build (query, chunk) pairs for cross-encoder
        pairs = [(query, chunk.text) for chunk in chunks]

        # BUG-6 fix: predict() is synchronous CPU-bound inference — offload to thread pool.
        scores = await asyncio.to_thread(reranker.predict, pairs)

        # Apply scores
        for chunk, score in zip(chunks, scores):
            chunk.rerank_score = float(score)

        # Sort by rerank score descending
        ranked = sorted(chunks, key=lambda c: c.rerank_score, reverse=True)

        # Filter by minimum score and limit to top_n
        filtered = [c for c in ranked if c.rerank_score >= min_score][:top_n]

        logger.debug(
            "Re-ranking complete",
            input_chunks=len(chunks),
            output_chunks=len(filtered),
            top_score=filtered[0].rerank_score if filtered else 0,
        )
        return filtered

    except Exception as e:
        logger.error("Re-ranker failed, returning top-K by cosine score", error=str(e))
        # Graceful fallback: return top chunks by initial similarity score
        return sorted(chunks, key=lambda c: c.score, reverse=True)[:top_n]