# -*- 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,
    PartStartEvent,
    ToolCallPart,
    FunctionToolResultEvent,
    ToolReturnPart,
)
from pydantic_ai.run import AgentRunResultEvent, AgentRunResult

from application.workshop.models import (
    DEEPSEEK_V4_FLASH_MODEL,
    DEEPSEEK_V4_FLASH_MODEL_SETTINGS,
)
from asyncpg import Pool
from application.workshop.events import AgentRuntimeEvent


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

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

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


class Output(BaseModel):
    """
    输出类
    """

    response: str = 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, Output](
    model=DEEPSEEK_V4_FLASH_MODEL,
    model_settings=DEEPSEEK_V4_FLASH_MODEL_SETTINGS,
    deps_type=Deps,
    output_type=Output,
    instructions="你是我方银行的客服智能体，为客户提供客服应答。",
    capabilities=[account_context],
)


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[Output] | 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:
            match event:
                case PartStartEvent(
                    part=part,
                ):
                    match part:
                        case ToolCallPart(tool_name=tool_name):
                            if tool_name.startswith("final_result"):
                                continue

                case FunctionToolResultEvent(part):
                    match part:
                        case ToolReturnPart(
                            tool_name=tool_name,
                        ):
                            if tool_name.startswith("final_result"):
                                continue

                case AgentRunResultEvent(result=result):
                    match result:
                        case AgentRunResult(output=output):
                            match output:
                                case Output(response=response):
                                    yield AgentRuntimeEvent(content=response)

            yield event