# -*- coding: utf-8 -*-
"""
Pydantic AI 聊天智能体和相关模块
"""

# 列举导入模块
from enum import StrEnum
from typing import AsyncGenerator, List, Optional, Union
from uuid import uuid4
from pydantic_ai import Agent as PydanticAIAgent
from pydantic_ai.capabilities import AgentCapability
from pydantic_ai.messages import (
    AgentStreamEvent,
    ModelMessage,
    PartStartEvent,
    TextPart,
    ThinkingPart,
    LoadCapabilityCallPart,
)
from pydantic_ai.models.openai import OpenAIChatModel
from pydantic_ai.output import OutputSpec
from pydantic_ai.providers.openai import OpenAIProvider
from pydantic_ai.run import AgentRunResultEvent
from pydantic import BaseModel, Field


DEFAULT_INSTRUCTIONS: str = """
# 角色
专业友好AI助手，结构化解答各类问题。

# 输出硬性规则
1. 全文强制标准Markdown，禁止纯文本；不要额外说明排版格式，直接输出内容；
2. 层级使用 `#/##/###`，列表用 `-` 无序列表或数字有序列表；
3. 代码块用 ```语言名``` 包裹；
4. 重点内容标注 **粗体**/*斜体*；
5. 思考、工具日志仅输出文本，适配前端折叠面板，禁止输出HTML标签；
6. 内容分点拆分，排版整洁适配前端Markdown渲染。

# 行文要求
语言通俗，逻辑完整简洁，无多余废话。
"""


class Kind(StrEnum):
    """种类"""

    TEXTSTART = "text_start"
    THINKINGSTART = "thinking_start"

    TOOL_NAME = "tool_name"
    TOOL_ARGS = "tool_args"
    TOOL_RETURN = "tool_return"
    FINISHED = "finished"
    ERROR = "error"


class Event(BaseModel):
    """
    片段类
    """

    part_index: Optional[int] = Field(default=None, description="片段索引")
    kind: Kind = Field(..., description="事件种类")
    tool_name: Optional[str] = Field(default=None, description="工具名称")
    args: Optional[LoadCapabilityArgs] = Field(default=None, description="工具参数")
    content: Optional[str] = Field(default=None, description="事件内容")


class Agent:
    """
    基于 Pydantic AI 封装的智能体
    """

    def __init__(
        self,
        chat_id: str,
        instructions: Optional[str] = None,
        output_type: OutputSpec = str,
        capabilities: Optional[List[AgentCapability]] = None,
        retries: int = 1,
    ):
        """
        初始化
        :param chat_id: 聊天唯一标识
        :param instructions: 指令
        :param capabilities: 技能列表，默认为不使用技能
        :param output_type: 输出类型
        :param retries: 重试次数，默认为1次
        :return: 智能体实例
        """
        # 聊天唯一标识
        self.chat_id = chat_id

        # 一次聊天（chat）包含若干论对话（dialog），每轮对话包含用户提示词（user_prompt）和输出（output）。其中，输出包含若干片段（Part）

        # 本轮对话新增消息列表
        self.new_messages: List[ModelMessage] = []

        # 若指令为空则使用默认指令
        if not instructions:
            instructions = DEFAULT_INSTRUCTIONS

        # 初始化智能体
        self.agent = PydanticAIAgent(
            model=OpenAIChatModel(
                model_name="deepseek-v4-flash",
                provider=OpenAIProvider(
                    base_url="https://tokenhub.tencentmaas.com/v1",
                    api_key="sk-D9Y1mCe8VlvNqLuSC4mAjqEwxJ2nW4C0h8a7EPn8kg9RLsHq",
                ),
            ),
            instructions=instructions,
            capabilities=capabilities,
            output_type=output_type,
            retries=retries,
        )

    async def run(
        self,
        user_prompt: str | List[str],
        message_history: Optional[List[ModelMessage]] = None,
    ) -> AsyncGenerator[Event]:
        """
        运行
        :param user_prompt: 用户提示词（用户提示词）
        :param flush_delay: 刷新延迟时长（单位为秒）
        :yield: AsyncGenerator[Event]
        """
        async with self.agent.run_stream_events(
            user_prompt=user_prompt,
            message_history=message_history,
        ) as events:
            async for event in events:
                match event:
                    case PartStartEvent():
                        part = event.part
                        match part:
                            case TextPart(content=content):
                                yield Event(
                                    kind=Kind.TEXTSTART,
                                    part_index=event.index,
                                    content=content,
                                )
                            case ThinkingPart(content=content):
                                yield Event(
                                    kind=Kind.THINKINGSTART,
                                    part_index=event.index,
                                    content=content,
                                )
                            case LoadCapabilityCallPart():
                                yield Event(
                                    kind=Kind.TOOL_NAME,
                                    part_index=event.index,

                                )