# -*- coding: utf-8 -*-
"""
银行客户服务
"""
from typing import Any, AsyncGenerator

from pydantic import BaseModel, Field
from pydantic_ai import (
    Agent,
    Capability,
    DeferredToolResults,
    ModelMessage,
    RunContext,
    RunUsage,
    UsageLimits,
)
from pydantic_ai.messages import AgentStreamEvent
from pydantic_ai.run import AgentRunResultEvent

from application.workshop.database import connection_pool
from application.workshop.models import (
    DEEPSEEK_V4_FLASH_CHAT_MODEL,
    DEEPSEEK_V4_FLASH_CHAT_MODEL_SETTINGS,
)


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

    user_id: str = Field(..., description="用户唯一标识")


# 用户上下文能力
user_context = Capability[Deps](
    id="user-context",
    description="提供用户信息，用于客服应答",
    defer_loading=False,  # 延迟加载置为否(立即加载)
)


@user_context.instructions
async def query_user_name(ctx: RunContext[Deps]) -> str:
    """
    查询用户姓名
    """
    if connection_pool is None:
        return "数据库连接池未初始化"

    async with connection_pool.acquire() as conn:
        user_name = await conn.fetchval(
            "SELECT user_name FROM bank_customer_service WHERE user_id = $1",
            ctx.deps.user_id,
        )
    if user_name is None:
        return "该用户不存在"
    return f"用户姓名是 {user_name!r}"


@user_context.tool(name="查询账户余额")
async def query_balance(ctx: RunContext[Deps]) -> str:
    """
    查询账户余额
    """
    if connection_pool is None:
        return "数据库连接池未初始化"

    async with connection_pool.acquire() as conn:
        account_balance = await conn.fetchval(
            "SELECT account_balance FROM bank_customer_service WHERE user_id = $1",
            ctx.deps.user_id,
        )
    if account_balance is None:
        return "该用户不存在"
    return f"该用户的账户余额: {account_balance:.2f}"


# 智能体
agent = Agent[Deps, str](
    model=DEEPSEEK_V4_FLASH_CHAT_MODEL,
    model_settings=DEEPSEEK_V4_FLASH_CHAT_MODEL_SETTINGS,
    deps_type=Deps,
    output_type=str,  # 若输出类型为 str，则智能体输出结果时不进行输出类型校验。若定义输出类型，需考虑处理相应事件
    instructions="你是我方银行的客服智能体，为客户提供客服应答。",  # 优先使用 instructions
    capabilities=[user_context],
)

