# -*- coding: utf-8 -*-
"""
研究
"""
from typing import AsyncGenerator

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

from application.workshop.models import DEEPSEEK_V4_FLASH_RESPONSE_MODEL


class Args(BaseModel):
    query: str = Field(description="查询语句")


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


@bocha_web_search.tool_plain()  # 普通工具：不使用运行上下文
async def web_search(args: Args) -> str:
    """
    联网搜索工具：
    用于查询互联网实时信息、新闻、政策、行业报告、企业公告、最新数据。
    适合查询模型静态知识库不存在的外部网页信息。
    回答时必须标注引用来源URL，禁止编造链接。
    不要用于简单常识、数学计算、代码编写等无需联网的任务。
    """
    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": args.query,
                    "summary": True,  # 显示文本摘要
                },
            )
            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", "未知业务错误")}"

    web_pages = response_payload.get("data", {}).get("webPages", {}).get("value", [])
    if not web_pages:
        return "未搜索到相关网页"

    result = []  # 搜索结果
    for idx, web_page in enumerate(web_pages, 1):
        # 解析网页的标题
        name = web_page.get("name", "")
        # 解析网页的URL
        url = web_page.get("url", "")
        # 解析网页的网站名称
        site_name = web_page.get("siteName", "")
        # 解析网页的发布时间
        date_published = web_page.get("datePublished") or web_page.get(
            "dateLastCrawled", ""
        )
        # 优先解析网页内容的文本摘要，若文本摘要为空则解析简短描述
        summary = web_page.get("summary") or web_page.get("snippet", "")

        result.append(
            f"### 搜索结果 {idx}\n"
            f"- 标题：{name}\n"
            f"- URL：{url}\n"
            f"- 来源网站：{site_name}\n"
            f"- 发布时间：{date_published}\n"
            f"- 摘要：{summary}"
        )
    return "\n\n".join(result)


instruction = """
你是研究助手，可调用联网搜索工具。
1. 事实、时效、外部资料类问题必须搜索；静态常识可直接回答，支持多次检索、改写query。
2. 优先采用一手资料、官方权威信源，慎用自媒体与转载内容。
3. 所有事实描述必须标注来源链接，引用格式：`[参考{序号}] {URL}`，序号匹配搜索结果编号。
4. 严格区分：【事实】来自网页原文；【推论】基于事实推导；资料不足时标注【推测】，不可当作事实。
5. 严禁编造链接和事实；证据不足直接说明，不要强行下结论。
6. 组织回答：先陈列证据，再总结结论，结论不能超出资料支撑范围。
"""

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


async def run_stream_events(
    user_prompt: str | None,
    message_history: list[ModelMessage],
    usage: RunUsage,
) -> AsyncGenerator[AgentStreamEvent | AgentRunResultEvent, None]:
    """
    运行并流式输出事件
    :param user_prompt: 用户提示词
    :param message_history: 消息历史
    :param usage: 使用量
    :return: AsyncGenerator
    """
    async with agent.run_stream_events(
        user_prompt=user_prompt,
        message_history=message_history,
        usage=usage,
    ) as events:
        async for event in events:
            print(event)
            yield event