# -*- coding: utf-8 -*-
"""
研究
"""
from datetime import datetime
from html import escape
from typing import Any, AsyncGenerator

import httpx
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):
    keyword: str = Field(description="查询语句")


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


@bocha_web_search.tool_plain()
async def web_search(args: Args) -> str:
    """
    联网搜索工具：
    用于查询互联网实时信息、新闻、政策法规、行业报告、企业公告、最新数据。
    """
    if connection_pool is None:
        return "数据库连接池未初始化"

    # 关键词
    keyword = args.keyword

    # 根据关键词查询联网搜索缓存
    async with connection_pool.acquire() as conn:
        results = await conn.fetchval(
            "SELECT results FROM web_search_caches WHERE keyword = $1",
            keyword,
        )
    # 若命中缓存则返回结果，否则先召回再精排、重构搜索结果
    if results is not None:
        return results

    else:
        # --- 召回阶段 --- #
        try:
            async with httpx.AsyncClient(timeout=60) as client:
                response = await client.post(
                    "https://api.bocha.cn/v1/web-search",
                    headers={
                        "Authorization": f"Bearer sk-a9ed73d8b26745799cf970f84d70f512",
                        "Content-Type": "application/json",
                    },
                    json={
                        "query": keyword,
                        "summary": True,  # 显示文本摘要
                        "count": 20,  # 返回结果数
                    },
                )
            # 检查响应状态
            response.raise_for_status()
            # 解析响应载荷
            response_payload = response.json()

        except httpx.HTTPStatusError as exception:
            return (
                f"请求错误 {exception.response.status_code}: {exception.response.text}"
            )
        except Exception as exception:
            return f"请求异常: {str(exception)}"

        if (code := response_payload.get("code")) != 200:
            return f"业务错误 {code}: {response_payload.get("msg", "未知业务错误")}"

        # 召回网页
        webpages = response_payload["data"]["webPages"]["value"]
        if not webpages:
            return "未搜索到相关网页"

        # --- 精排阶段 --- #
        try:
            async with httpx.AsyncClient(timeout=60) as client:
                response = await client.post(
                    "https://api.siliconflow.cn/v1/rerank",
                    headers={
                        "Authorization": f"Bearer sk-xsnuwirjjphhfdbvznfdfjqlinfdlrnlxuhkbbqynfnbhiqz",
                        "Content-Type": "application/json",
                    },
                    json={
                        "model": "BAAI/bge-reranker-v2-m3",
                        "query": keyword,
                        "documents": [
                            called_webpage["summary"] or called_webpage["snippet"]
                            for called_webpage in webpages
                        ],
                        "return_documents": True,  # 返回文档内容
                        "top_n": 10,  # 返回结果数
                    },
                )
            # 检查响应状态
            response.raise_for_status()
            # 解析响应载荷
            response_payload = response.json()

        except httpx.HTTPStatusError as exception:
            return (
                f"请求错误 {exception.response.status_code}: {exception.response.text}"
            )
        except Exception as exception:
            return f"请求异常: {str(exception)}"

        # 精排后网页
        reranked_webpages = [
            (result["index"], result["document"]["text"])
            for result in response_payload["results"]
        ]

        # --- 重构搜索结果 --- #
        results = []
        results.append(f"搜索内容：**{keyword}**，搜索结果：\n")
        for index, summary in reranked_webpages:
            # 根据重排后网页索引获取网页
            webpage = webpages[index]
            # 解析网页的标题
            name = escape(webpage["name"])
            # 解析网页的URL
            url = escape(webpage["url"])
            # 解析网页的网站名称
            site_name = escape(webpage["siteName"])
            # 解析网页的发布时间
            date_published = escape(
                datetime.fromisoformat(webpage["datePublished"]).date().isoformat()
            )

            results.append(
                f"- <a href='{url}'>**{name}**</a>({site_name}・{date_published})：{escape(summary)}"
            )

        results = "\n".join(results)

        # 新增联网搜索缓存
        async with connection_pool.acquire() as conn:
            await conn.execute(
                "INSERT INTO web_search_caches (keyword, results) VALUES ($1, $2)",
                keyword,
                results,
            )

    return results


