"""
handler.py — HuggingFace Inference Endpoint handler for SriRamanaAtmic/AtmicReRanker (BERT-large
cross-encoder, BertForSequenceClassification, 2 labels). Upload to the root of the model repo.

Scores match backend/reranker.py exactly, so the backend can fall back to local CPU scoring without
changing results: same num_labels=2 load, same tokenisation ([query, passage] pairs, passage-only
truncation at 512 tokens), fp32, softmax(logits)[:, 1], passages scored shortest first. TF32 is off
(L4 / A10G would otherwise round matmuls and move scores by ~1e-3).

Input:  {"inputs": {"query": "Who wrote Guru Ramana?", "passages": ["...", "..."]}}
Output: {"scores": [0.93, 0.12, ...]}   P(relevant) in [0, 1], in the order of passages
A warm-up is an ordinary call with one short pair.

Env (endpoint settings, optional): RERANK_BATCH (pairs per forward pass, default 32),
RERANK_MAX_PASSAGES (largest request accepted, default 128).
"""
import os

import torch
from transformers import BertConfig, BertForSequenceClassification, BertTokenizerFast

torch.backends.cuda.matmul.allow_tf32 = False   # exact fp32, as on the backend's CPU
torch.backends.cudnn.allow_tf32 = False

BATCH = int(os.environ.get("RERANK_BATCH", "32"))
MAX_PASSAGES = int(os.environ.get("RERANK_MAX_PASSAGES", "128"))


class EndpointHandler:
    def __init__(self, path=""):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.tokenizer = BertTokenizerFast.from_pretrained(path)
        self.model = BertForSequenceClassification.from_pretrained(
            path, config=BertConfig.from_pretrained(path, num_labels=2),   # config.json leaves num_labels unset
        ).to(self.device).eval()

    def __call__(self, data: dict) -> dict:
        inputs = data.get("inputs", data)
        if not isinstance(inputs, dict):
            raise ValueError('expected {"inputs": {"query": str, "passages": [str, ...]}}')
        query, passages = inputs.get("query"), inputs.get("passages")
        if not isinstance(query, str) or not query.strip():
            raise ValueError("query must be a non-empty string")
        if not isinstance(passages, list) or not all(isinstance(p, str) for p in passages):
            raise ValueError("passages must be a list of strings")
        if len(passages) > MAX_PASSAGES:
            raise ValueError(f"at most {MAX_PASSAGES} passages per request")

        # Shortest first: a batch is padded to its longest pair, so similar lengths waste least work.
        order = sorted(range(len(passages)), key=lambda i: len(passages[i]))
        scores = [0.0] * len(passages)
        with torch.inference_mode():
            for s in range(0, len(order), BATCH):
                idx = order[s:s + BATCH]
                enc = self.tokenizer(
                    [[query, passages[i]] for i in idx],
                    padding=True, truncation="only_second", max_length=512, return_tensors="pt",
                ).to(self.device)
                probs = self.model(**enc).logits.float().softmax(-1)[:, 1].tolist()
                for i, p in zip(idx, probs):
                    scores[i] = p
        return {"scores": scores}