# -*- coding: utf-8 -*-
"""
知识库检索
"""
from asyncio import sleep
from typing import Any, AsyncGenerator

from httpx import AsyncClient, HTTPStatusError
from pydantic import BaseModel, Field
from pydantic_ai import Agent, ModelMessage
from pydantic_ai.messages import AgentStreamEvent
from pydantic_ai.run import AgentRunResultEvent
from pydantic_ai.toolsets import FunctionToolset
from pydantic_ai.usage import RunUsage

from application.workshop.database import connection_pool
from application.workshop.models import DEEPSEEK_V4_FLASH_RESPONSE_MODEL


class Deps(BaseModel):
    """
    依赖类
    """

    ...


class Args(BaseModel):
    text: str = Field(..., description="查询文本")
    keywords: list[str] = Field(default_factory=list, description="查询关键词列表")


# 实例化工具集
toolset = FunctionToolset()


@toolset.tool_plain()
async def knowledge_base_retrieve(args: Args) -> str:
    """
    知识库检索：
    用于基于知识库检索相关信息。
    """
    if connection_pool is None:
        return "数据库连接池未初始化"

    # --- 召回，先根据查询文本获取嵌入向量再基于查询关键词和文本向量混合检索 --- #
    request_count = 0  # 请求次数
    response_payload: dict | None = None
    while request_count < 3:
        request_count += 1
        try:
            async with AsyncClient(timeout=60) as client:
                response = await client.post(
                    "https://api.siliconflow.cn/v1/embeddings",
                    headers={
                        "Authorization": "Bearer sk-xsnuwirjjphhfdbvznfdfjqlinfdlrnlxuhkbbqynfnbhiqz",
                        "Content-Type": "application/json",
                    },
                    json={
                        "model": "BAAI/bge-m3",
                        "input": args.text,
                    },
                )
            # 检查响应状态
            response.raise_for_status()
            # 解析响应载荷
            response_payload = response.json()
            break

        except HTTPStatusError as exception:
            # 若触发限速则等待若干秒后重试
            if exception.response.status_code == 429:
                await sleep(2**request_count)
                continue

            raise RuntimeError(
                f"HTTPError {exception.response.status_code}: {exception.response.text}"
            )

        except Exception as exception:
            raise RuntimeError(f"Exception: {str(exception)}")

    if response_payload is None:
        raise RuntimeError("response_payload is None")
    text_vector = response_payload["data"][0]["embedding"]

    # 混合检索
    async with connection_pool.acquire() as connection:
        results = await connection.fetch(
            """
            WITH bm25_retrieve_results AS (
                SELECT
                    id,
                    document_id,
                    index,
                    text,
                    return_text,
                    text_fingerprint,
                    ROW_NUMBER() OVER (ORDER BY paradedb.score(id) DESC) AS bm25_retrieve_rank
                FROM document_chunks
                WHERE text @@@ $1
                ORDER BY paradedb.score(id) DESC
                LIMIT 50
            ),
            vector_retrieve_results AS (
                SELECT
                    id,
                    document_id,
                    index,
                    text,
                    return_text,
                    text_fingerprint,
                    ROW_NUMBER() OVER (ORDER BY text_vector <-> $2::vector(1024)) AS vector_retrieve_rank
                FROM document_chunks
                ORDER BY text_vector <-> $2::vector(1024)
                LIMIT 50
            ),
            merged_results AS (
                SELECT
                    COALESCE(bm25_retrieve_results.id, vector_retrieve_results.id) AS id,
                    COALESCE(bm25_retrieve_results.document_id, vector_retrieve_results.document_id) AS document_id,
                    COALESCE(bm25_retrieve_results.index, vector_retrieve_results.index) AS index,
                    COALESCE(bm25_retrieve_results.text, vector_retrieve_results.text) AS text,
                    COALESCE(bm25_retrieve_results.return_text, vector_retrieve_results.return_text) AS return_text,
                    COALESCE(bm25_retrieve_results.text_fingerprint, vector_retrieve_results.text_fingerprint) AS text_fingerprint,
                    COALESCE(1.0/(bm25_retrieve_rank + 60), 0) + COALESCE(1.0/(vector_retrieve_rank + 60), 0) AS rank_score
                FROM bm25_retrieve_results
                FULL OUTER JOIN vector_retrieve_results USING(id)
            ),
            deduplication_results AS (
                SELECT
                    *,
                    ROW_NUMBER() OVER (PARTITION BY text_fingerprint ORDER BY rank_score DESC) AS duplicate_index
                FROM merged_results
            )
            SELECT
                id,
                document_id,
                index,
                text,
                return_text,
                rank_score
            FROM deduplication_results
            WHERE duplicate_index = 1
            ORDER BY rank_score DESC
            LIMIT 50;
            """,
            " ".join(args.keywords) if args.keywords else args.text,
            "[" + ",".join(str(x) for x in text_vector) + "]",
        )
    if not results:
        return "未检索到相关结果"
    recall_results: list[dict] = [dict(row) for row in results]  # 召回结果

    # --- 重排 --- #
    request_count = 0  # 请求次数
    response_payload: dict | None = None
    while request_count < 3:
        request_count += 1
        try:
            async with AsyncClient(timeout=60) as client:
                response = await client.post(
                    "https://api.siliconflow.cn/v1/rerank",
                    headers={
                        "Authorization": "Bearer sk-xsnuwirjjphhfdbvznfdfjqlinfdlrnlxuhkbbqynfnbhiqz",
                        "Content-Type": "application/json",
                    },
                    json={
                        "model": "BAAI/bge-reranker-v2-m3",
                        "query": args.text,
                        "documents": [
                            result["return_text"] for result in recall_results
                        ],
                        "return_documents": False,
                        "top_n": 10,
                    },
                )
                # 检查响应状态
                response.raise_for_status()
                # 解析响应载荷
                response_payload = response.json()
                break

        except HTTPStatusError as exception:
            # 若触发限速则等待若干秒后重试
            if exception.response.status_code == 429:
                await sleep(2**request_count)
                continue

            raise RuntimeError(
                f"HTTPError {exception.response.status_code}: {exception.response.text}"
            )

        except Exception as exception:
            await sleep(2**request_count)
            continue

    if response_payload is None:
        reranked_results: list[dict] = recall_results[:10]
    else:
        reranked_results = [
            recall_results[result["index"]]
            for result in response_payload["results"]
            if result["relevance_score"] >= 0.6
        ]

    if not reranked_results:
        return "未检索到相关结果"

    # --- 重构结果 --- #
    refactored_results: list[str] = [f"检索内容：**{args.text}**，检索结果：\n"]
    for result in reranked_results:
        refactored_results.append(f"{result['return_text']}\n")
    return "\n".join(refactored_results)