instruction = """
你是研究助手，可调用联网搜索工具。
## 一、联网判断
需要联网：时效类信息、新闻、声明公告、政策法规、最新数据，以及模型静态知识库无法可靠回答的问题。
无需联网：基础常识、数学计算、代码编写与解释、用户已提供完整材料的分析任务、无时间依赖的静态问题。

【搜索硬约束】
1. 最多**2轮联网搜索**，严禁发起第3轮。一轮定义：单次思考决策发起1条搜索Query；**单次思考步骤，最多只能生成1条搜索query，禁止同一思考周期并行发起多条联网搜索调用**。必须串行执行：执行1次搜索，拿到结果评估之后，才允许在后续思考中决定是否发起下一轮搜索。不允许一步同时调用多次web_search。
2. 发起新一轮搜索前，必须评估已获取的搜索摘要，现有信息足以回答问题则立即停止调用工具。
3. 禁止仅改写同义词、做语义几乎等价的重复搜索；若存在局部信息缺失，优先从已有结果中提取可用内容，不要直接发起新搜索。
4. 默认只执行1次真实联网搜索；拿到搜索结果后评估证据，满足以下任意条件，必须停止搜索：
- 现有信息足以作答；官方一手资料已给出答案；多个可靠来源结论一致；继续搜索仅起到重复确认的作用。
❌禁止以“再次确认、查找更多来源、补充更多资料”等理由发起额外搜索。
✅仅在以下4种场景，才允许执行第2轮搜索：
①第一轮无有效搜索结果；②回答问题必需的关键事实完全缺失；③权威信源之间事实存在明显冲突；④第一轮搜索结果严重偏题。
第二轮搜索必须瞄准明确的信息缺口，禁止全盘复刻上一轮查询。
> 重要：第二轮搜索执行完成后，无论信息是否补齐，都必须终止搜索，不再发起任何web_search调用。

## 二、时间与Query规则
时效判断严格以运行时传入的当前日期、年份、时区为准，禁止依靠模型内置记忆猜测时间。
- 用户未指定年份：**禁止自行向搜索Query中追加年份**。
- 用户明确给出年份：搜索Query必须完整保留用户指定的年份。
- 若检索到网页发布时间晚于运行时当前日期，直接舍弃该结果，不采信。

## 三、信源优先级
1级（最高）：官网、官方文档、官方公告、一手原始资料
2级：权威媒体、专业机构、科研机构
3级：其他可信赖网页
> 媒体转述不等于官方原文，若官方原文存在，优先采信官方原文，不采用媒体转述版本。
官方资料已经可以完整回答问题时，无需再搜索其他来源做重复验证。
自媒体、论坛、转载内容，不可作为重大事实的唯一证据。

## 四、最新内容校验
查询“最新”类问题，不可仅凭网页排序判断时效性；必须核对资料的**发布时间、版本时间，结合运行时提供的当前日期综合判断**。
当存在多份结果时，优先选用发布时间更近、信源等级更高的资料。

## 五、事实规范
- 【事实】：搜索返回资料原文直接支持的内容
- 【推论】：基于已确认事实做合理推导
- 【推测】：证据不足前提下的主观判断
严禁编造URL、日期、新闻事件、数据、官方声明等任何信息；证据不足时直接如实说明，不得强行得出结论。
> 若2轮搜索结束，关键信息仍然缺失，直接告知用户：现有检索资料无法获取目标信息，不编造内容。

## 六、引用规范
引用编号区分搜索轮次：S1‑1、S1‑2 代表第1轮搜索结果；S2‑1、S2‑2 代表第2轮搜索结果。
只引用回答中实际用到的资料。
引用格式：`[参考 S1‑1] https://xxx`，链接必须来源于 **web_search** 工具返回结果，禁止编造链接。

## 七、输出规则
简单事实类问题：优先给出结论，附带关键证据与来源，回答简洁，避免冗余。
复杂研究类问题，按「关键证据 → 分析 → 结论 → 参考来源」结构组织内容；先陈列证据再输出总结，结论不能超出现有资料可支撑的范围。
"""

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


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:
            print(event)
            yield event