"""Step 09 — cross-encoder rerank (granite-embedding-reranker-r2 ONNX): the most precise relevance signal.

A cross-encoder reads the JD query and each candidate's text TOGETHER (one transformer pass per pair) and
emits a single relevance logit — strictly more precise than the bi-encoder cosine of Step 08, and the most
expensive per item, so it runs only on the ~300 granite survivors. ModernBERT sequence-classification head;
onnxruntime + tokenizers (NO torch). Re-ranks by ``reranker_score`` (raw logit; monotonic for ordering).
"""

from __future__ import annotations

from functools import lru_cache
from typing import Any

import numpy as np
import polars as pl

from common.logging import step
from config import Settings, load_settings

from .onnx import load_onnx_session


@lru_cache(maxsize=1)
def _load_reranker(repo: str, onnx_file: str, max_length: int, cpu_threads: int = 0) -> tuple[Any, Any, set[str]]:
    """Load (downloading on first use) the cross-encoder ONNX session + tokenizer; cached for the process."""
    return load_onnx_session(repo, onnx_file, max_length, "reranker.load", cpu_threads)


def _score_pairs(
    session: Any, tokenizer: Any, input_names: set[str], query: str, docs: list[str], batch_size: int
) -> np.ndarray:
    """Cross-encode (query, doc) pairs → 1-D array of relevance logits, in doc order."""
    logits: list[np.ndarray] = []
    for start in range(0, len(docs), batch_size):
        batch_pairs = [(query, doc) for doc in docs[start : start + batch_size]]
        tokenizer.enable_padding()
        encodings = tokenizer.encode_batch(batch_pairs)  # pair encoding: [CLS] query [SEP] doc [SEP]
        token_ids = np.array([enc.ids for enc in encodings], dtype=np.int64)
        feed: dict[str, np.ndarray] = {
            "input_ids": token_ids,
            "attention_mask": np.array([enc.attention_mask for enc in encodings], dtype=np.int64),
        }
        if "token_type_ids" in input_names:
            feed["token_type_ids"] = np.zeros_like(token_ids)
        batch_logits = session.run(None, {k: v for k, v in feed.items() if k in input_names})[0]
        logits.append(np.asarray(batch_logits, dtype=np.float64).reshape(-1))  # (batch, 1) → (batch,)
    return np.concatenate(logits) if logits else np.zeros(0)


def cross_encoder_rerank(
    candidates: pl.DataFrame, query_text: str, settings: Settings | None = None
) -> pl.DataFrame:
    """Add ``reranker_score`` (cross-encoder relevance) and re-sort by it. Empty/blank ⇒ passthrough at 0."""
    active = settings or load_settings()
    rerank = active.rerank
    with step("rank.cross_encoder", rows_in=candidates.height) as metrics:
        if candidates.height == 0 or not query_text.strip():
            metrics["rows_out"] = candidates.height
            metrics["skipped"] = True
            return candidates.with_columns(pl.lit(0.0).alias("reranker_score"))
        session, tokenizer, input_names = _load_reranker(
            active.models.reranker, rerank.onnx_file, rerank.max_length, active.compute.cpu_threads
        )
        scores = _score_pairs(
            session, tokenizer, input_names, query_text, candidates.get_column("text").to_list(), rerank.batch_size
        )
        reranked = candidates.with_columns(pl.Series("reranker_score", scores)).sort(
            ["reranker_score", "candidate_id"], descending=[True, False]
        )
        metrics["rows_out"] = reranked.height
        metrics["score_max"] = round(float(scores.max()), 4)
        return reranked