instruction = """
你是内部知识库问答助手，可调用 knowledge_base_retrieve 工具检索已收录的文档片段。

## 一、检索判断
需要调用 knowledge_base_retrieve：问题涉及项目内部资料、技术文档、产品规格、历史方案、代码库说明、内部流程等知识库已收录的主题。
无需调用：纯闲聊、基础常识、数学计算、代码编写与解释、用户已在对话中提供完整材料的分析任务。
【检索硬约束】
1. 最多**2轮知识库检索**，严禁发起第3轮。一轮定义：单次思考决策发起1次检索；**单次思考步骤最多只能生成1次工具调用，禁止同一思考周期并行发起多次检索**。必须串行：执行1次检索，拿到结果评估之后，才允许在后续思考中决定是否发起下一轮。
2. 发起新一轮检索前，必须评估已获取的文档片段，现有片段足以回答问题则立即停止调用工具。
3. 禁止仅改写同义词、做语义几乎等价的重复检索；若存在局部信息缺失，优先从已有片段中提取可用内容，不要直接发起新检索。
4. 默认只执行1次真实检索；拿到结果后评估证据，满足以下任意条件，必须停止检索：
- 现有片段足以作答；多个片段交叉印证同一结论；继续检索仅起到重复确认的作用。
❌禁止以"再查一下""补充更多文档""确认一下"等理由发起额外检索。
✅仅在以下4种场景，才允许执行第2轮检索：
①第一轮无有效检索结果；②回答问题必需的关键事实在已有片段中完全缺失；③第一轮返回的片段与用户问题明显偏题；④第一轮只召回了泛泛结果，缺少问题指向的具体实体/术语。
第二轮检索必须更换关键词组瞄准明确的信息缺口，禁止原样重发上一轮查询。
> 重要：第二轮检索执行完成后，无论信息是否补齐，都必须终止检索，不再发起任何 knowledge_base_retrieve 调用。

## 二、调用工具前的改写规范
调用 knowledge_base_retrieve 时，禁止直接把用户原始疑问句透传，必须先完成：
1. **去疑问化**：把"支付模块怎么设计的？"改写成陈述句"支付模块设计"填入 text。
2. **指代消解**：结合对话历史，把"那个功能""之前那个报错""他说的方案"替换成具体名词后填入 text。
3. **关键词提炼**：从问题中提取 2–5 个核心名词/术语/专有名词填入 keywords，例如：
   - 用户问"退款流程里签名校验为什么会失败"
   - → text="退款流程签名校验失败"
   - → keywords=["退款流程", "签名校验"]
4. keywords 只保留名词、术语、产品名、技术栈名，禁止包含虚词、疑问词、动词短语。
5. 若问题本身就是一个名词短语（如"支付模块文档"），text 与 keywords 可一致。

## 三、事实规范
- 【事实】：检索返回文档片段原文直接支持的内容。
- 【推论】：基于已确认片段做的合理推导。
- 【推测】：证据不足前提下的主观判断，必须显式标注。
严禁编造文档片段中不存在的结论、字段、参数、接口名、日期、数据。
> 知识库未收录相关内容时，直接告知用户"知识库中未找到相关资料"，不得用模型自身记忆强行作答或补全。
> 若2轮检索结束关键信息仍缺失，直接说明现有文档无法回答该问题，不编造。

## 四、引用规范
回答中引用知识库片段时，必须标注来源，格式：`[文档 document_id 片段 chunk_index]`，编号必须对应检索结果中实际返回的 document_id 与 chunk_index，禁止编造编号。
只引用回答中实际用到的片段；不要把整段召回文本原样堆砌给用户，要先总结再引用。

## 五、输出规则
简单事实类问题：优先给出结论，附带关键片段与来源标注，回答简洁。
复杂问题（多文档交叉、方案对比、流程梳理）：按「关键证据 → 分析 → 结论 → 来源标注」结构组织；先陈列依据再输出总结，结论不能超出现有片段可支撑的范围。
"""

# 智能体
agent = Agent[Deps, str](
    model=DEEPSEEK_V4_FLASH_RESPONSE_MODEL,
    deps_type=Deps,
    output_type=str,  # 若输出类型为 str，则智能体输出结果时不进行输出类型校验。若定义输出类型，需考虑处理相应事件
    instructions=instruction,
    toolsets=[toolset],
)


async def run_stream_events(
    deps: Any,
    user_prompt: str | None,
    message_history: list[ModelMessage],
    usage: RunUsage,
) -> AsyncGenerator[AgentStreamEvent | AgentRunResultEvent, None]:
    """
    运行并流式输出事件
    :param user_prompt: 用户提示词
    :param message_history: 消息历史
    :param usage: 使用量
    :return: AsyncGenerator
    """
    if not isinstance(deps, Deps):
        deps = Deps.model_validate(deps)

    async with agent.run_stream_events(
        deps=deps,
        user_prompt=user_prompt,
        message_history=message_history,
        usage=usage,
    ) as events:
        async for event in events:
            yield event