"""双索引检索引擎 — embedding/keywords/blend 三模式。

- embedding: pgvector 语义检索（LlamaIndex PGVectorStore.aquery）
- keywords:  PG 全文检索（jieba 分词 + tsquery_rank）
- blend:     并行执行两种检索，按分数归一化后合并去重排序

参考 LlamaIndex hybrid retriever 模式（VectorStoreRetriever + 全文检索）。
ACL 过滤由调用方在传入 `kb_ids` 前完成（见 acl.filter_kb_by_user_acl）。
"""

from __future__ import annotations

import asyncio
import logging
from typing import TYPE_CHECKING, Any

from sqlalchemy import text

from agentkit.rag_platform.fulltext import KB_CHUNKS_TABLE, build_tsquery
from agentkit.rag_platform.models import QueryMode, QueryResult

if TYPE_CHECKING:
    from llama_index.core.embeddings import BaseEmbedding
    from llama_index.core.vector_stores.types import VectorStoreQuery
    from llama_index.vector_stores.postgres import PGVectorStore

logger = logging.getLogger(__name__)

# ponytail: 默认归一化策略 — min-max，将每种检索的分数映射到 [0,1]
# 升级路径：若需更精细的融合权重，可引入 RRF (Reciprocal Rank Fusion) 或学习权重
_BLEND_WEIGHT_EMBEDDING = 0.6
_BLEND_WEIGHT_KEYWORDS = 0.4


def _normalize_scores(scores: list[float]) -> list[float]:
    """min-max 归一化分数到 [0, 1]。

    空列表返回空。常数列表（含单元素）返回全 1.0 — 所有结果等价于最高相关度。
    """
    if not scores:
        return []
    lo, hi = min(scores), max(scores)
    if hi - lo < 1e-9:
        return [1.0 for _ in scores]
    return [(s - lo) / (hi - lo) for s in scores]


