"""
Minimal BGE reranker for pairwise scoring (query, doc).

Features:
- Model loading with optional FP16
- Length-based sorting to reduce padding waste
- Deduplication to avoid redundant inference
- Scores returned in original doc order
"""

import logging
import math
import threading
import time
from typing import Any, Dict, List, Optional, Tuple

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

logger = logging.getLogger("reranker.core")


class BGEReranker:
    def __init__(
        self,
        model_name: str = "BAAI/bge-reranker-v2-m3",
        device: Optional[str] = None,
        batch_size: int = 64,
        use_fp16: bool = True,
        max_length: int = 512,
        cache_dir: str = "./model_cache",
        enable_warmup: bool = True,
    ) -> None:
        self.model_name = model_name
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self.batch_size = max(1, int(batch_size))
        self.max_length = int(max_length)
        self.use_fp16 = bool(use_fp16 and self.device == "cuda")
        self._lock = threading.Lock()

        logger.info(
            "[BGE_RERANKER] Loading model %s on %s (fp16=%s)",
            self.model_name,
            self.device,
            self.use_fp16,
        )

        self.tokenizer = AutoTokenizer.from_pretrained(
            self.model_name, trust_remote_code=True, cache_dir=cache_dir
        )
        self.model = AutoModelForSequenceClassification.from_pretrained(
            self.model_name, trust_remote_code=True, cache_dir=cache_dir
        )

        self.model = self.model.to(self.device)
        self.model.eval()

        if self.use_fp16:
            self.model = self.model.half()

        if self.device == "cuda":
            torch.backends.cudnn.benchmark = True

        if enable_warmup:
            self._warmup()

        logger.info(
            "[BGE_RERANKER] Model ready | model=%s device=%s fp16=%s batch=%s max_len=%s",
            self.model_name,
            self.device,
            self.use_fp16,
            self.batch_size,
            self.max_length,
        )

    def _warmup(self) -> None:
        try:
            with torch.inference_mode():
                pairs = [["warmup", "warmup"]]
                inputs = self.tokenizer(
                    pairs,
                    padding=True,
                    truncation=True,
                    return_tensors="pt",
                    max_length=self.max_length,
                )
                inputs = {k: v.to(self.device) for k, v in inputs.items()}
                if self.use_fp16:
                    inputs = {
                        k: (v.half() if v.dtype == torch.float32 else v)
                        for k, v in inputs.items()
                    }
                _ = self.model(**inputs, return_dict=True).logits
                if self.device == "cuda":
                    torch.cuda.synchronize()
        except Exception as exc:
            logger.warning("[BGE_RERANKER] Warmup failed: %s", exc)

    def score(self, query: str, docs: List[str], normalize: bool = True) -> List[float]:
        scores, _meta = self.score_with_meta(query, docs, normalize=normalize)
        return scores

    def score_with_meta(
        self, query: str, docs: List[str], normalize: bool = True
    ) -> Tuple[List[float], Dict[str, Any]]:
        start_ts = time.time()

        if docs is None:
            docs = []

        query = "" if query is None else str(query).strip()
        total_docs = len(docs)
        output_scores: List[float] = [0.0] * total_docs

        # Log request summary (query + first 3 docs preview)
        preview_docs: List[str] = []
        for d in docs[:3]:
            preview_docs.append("" if d is None else str(d))
        logger.info(
            "[BGE_RERANKER] Request | query=%r | docs=%d | docs_preview=%s",
            query,
            total_docs,
            preview_docs,
        )

        indexed_docs: List[Tuple[int, str]] = []
        for i, doc in enumerate(docs):
            if doc is None:
                continue
            text = str(doc).strip()
            if not text:
                continue
            indexed_docs.append((i, text))

        if not query or not indexed_docs:
            elapsed_ms = (time.time() - start_ts) * 1000.0
            return output_scores, {
                "input_docs": total_docs,
                "usable_docs": len(indexed_docs),
                "unique_docs": 0,
                "dedup_ratio": 0.0,
                "elapsed_ms": round(elapsed_ms, 3),
            }

        # Sort by estimated length + text to cluster similar lengths
        indexed_docs.sort(key=lambda x: (len(x[1]), x[1]))

        unique_texts: List[str] = []
        position_to_unique: List[int] = []
        prev_text: Optional[str] = None

        for _idx, text in indexed_docs:
            if text != prev_text:
                unique_texts.append(text)
                prev_text = text
            position_to_unique.append(len(unique_texts) - 1)

        logger.debug(
            "[BGE_RERANKER] Preprocess | input=%d usable=%d unique=%d",
            total_docs,
            len(indexed_docs),
            len(unique_texts),
        )

        unique_scores = self._score_unique(
            query=query, passages=unique_texts, normalize=normalize
        )

        for (orig_idx, _text), unique_idx in zip(indexed_docs, position_to_unique):
            output_scores[orig_idx] = float(unique_scores[unique_idx])

        # Log per-doc scores (aligned to original docs order)
        if 0:
            try:
                lines = []
                for i, d in enumerate(docs[:100]):
                    lines.append(f"{output_scores[i]},{'' if d is None else str(d)}")
                logger.info("[BGE_RERANKER] query:%s Scores (score,doc):\n%s", query, "\n".join(lines))
            except Exception:
                pass

        elapsed_ms = (time.time() - start_ts) * 1000.0
        dedup_ratio = 0.0
        if indexed_docs:
            dedup_ratio = 1.0 - (len(unique_texts) / float(len(indexed_docs)))

        meta = {
            "input_docs": total_docs,
            "usable_docs": len(indexed_docs),
            "unique_docs": len(unique_texts),
            "dedup_ratio": round(dedup_ratio, 4),
            "elapsed_ms": round(elapsed_ms, 3),
            "model": self.model_name,
            "device": self.device,
            "fp16": self.use_fp16,
            "batch_size": self.batch_size,
            "max_length": self.max_length,
            "normalize": normalize,
        }

        logger.info(
            "[BGE_RERANKER] Done | input=%d usable=%d unique=%d dedup=%s elapsed_ms=%s",
            meta["input_docs"],
            meta["usable_docs"],
            meta["unique_docs"],
            meta["dedup_ratio"],
            meta["elapsed_ms"],
        )

        return output_scores, meta

    def _compute_optimal_batch_size(self, total: int) -> int:
        if total <= 0:
            return 1
        current_batch_size = self.batch_size + 8
        current_batch_count = math.ceil(total / current_batch_size)

        optimal_batch_size = current_batch_size
        test_batch_size = current_batch_size - 4

        while test_batch_size > 0:
            test_batch_count = math.ceil(total / test_batch_size)
            if test_batch_count <= current_batch_count:
                optimal_batch_size = test_batch_size
                test_batch_size -= 4
            else:
                break

        return max(1, optimal_batch_size)

    def _score_unique(
        self, query: str, passages: List[str], normalize: bool = True
    ) -> List[float]:
        if not passages:
            return []

        optimal_batch_size = self._compute_optimal_batch_size(len(passages))

        logger.info(
            "[BGE_RERANKER] Reranking %d unique passages | batch=%d | device=%s | fp16=%s",
            len(passages),
            optimal_batch_size,
            self.device,
            self.use_fp16,
        )

        scores: List[float] = []

        with self._lock:
            for i in range(0, len(passages), optimal_batch_size):
                batch_passages = passages[i : i + optimal_batch_size]
                pairs = [[query, passage] for passage in batch_passages]

                with torch.inference_mode():
                    inputs = self.tokenizer(
                        pairs,
                        padding=True,
                        truncation=True,
                        return_tensors="pt",
                        max_length=self.max_length,
                        add_special_tokens=True,
                    )
                    inputs = {k: v.to(self.device) for k, v in inputs.items()}

                    if self.use_fp16:
                        inputs = {
                            k: (v.half() if v.dtype == torch.float32 else v)
                            for k, v in inputs.items()
                        }

                    logits = self.model(**inputs, return_dict=True).logits.view(-1).float()
                    if normalize:
                        logits = torch.sigmoid(logits)
                    batch_scores = logits.detach().cpu().numpy().tolist()
                    scores.extend(batch_scores)

        return scores