"""Minimal TEI-compatible cross-encoder rerank server (GPU).

Why this exists: HuggingFace's official Text-Embeddings-Inference GPU images
require CUDA compute capability >= 7.5 (Turing+). This box has a GTX 1070
(Pascal, 6.1), so the stock TEI image won't run. Plain CUDA torch DOES support
Pascal (that's why ollama works here), so we serve the same
`jina-reranker-v2-base-multilingual` cross-encoder via torch and expose only the
two endpoints Hindsight's `tei` reranker provider calls:
  GET  /info    -> JSON (init/health probe)
  POST /rerank  -> {"query": str, "texts": [str], ...}
                   -> bare list [{"index": i, "score": f}, ...] sorted desc
See hindsight_api/engine/cross_encoder.py::RemoteTEICrossEncoder for the client.
"""

import os
import torch
from fastapi import FastAPI
from pydantic import BaseModel
from transformers import AutoModelForSequenceClassification

MODEL_ID = os.environ.get("RERANKER_MODEL", "jinaai/jina-reranker-v2-base-multilingual")
DEVICE = os.environ.get("RERANKER_DEVICE", "cuda")
MAX_LENGTH = int(os.environ.get("RERANKER_MAX_LENGTH", "1024"))
# fp16 on GPU halves the ~1.1GB fp32 footprint; Pascal supports fp16 storage.
DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32

app = FastAPI(title="tei-reranker")
_model = None


def _load():
    global _model
    if _model is not None:
        return
    m = AutoModelForSequenceClassification.from_pretrained(
        MODEL_ID, torch_dtype=DTYPE, trust_remote_code=True
    )
    m.to(DEVICE)
    m.eval()
    _model = m


@app.on_event("startup")
def startup():
    _load()


class RerankRequest(BaseModel):
    query: str
    texts: list[str]
    return_text: bool = False
    truncate: bool | None = None
    raw_scores: bool | None = None


@app.get("/info")
def info():
    # Hindsight only needs a 200 JSON here to consider the server initialized.
    return {
        "model_id": MODEL_ID,
        "model_dtype": str(DTYPE).replace("torch.", ""),
        "model_type": {"reranker": {}},
        "max_input_length": MAX_LENGTH,
        "device": DEVICE,
    }


@app.get("/health")
def health():
    return {"status": "ok" if _model is not None else "loading"}


@app.post("/rerank")
def rerank(req: RerankRequest):
    if not req.texts:
        return []
    pairs = [[req.query, t] for t in req.texts]
    with torch.no_grad():
        # jina-reranker-v2 exposes compute_score (batches + moves to device).
        scores = _model.compute_score(pairs, max_length=MAX_LENGTH)
    if not isinstance(scores, list):
        scores = [scores]
    results = [{"index": i, "score": float(s)} for i, s in enumerate(scores)]
    results.sort(key=lambda r: r["score"], reverse=True)
    return results