"""
此为设置输出类型为 Output，智能体输出结果时需校验输出类型，同工具调用：
PartStartEvent(index=0, part=ToolCallPart(tool_name='final_result', args='', tool_call_id='call_3ab0ed20b1114fa0be726723'))
FinalResultEvent(tool_name='final_result', tool_call_id='call_3ab0ed20b1114fa0be726723')
PartDeltaEvent(index=0, delta=ToolCallPartDelta(args_delta='', tool_call_id='call_3ab0ed20b1114fa0be726723'))
PartDeltaEvent(index=0, delta=ToolCallPartDelta(args_delta='{', tool_call_id='call_3ab0ed20b1114fa0be726723'))
PartDeltaEvent (index=0, delta=ToolCallPartDelta (args_delta='"response": " 您好 ', tool_call_id='call_3ab0ed20b1114fa0be726723'))
PartDeltaEvent (index=0, delta=ToolCallPartDelta (args_delta='，刘 ', tool_call_id='call_3ab0ed20b1114fa0be726723'))
PartDeltaEvent (index=0, delta=ToolCallPartDelta (args_delta=' 弼仁 ', tool_call_id='call_3ab0ed20b1114fa0be726723'))
PartDeltaEvent (index=0, delta=ToolCallPartDelta (args_delta=' 先生。', tool_call_id='call_3ab0ed20b1114fa0be726723'))
PartDeltaEvent (index=0, delta=ToolCallPartDelta (args_delta=' 您账户 ', tool_call_id='call_3ab0ed20b1114fa0be726723'))
PartDeltaEvent (index=0, delta=ToolCallPartDelta (args_delta=' 的 ', tool_call_id='call_3ab0ed20b1114fa0be726723'))
PartDeltaEvent (index=0, delta=ToolCallPartDelta (args_delta=' 当前余额 ', tool_call_id='call_3ab0ed20b1114fa0be726723'))
PartDeltaEvent (index=0, delta=ToolCallPartDelta (args_delta=' 为 **', tool_call_id='call_3ab0ed20b1114fa0be726723'))
PartDeltaEvent(index=0, delta=ToolCallPartDelta(args_delta='9.', tool_call_id='call_3ab0ed20b1114fa0be726723'))
PartDeltaEvent (index=0, delta=ToolCallPartDelta (args_delta='99 元 ', tool_call_id='call_3ab0ed20b1114fa0be726723'))
PartDeltaEvent(index=0, delta=ToolCallPartDelta(args_delta='**。', tool_call_id='call_3ab0ed20b1114fa0be726723'))
PartDeltaEvent (index=0, delta=ToolCallPartDelta (args_delta=' 请问还有什么 ', tool_call_id='call_3ab0ed20b1114fa0be726723'))
PartDeltaEvent (index=0, delta=ToolCallPartDelta (args_delta=' 可以帮 ', tool_call_id='call_3ab0ed20b1114fa0be726723'))
PartDeltaEvent (index=0, delta=ToolCallPartDelta (args_delta=' 您的 ', tool_call_id='call_3ab0ed20b1114fa0be726723'))
PartDeltaEvent (index=0, delta=ToolCallPartDelta (args_delta=' 吗？"}', tool_call_id='call_3ab0ed20b1114fa0be726723'))
PartEndEvent (index=0, part=ToolCallPart (tool_name='final_result', args='{"response": " 您好，刘弼仁先生。您账户的当前余额为 **9.99 元**。请问还有什么可以帮您的吗？"}', tool_call_id='call_3ab0ed20b1114fa0be726723'))
OutputToolCallEvent (part=ToolCallPart (tool_name='final_result', args='{"response": " 您好，刘弼仁先生。您账户的当前余额为 **9.99 元**。请问还有什么可以帮您的吗？"}', tool_call_id='call_3ab0ed20b1114fa0be726723'), args_valid=True)
OutputToolResultEvent(part=ToolReturnPart(tool_name='final_result', content='Final result processed.', tool_call_id='call_3ab0ed20b1114fa0be726723', timestamp=datetime.datetime(2026, 9, 17, 10, 34, 0, 799329, tzinfo=datetime.timezone.utc)))

此为设置输出类型为 str，智能体输出结果时不进行输出类型校验：
PartStartEvent(index=0, part=TextPart(content='您好'))
FinalResultEvent(tool_name=None, tool_call_id=None)
PartDeltaEvent(index=0, delta=TextPartDelta(content_delta='，刘'))
PartDeltaEvent(index=0, delta=TextPartDelta(content_delta='弼仁'))
PartDeltaEvent(index=0, delta=TextPartDelta(content_delta='先生。'))
PartDeltaEvent(index=0, delta=TextPartDelta(content_delta='您账户'))
PartDeltaEvent(index=0, delta=TextPartDelta(content_delta='的'))
PartDeltaEvent(index=0, delta=TextPartDelta(content_delta='当前余额'))
PartDeltaEvent(index=0, delta=TextPartDelta(content_delta='为 **'))
PartDeltaEvent(index=0, delta=TextPartDelta(content_delta='9.'))
PartDeltaEvent(index=0, delta=TextPartDelta(content_delta='99'))
PartDeltaEvent(index=0, delta=TextPartDelta(content_delta='元**'))
PartDeltaEvent(index=0, delta=TextPartDelta(content_delta='。请问'))
PartDeltaEvent(index=0, delta=TextPartDelta(content_delta='还有什么'))
PartDeltaEvent(index=0, delta=TextPartDelta(content_delta='可以帮'))
PartDeltaEvent(index=0, delta=TextPartDelta(content_delta='您的吗'))
PartDeltaEvent(index=0, delta=TextPartDelta(content_delta='？'))
PartEndEvent(index=0, part=TextPart(content='您好，刘弼仁先生。您账户的当前余额为 **9.99元**。请问还有什么可以帮您的吗？'))
"""


async def run_stream_events(
    deps: Any,
    user_prompt: str | None,
    message_history: list[ModelMessage],
    usage: RunUsage,
    usage_limits: UsageLimits | None = None,
    deferred_tool_results: DeferredToolResults | None = None,
) -> AsyncGenerator[AgentStreamEvent | AgentRunResultEvent[str] | None,]:
    """
    运行并流式输出事件
    :param user_prompt: 用户提示词
    :param message_history: 消息历史记录
    :param usage: 运行使用情况
    :param usage_limits: 运行使用限制
    :param deferred_tool_results: 延迟工具结果
    :return: 异步生成器，生成智能体流式事件或运行结果事件
    """
    if not isinstance(deps, Deps):
        deps = Deps.model_validate(deps)

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