"""Rerank 模型集成 — 支持 Cohere Rerank 和 BGE-Reranker（本地部署）。

数据出境风险：Cohere Rerank 将文档 chunks 发送到第三方 API。
敏感数据 KB 应使用 BGE-Reranker via Xinference（本地部署）。
"""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING

from pydantic import BaseModel, ConfigDict

from agentkit.rag_platform.models import QueryResult

if TYPE_CHECKING:
    from llama_index.core.schema import NodeWithScore

logger = logging.getLogger(__name__)


class RerankConfig(BaseModel):
    """Rerank 配置 — 可按 KB 覆盖。

    provider:
        - "cohere": Cohere Rerank API（数据出境）
        - "bge": BGE-Reranker via Xinference（本地部署，敏感数据 KB 推荐）
        - "none": 不重排
    """

    model_config = ConfigDict()

    provider: str = "none"
    api_key: str | None = None
    base_url: str | None = None  # Xinference URL for BGE
    model_name: str | None = None  # 模型名（如 "bge-reranker-base"）
    top_n: int = 5
    # True 表示当前 KB 使用了云端 rerank，存在数据出境风险
    data_export_warning: bool = False


class Reranker:
    """Rerank 引擎 — 包装 LlamaIndex rerankers。

    使用方式：
        config = RerankConfig(provider="cohere", api_key="...", top_n=5)
        reranker = Reranker(config)
        reranked = await reranker.rerank(query, results)
    """

    def __init__(self, config: RerankConfig) -> None:
        self._config = config
        self._reranker: object | None = None  # 延迟初始化，避免 import 失败

    def _get_reranker(self) -> object | None:
        """延迟加载 reranker 实例 — 避免在 import 时失败。"""
        if self._reranker is not None:
            return self._reranker

        cfg = self._config
        if cfg.provider == "cohere":
            self._reranker = self._build_cohere_reranker(cfg)
        elif cfg.provider == "bge":
            self._reranker = self._build_bge_reranker(cfg)
        elif cfg.provider == "none":
            self._reranker = None
        else:  # pragma: no cover — 配置校验已穷尽
            raise ValueError(f"Unsupported rerank provider: {cfg.provider}")

        return self._reranker

    def _build_cohere_reranker(self, cfg: RerankConfig) -> object:
        """构建 CohereRerank — 数据出境，需 api_key。"""
        if not cfg.api_key:
            raise ValueError("Cohere rerank requires api_key")
        try:
            from llama_index.postprocessor.cohere_rerank import CohereRerank
        except ImportError as e:
            raise ImportError(
                "CohereRerank requires llama-index-postprocessor-cohere-rerank. "
                "Install: pip install llama-index-postprocessor-cohere-rerank"
            ) from e

        kwargs: dict[str, object] = {
            "api_key": cfg.api_key,
            "top_n": cfg.top_n,
        }
        if cfg.model_name:
            kwargs["model"] = cfg.model_name
        return CohereRerank(**kwargs)

    def _build_bge_reranker(self, cfg: RerankConfig) -> object:
        """构建 BGE-Reranker via Xinference（本地部署，无数据出境）。

        使用 SentenceTransformerRerank 作为本地 BGE-Reranker 的封装。
        若 base_url 指向 Xinference，调用方应在 KB 设置中标注为本地部署。
        """
        try:
            from llama_index.postprocessor.sentence_transformers_rerank import (
                SentenceTransformerRerank,
            )
        except ImportError as e:
            raise ImportError(
                "SentenceTransformerRerank requires "
                "llama-index-postprocessor-sentence-transformers-rerank. "
                "Install: pip install llama-index-postprocessor-sentence-transformers-rerank"
            ) from e

        model = cfg.model_name or "BAAI/bge-reranker-base"
        kwargs: dict[str, object] = {
            "model": model,
            "top_n": cfg.top_n,
        }
        return SentenceTransformerRerank(**kwargs)

    async def rerank(
        self,
        query: str,
        results: list[QueryResult],
    ) -> list[QueryResult]:
        """对检索结果重排，返回按相关性排序的 top_n 结果。

        Args:
            query: 查询文本
            results: 原始检索结果列表

        Returns:
            重排后的 QueryResult 列表（按相关性降序），分数已更新为 rerank 分数。
            若 provider == "none" 或 results 为空，原样返回。
        """
        # 空结果或关闭重排 — 直接返回
        if not results or self._config.provider == "none":
            return list(results)

        reranker = self._get_reranker()
        if reranker is None:
            return list(results)

        # 将 QueryResult 转为 LlamaIndex NodeWithScore
        nodes_with_scores = self._to_nodes_with_scores(results)

        try:
            # LlamaIndex reranker 的 postprocessnodes 是同步方法
            reranked_nodes = reranker.postprocessnodes(
                nodes_with_scores,
                query_str=query,
            )
        except Exception as e:
            logger.warning("Rerank failed, returning original results: %s", e)
            return list(results)

        # 将重排结果映射回 QueryResult，更新分数
        return self._from_nodes_with_scores(reranked_nodes, results)

    @staticmethod
    def _to_nodes_with_scores(
        results: list[QueryResult],
    ) -> "list[NodeWithScore]":
        """将 QueryResult 列表转为 LlamaIndex NodeWithScore 列表。"""
        from llama_index.core.schema import NodeWithScore, TextNode

        out: list[NodeWithScore] = []
        for r in results:
            node = TextNode(
                id_=r.chunk_id,
                text=r.content,
                metadata=r.metadata,
            )
            out.append(NodeWithScore(node=node, score=r.score))
        return out

    @staticmethod
    def _from_nodes_with_scores(
        nodes: "list[NodeWithScore]",
        original: list[QueryResult],
    ) -> list[QueryResult]:
        """将重排后的 NodeWithScore 列表转回 QueryResult，更新分数。

        通过 node_id 匹配原始 QueryResult 的元数据（document_id, kb_id）。
        """
        original_map = {r.chunk_id: r for r in original}
        out: list[QueryResult] = []
        for nws in nodes:
            node_id = nws.node.node_id if hasattr(nws.node, "node_id") else None
            original_r = original_map.get(node_id) if node_id else None
            if original_r is None:
                # 兜底：通过内容匹配（理论上不应触发）
                content = nws.node.get_content() if hasattr(nws.node, "get_content") else ""
                original_r = next(
                    (r for r in original if r.content == content),
                    None,
                )
                if original_r is None:
                    continue

            new_score = float(nws.score) if nws.score is not None else original_r.score
            out.append(
                QueryResult(
                    chunk_id=original_r.chunk_id,
                    content=original_r.content,
                    score=new_score,
                    metadata=original_r.metadata,
                    document_id=original_r.document_id,
                    kb_id=original_r.kb_id,
                )
            )
        return out


__all__ = ["RerankConfig", "Reranker"]