from dataclasses import replace
from threading import Lock
from typing import Any

from zbt.core.config import Settings
from zbt.core.errors import AppError
from zbt.domains.knowledge.models import KnowledgeSearchHit
from zbt.infrastructure.embedding.bge_model import resolve_local_huggingface_snapshot


class BgeKnowledgeReranker:
    """使用本地BGE Reranker对混合检索候选结果做语义重排。"""

    def __init__(self, settings: Settings, model: Any | None = None) -> None:
        self._settings = settings
        self._model = model
        self._model_lock = Lock()

    def warmup(self) -> None:
        """在接收首个知识问答请求前完成模型装载与一次最小推理。"""
        self._get_model().compute_score(
            [["智保通", "智保通保险知识服务"]],
            normalize=True,
        )

    def rerank(
        self,
        query: str,
        hits: list[KnowledgeSearchHit],
        *,
        limit: int,
    ) -> list[KnowledgeSearchHit]:
        if not hits:
            return []
        try:
            raw_scores = self._get_model().compute_score(
                [[query, hit.content] for hit in hits],
                normalize=True,
            )
            scores = (
                raw_scores.tolist()
                if hasattr(raw_scores, "tolist")
                else list(raw_scores)
                if isinstance(raw_scores, (list, tuple))
                else [raw_scores]
            )
            if len(scores) != len(hits):
                raise ValueError("BGE Reranker返回数量与候选数量不一致")
            ranked = [
                replace(hit, score=float(score))
                for hit, score in zip(hits, scores, strict=True)
            ]
            return sorted(ranked, key=lambda hit: hit.score, reverse=True)[:limit]
        except AppError:
            raise
        except Exception as error:
            raise AppError(
                "KNOWLEDGE_RERANK_FAILED",
                "知识重排失败，请检查本地BGE Reranker配置",
                503,
                retryable=True,
            ) from error

    def _get_model(self) -> Any:
        if self._model is None:
            with self._model_lock:
                if self._model is None:
                    from FlagEmbedding import FlagReranker  # type: ignore[import-untyped]

                    self._model = FlagReranker(
                        resolve_local_huggingface_snapshot(
                            self._settings.bge_reranker_model_path,
                        ),
                        devices=self._settings.bge_reranker_device,
                        use_fp16=self._settings.bge_reranker_device != "cpu",
                    )
        return self._model
