# -*- coding: utf-8 -*-
"""
会话状态
"""
from typing import Any, AsyncGenerator, Dict
from typing import Dict

from pydantic_ai import Agent, ThinkingPartDelta
from pydantic_ai.messages import (
    FunctionToolCallEvent,
    FunctionToolResultEvent,
    LoadCapabilityCallPart,
    PartDeltaEvent,
    PartEndEvent,
    PartStartEvent,
    TextPart,
    TextPartDelta,
    ThinkingPart,
    ToolCallPart,
    ToolSearchCallPart,
)
from pydantic_ai.models.openai import OpenAIChatModel
from pydantic_ai.providers.openai import OpenAIProvider
from pydantic_ai.run import AgentRunResultEvent
import reflex as rx

from application.models import Conversation, Dialog, ThoughtNode
from application.states.auth import AuthState
from application.states.database import DatabaseState


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

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

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

# 实例化智能体（因无法序列化故剥离出状态管理）
agent: Agent = Agent(
    model=OpenAIChatModel(
        model_name="deepseek-v4-flash",
        provider=OpenAIProvider(
            base_url="https://tokenhub.tencentmaas.com/v1",
            api_key="sk-D9Y1mCe8VlvNqLuSC4mAjqEwxJ2nW4C0h8a7EPn8kg9RLsHq",
        ),
    ),
    instructions=instructions,
    capabilities=None,
    output_type=str,
    retries=1,
)


