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

from asyncpg import Pool
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.models import (
    DEEPSEEK_V4_FLASH_CHAT_MODEL,
    DEEPSEEK_V4_FLASH_CHAT_MODEL_SETTINGS,
)


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

    model_config = {"arbitrary_types_allowed": True}  # 允许任意类型

    account_id: str = Field(..., description="账户唯一标识")
    connection_pool: Pool = Field(..., description="连接池")


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


@account_context.instructions
async def query_account_name(ctx: RunContext[Deps]) -> str:
    """
    查询账户户名
    """
    async with ctx.deps.connection_pool.acquire() as conn:
        account_name = await conn.fetchval(
            "SELECT account_name FROM bank_customer_service WHERE account_id = $1",
            ctx.deps.account_id,
        )
    if account_name is None:
        return "该账户不存在"
    return f"该账户户名是 {account_name!r}"


@account_context.tool(name="查询账户余额")
async def query_account_balance(ctx: RunContext[Deps]) -> str:
    """
    查询账户余额
    """
    async with ctx.deps.connection_pool.acquire() as conn:
        account_balance = await conn.fetchval(
            "SELECT account_balance FROM bank_customer_service WHERE account_id = $1",
            ctx.deps.account_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=[account_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,]:
    """
    运行并流式输出事件
    """
    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:
            print(event)
            yield event