class RetrievalEngine:
    """双索引检索引擎 — embedding/keywords/blend 三模式。

    Args:
        vector_store: LlamaIndex PGVectorStore 实例（embedding 模式使用）
        session_factory: SQLAlchemy async session factory（keywords 模式使用）
        embed_model: LlamaIndex BaseEmbedding（embedding 模式使用）
    """

    def __init__(
        self,
        vector_store: "PGVectorStore",
        session_factory,
        embed_model: "BaseEmbedding | None" = None,
    ) -> None:
        self._vector_store = vector_store
        self._sf = session_factory
        self._embed_model = embed_model

    async def retrieve(
        self,
        query: str,
        kb_ids: list[str],
        mode: QueryMode,
        top_k: int = 5,
    ) -> list[QueryResult]:
        """检索入口 — 根据 mode 分发到具体实现。

        Args:
            query: 查询文本
            kb_ids: 限定检索的知识库 ID 列表（已通过 ACL 过滤）
            mode: 检索模式
            top_k: 返回结果数

        Returns:
            QueryResult 列表（按分数降序）。无结果时返回空列表。
        """
        if not kb_ids:
            return []

        if mode == QueryMode.embedding:
            return await self._retrieve_embedding(query, kb_ids, top_k)
        elif mode == QueryMode.keywords:
            return await self._retrieve_keywords(query, kb_ids, top_k)
        elif mode == QueryMode.blend:
            return await self._retrieve_blend(query, kb_ids, top_k)
        else:  # pragma: no cover — 枚举已穷尽
            raise ValueError(f"Unsupported query mode: {mode}")

    async def _retrieve_embedding(
        self,
        query: str,
        kb_ids: list[str],
        top_k: int,
    ) -> list[QueryResult]:
        """pgvector 语义检索。

        使用 LlamaIndex vector_store.aquery 执行向量检索，结果按 metadata.kb_id
        过滤到 kb_ids 子集。
        """
        if self._embed_model is None:
            raise ValueError("embed_model is required for embedding retrieval mode")

        from llama_index.core.vector_stores.types import VectorStoreQuery

        query_embedding = await self._embed_model.aget_text_embedding(query)
        vs_query: VectorStoreQuery = VectorStoreQuery(
            query_embedding=query_embedding,
            similarity_top_k=top_k,
        )
        result = await self._vector_store.aquery(vs_query)

        kb_set = set(kb_ids)
        out: list[QueryResult] = []
        for node, score in zip(result.nodes, result.similarities):
            meta = node.metadata or {}
            node_kb = meta.get("kb_id", "")
            if node_kb not in kb_set:
                continue
            out.append(
                QueryResult(
                    chunk_id=node.node_id,
                    content=node.get_content(),
                    score=float(score) if score is not None else 0.0,
                    metadata=meta,
                    document_id=meta.get("document_id", ""),
                    kb_id=node_kb,
                )
            )
            if len(out) >= top_k:
                break
        return out

    async def _retrieve_keywords(
        self,
        query: str,
        kb_ids: list[str],
        top_k: int,
    ) -> list[QueryResult]:
        """PG 全文检索（jieba 分词 + tsquery）。

        使用 `ts_rank(search_vector, tsquery)` 排序，按 kb_id 过滤。
        kb_id 存储在 metadata_ JSON 列中（LlamaIndex PGVectorStore schema），
        用 `metadata_->>'kb_id'` 提取。
        """
        tsquery = build_tsquery(query)
        if not tsquery:
            return []

        # ponytail: 用 ANY(%s) 传 kb_ids 列表，避免字符串拼接
        # 升级路径：若 kb_ids 数量超过 PG 参数限制（32k），需分批查询
        # 列名参考 LlamaIndex PGVectorStore 默认 schema：
        #   id / embedding / text / metadata_ (JSON) / document_id / search_vector
        sql = text(
            f"""
            SELECT id, text, metadata_, document_id,
                   ts_rank(search_vector, to_tsquery('simple', :tsquery)) AS score
            FROM {KB_CHUNKS_TABLE}
            WHERE search_vector @@ to_tsquery('simple', :tsquery)
              AND metadata_->>'kb_id' = ANY(:kb_ids)
            ORDER BY score DESC
            LIMIT :top_k
            """  # noqa: S608 — 表名为常量
        )

        async with self._sf() as db:
            result = await db.execute(
                sql,
                {
                    "tsquery": tsquery,
                    "kb_ids": list(kb_ids),
                    "top_k": top_k,
                },
            )
            rows = result.all()

        out: list[QueryResult] = []
        for row in rows:
            chunk_id, content, metadata, document_id, score = row
            meta: dict[str, Any] = metadata if isinstance(metadata, dict) else {}
            kb_id = meta.get("kb_id", "")
            out.append(
                QueryResult(
                    chunk_id=str(chunk_id),
                    content=content or "",
                    score=float(score) if score is not None else 0.0,
                    metadata=meta,
                    document_id=str(document_id) if document_id is not None else "",
                    kb_id=str(kb_id) if kb_id is not None else "",
                )
            )
        return out

    async def _retrieve_blend(
        self,
        query: str,
        kb_ids: list[str],
        top_k: int,
    ) -> list[QueryResult]:
        """双索引合并 — 语义 + 全文结果去重排序。

        并行执行两种检索（各取 top_k），按 chunk_id 去重，分数归一化后加权融合。
        若任一检索失败（如 embed_model 未配置），降级为另一种。
        """
        # 并行执行两种检索
        embed_task = self._safe_retrieve_embedding(query, kb_ids, top_k)
        kw_task = self._retrieve_keywords(query, kb_ids, top_k)
        embed_results, kw_results = await asyncio.gather(
            embed_task, kw_task, return_exceptions=False
        )

        # 归一化分数
        embed_scores = _normalize_scores([r.score for r in embed_results])
        for r, s in zip(embed_results, embed_scores):
            r.score = s * _BLEND_WEIGHT_EMBEDDING

        kw_scores = _normalize_scores([r.score for r in kw_results])
        for r, s in zip(kw_results, kw_scores):
            r.score = s * _BLEND_WEIGHT_KEYWORDS

        # 合并去重 — 同 chunk_id 取最高分
        merged: dict[str, QueryResult] = {}
        for r in (*embed_results, *kw_results):
            existing = merged.get(r.chunk_id)
            if existing is None or r.score > existing.score:
                merged[r.chunk_id] = r

        # 按分数降序，取 top_k
        results = sorted(merged.values(), key=lambda x: x.score, reverse=True)
        return results[:top_k]

    async def _safe_retrieve_embedding(
        self,
        query: str,
        kb_ids: list[str],
        top_k: int,
    ) -> list[QueryResult]:
        """embedding 检索的容错包装 — 失败时返回空列表（降级为纯关键词）。"""
        if self._embed_model is None:
            return []
        try:
            return await self._retrieve_embedding(query, kb_ids, top_k)
        except Exception as e:
            logger.warning("Embedding retrieval failed, falling back to keywords only: %s", e)
            return []


__all__ = ["RetrievalEngine"]