class ConversationState(rx.State):
    """
    会话状态
    """

    # 当前用户唯一标识
    user_id: str = ""

    # 会话字典
    conversations: Dict[str, Conversation] = {}
    # 当前会话唯一标识
    conversation_id: str = ""

    # 会话历史展示状态，True表示展示，False表示隐藏
    is_conversation_history_shown: bool = False

    # 正在新建会话状态，True表示正在新建，False表示未正在新建
    is_creating_conversation: bool = False

    @rx.event
    async def on_mount(self):
        """
        页面挂载时初始化会话列表和会话唯一标识
        """
        # 获取当前用户唯一标识
        auth_state = await self.get_state(AuthState)
        self.user_id = auth_state.user_id
        if not self.user_id:
            return

        # 获取数据库状态
        database_state = await self.get_state(DatabaseState)

        # 检索会话字典并赋值
        self.conversations = await database_state.retrieve_conversations(
            user_id=self.user_id
        )
        # 若会话字典为空则先创建会话记录再在会话字典中添加默认会话
        if not self.conversations:
            conversation_id, created_at = (
                await database_state.create_conversations_record(user_id=self.user_id)
            )
            self.conversations[conversation_id] = Conversation(created_at=created_at)

        # 将最后一个会话的唯一标识设置为当前会话唯一标识
        self.conversation_id = next(reversed(self.conversations.keys()))

    @rx.event
    def toggle_conversation_history_shown(self) -> None:
        """
        切换会话历史展示状态
        :return: None
        """
        self.is_conversation_history_shown = not self.is_conversation_history_shown

    @rx.var
    def conversation_history(self) -> Dict[str, Conversation]:
        """
        获取会话历史，用于前端渲染会话历史
        :return: 会话历史
        """
        return dict(reversed(self.conversations.items()))

    @rx.event
    async def create_conversation(self, form_data: Dict[str, Any]) -> None:
        """
        新建会话
        :param form_data: 表单数据
        :return: None
        """
        # 获取会话描述
        description = form_data["description"].strip() or "新会话"

        # 先创建会话记录再在会话字典中添加会话
        database_state = await self.get_state(DatabaseState)
        conversation_id, created_at = await database_state.create_conversations_record(
            user_id=self.user_id, description=description
        )
        self.conversations[conversation_id] = Conversation(
            description=description, created_at=created_at
        )

        # 将所添加的会话唯一标识设置为当前会话唯一标识
        self.conversation_id = conversation_id
        # 将正在新建会话状态设置为已完成
        self.is_creating_conversation = False

    @rx.event
    async def delete_conversation(self, conversation_id: str) -> None:
        """
        删除指定会话
        :param conversation_id: 指定会话唯一标识
        :return: None
        """
        if conversation_id not in self.conversations:
            return

        # 先逻辑删除会话记录再在会话字典中删除会话
        database_state = await self.get_state(DatabaseState)
        await database_state.delete_conversations_record(conversation_id=conversation_id)
        del self.conversations[conversation_id]

        # 删除后，若会话字典为空则先创建会话记录再在会话字典中添加默认会话
        if not self.conversations:
            conversation_id, created_at = (
                await database_state.create_conversations_record(user_id=self.user_id)
            )
            self.conversations[conversation_id] = Conversation(created_at=created_at)

        # 删除后，若当前会话唯一标识不存在则将最后一个会话的唯一标识设置为当前会话唯一标识
        if self.conversation_id not in self.conversations:
            self.conversation_id = next(reversed(self.conversations.keys()))

    @rx.event
    def switch_conversation(self, conversation_id: str) -> None:
        """
        将指定会话的唯一标识设置为当前会话唯一标识
        :param conversation_id: 指定会话唯一标识
        :return: None
        """
        if conversation_id not in self.conversations:
            return
        self.conversation_id = conversation_id

    @rx.event
    def toggle_creating_conversation(self) -> None:
        """
        切换正在新建会话状态
        """
        self.is_creating_conversation = not self.is_creating_conversation

    @rx.var
    def dialog_history(self) -> Dict[str, Dialog]:
        """
        获取当前会话的对话历史，用于前端渲染对话历史
        :return: 当前会话的对话历史
        """
        # 当前会话
        conversation = self.conversations.get(self.conversation_id)
        return conversation.dialogs if conversation else {}

    @rx.var
    def running_status(self) -> bool:
        """
        获取当前会话的运行状态
        :return: 当前会话的运行状态（True 表示正在运行，False 表示运行完成）
        """
        # 当前会话
        conversation = self.conversations.get(self.conversation_id)
        if not conversation:
            return False
        return conversation.is_running

    @rx.event
    async def run(self, form_data: dict[str, Any]) -> AsyncGenerator[None]:
        """
        运行
        :param form_data: 表单数据
        :return: AsyncGenerator
        """
        # 获取问题
        question = form_data.get("question", "").strip()
        if not question:
            return

        # 当前会话
        conversation = self.conversations.get(self.conversation_id)
        if not conversation:
            return

        # 将当前会话的运行状态设置为正在运行
        conversation.is_running = True

        # 获取数据库状态
        database_state = await self.get_state(DatabaseState)

        # 检索运行结果记录并作为消息历史
        message_history = await database_state.retrieve_message_history(
            conversation_id=self.conversation_id
        )

        # 先创建对话记录再在当前对话字典中新增对话
        dialog_id = await database_state.create_dialogs_record(
            conversation_id=self.conversation_id, question=question
        )
        dialog = conversation.dialogs.setdefault(dialog_id, Dialog(question=question))

        # 初始化工具调用唯一标识和片段索引映射字典
        tool_call_ids: Dict[str, int] = {}

        yield  # 通知前端渲染

        async with agent.run_stream_events(
            conversation_id=self.conversation_id,
            user_prompt=question,
            message_history=message_history,
        ) as events:
            async for event in events:
                match event:
                    # ========== 开始事件 ==========
                    case PartStartEvent(
                        index=index,
                        part=part,
                        previous_part_kind=previous_part_kind,
                    ):
                        match part:
                            # 思考分片开始事件
                            case ThinkingPart(content=content):
                                # 若上一分片种类为空则将思考状态设置为正在思考、思考折叠面板展开状态设置为展开
                                if not previous_part_kind:
                                    dialog.is_thinking = True
                                    dialog.is_collapse_expanded = True

                                dialog.thought_nodes[index] = ThoughtNode(
                                    kind="thinking", content=content
                                )
                                yield

                            # 工具检索分片开始事件
                            case ToolSearchCallPart(tool_call_id=tool_call_id):
                                # 创建工具调用唯一标识与片段索引映射
                                tool_call_ids[tool_call_id] = index

                                dialog.thought_nodes[index] = ThoughtNode(
                                    kind="tool_search",
                                    content="正在生成检索关键词",
                                )
                                yield

                            # 能力加载分片开始事件
                            case LoadCapabilityCallPart(tool_call_id=tool_call_id):
                                tool_call_ids[tool_call_id] = index

                                dialog.thought_nodes[index] = ThoughtNode(
                                    kind="capability_load",
                                    content="正在生成加载参数",
                                )
                                yield

                            # 工具调用分片开始事件
                            case ToolCallPart(tool_call_id=tool_call_id):
                                tool_call_ids[tool_call_id] = index

                                dialog.thought_nodes[index] = ThoughtNode(
                                    kind="tool_call",
                                    content="正在生成调用参数",
                                )
                                yield

                            # 文本分片开始事件
                            case TextPart(content=content):
                                dialog.answer = content
                                yield

                    # ========== 增量事件 ==========
                    case PartDeltaEvent(index=index, delta=delta):
                        match delta:
                            # 思考分片增量事件
                            case ThinkingPartDelta(
                                content_delta=content_delta,
                            ):
                                dialog.thought_nodes[index].content += (
                                    content_delta or ""
                                )
                                yield

                            # 文本分片增量事件
                            case TextPartDelta(
                                content_delta=content_delta,
                            ):
                                dialog.answer += content_delta
                                yield

                    # ========== 结束事件 ==========
                    case PartEndEvent(
                        index=index,
                        part=part,
                        next_part_kind=next_part_kind,
                    ):
                        match part:
                            # 思考分片结束事件
                            case ThinkingPart(part_kind=part_kind, content=content):
                                # 若下一分片种类为文本则将思考状态设置为思考完成、思考面板展开状态设置为折叠
                                if next_part_kind == "text":
                                    dialog.is_thinking = False
                                    dialog.is_collapse_expanded = False
                                    yield

                    # ========== 函数工具调用事件 ==========
                    case FunctionToolCallEvent(tool_call_id=tool_call_id, part=part):
                        # 获取分片索引
                        index = tool_call_ids[tool_call_id]
                        match dialog.thought_nodes[index].kind:
                            # 工具检索
                            case "tool_search":
                                dialog.thought_nodes[index].content = "正在检索"
                                yield

                            # 能力加载
                            case "capability_load":
                                dialog.thought_nodes[index].content = (
                                    f"正在加载能力 {part.tool_name}"
                                )
                                yield

                            # 工具调用
                            case "tool_call":
                                dialog.thought_nodes[index].content = (
                                    f"正在调用工具 {part.tool_name}"
                                )
                                yield

                    # ========== 函数工具结果事件 ==========
                    case FunctionToolResultEvent(
                        tool_call_id=tool_call_id,
                        content=content,
                    ):
                        index = tool_call_ids[tool_call_id]
                        match dialog.thought_nodes[index].kind:
                            # 工具检索
                            case "tool_search":
                                dialog.thought_nodes[index].content = (
                                    content if isinstance(content, str) else ""
                                )  # 暂仅考虑文本内容
                                yield

                            # 能力加载
                            case "capability_load":
                                dialog.thought_nodes[index].content = "已加载"
                                yield

                            # 工具调用
                            case "tool_call":
                                dialog.thought_nodes[index].content = f"已调用"
                                yield

                    # ========== 智能体运行结果事件 ==========
                    case AgentRunResultEvent(result=result):
                        # 更新对话记录
                        await database_state.update_dialog_record(
                            dialog_id=dialog_id,
                            thought_nodes=dialog.thought_nodes,
                            answer=dialog.answer,
                        )
                        # 创建运行结果记录
                        await database_state.create_run_result_record(
                            conversation_id=self.conversation_id,
                            dialog_id=dialog_id,
                            new_messages=result.new_messages(),
                        )
                        # 将当前会话的运行状态设置为运行完成
                        conversation.is_running = False
                        yield

    @rx.event
    def toggle_collapse(self, dialog_id: str) -> None:
        """
        展开/折叠思考折叠面板
        """
        # 指定运行
        dialog = self.conversations[self.conversation_id].dialogs[dialog_id]
        dialog.is_collapse_expanded = not dialog.is_collapse_expanded