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

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

from application.states.database import DatabaseState
from application.states.models import (
    Conversation,
    Dialog,
    Task,
    TaskResultEvent,
    TaskType,
    Thought,
    thoughts_to_dict,
    usage_to_dict,
)
from application.tasks import run_stream_events


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

    # 当前用户唯一标识
    user_id: str = ""
    # 键为会话唯一标识，值为会话实例的会话字典
    conversations: dict[str, Conversation] = {}  # 按照会话唯一标识顺序排序
    # 当前会话唯一标识
    conversation_id: str = ""

    # 会话历史展示状态，True表示展示，False表示隐藏
    is_conversation_history_shown: bool = False
    # 显示更多的会话唯一标识
    shown_more_conversation_id: str = ""

    # 任务类型
    task_type: TaskType = TaskType.NONE
    # 用户提示词
    user_prompt: str = ""

    # 数据库状态（私有变量，不予序列化）
    _db_state: DatabaseState | None = None

    async def get_db_state(self) -> DatabaseState:
        """
        获取数据库状态
        :return: 数据库状态
        """
        if not self._db_state:
            self._db_state = await self.get_state(DatabaseState)
        return self._db_state

    async def resume(self, user_id: str) -> None:
        """
        恢复当前用户会话状态
        :param user_id: 用户唯一标识
        :return: None
        """
        self.user_id = user_id
        if not self.user_id:
            return

        # 获取数据库状态
        db_state = await self.get_db_state()
        # 获取会话字典
        self.conversations = await db_state.get_conversations(user_id=self.user_id)
        # 若会话字典为空则先创建会话记录再在会话字典中添加会话实例
        if not self.conversations:
            self.conversations.update(
                await db_state.create_conversations_record(user_id=self.user_id)
            )
        # 将最后一个会话作为当前会话并更新会话唯一标识
        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.event
    def set_shown_more_conversation_id(
        self, conversation_id: str, is_shown: bool
    ) -> None:
        """
        设置显示更多的会话唯一标识
        :return: None
        """
        self.shown_more_conversation_id = conversation_id if is_shown else ""

    @rx.event
    async def delete_conversation(self, conversation_id: str) -> None:
        """
        删除会话
        :param conversation_id: 需删除的会话唯一标识
        :return: None
        """
        # 获取数据库状态
        db_state = await self.get_db_state()
        # 先设置会话记录为已删除再在会话字典中删除会话实例
        await db_state.delete_conversations_record(conversation_id=conversation_id)
        del self.conversations[conversation_id]

        # 删除后，若会话字典为空则先创建会话记录再添加会话实例
        if not self.conversations:
            self.conversations.update(
                await db_state.create_conversations_record(user_id=self.user_id)
            )

        # 删除后，若当前会话不存在则将最后一个会话作为当前会话并更新会话唯一标识
        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
        """
        self.conversation_id = conversation_id

    @rx.event
    async def create_conversation(self) -> None:
        """
        创建会话
        :return: None
        """
        # 获取数据库状态
        db_state = await self.get_db_state()
        # 创建会话记录再添加会话实例
        self.conversations.update(
            await db_state.create_conversations_record(
                user_id=self.user_id, description="新会话"
            )
        )
        # 将最后一个会话作为当前会话并更新会话唯一标识
        self.conversation_id = next(reversed(self.conversations.keys()))

    @rx.event
    def set_user_prompt(self, user_prompt: str) -> None:
        """
        设置用户提示词
        :param user_prompt: 用户提示词
        :return: None
        """
        self.user_prompt = user_prompt.strip()

    @rx.var
    def is_user_prompt_sending_disabled(self) -> bool:
        """
        用户提示词发送禁用状态
        :return: bool
        """
        # 当前会话
        conversation = self.conversations.get(self.conversation_id)
        if not conversation:
            return True
        return conversation.is_running or not self.user_prompt

    @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.var
    def dialogs(self) -> dict[str, Dialog]:
        """
        当前会话的对话字典
        :return: 当前会话的对话字典
        """
        # 当前会话
        conversation = self.conversations.get(self.conversation_id)
        if not conversation:
            return {}
        return conversation.dialogs

    @rx.event
    async def handle_user_prompt(self) -> AsyncGenerator[None]:
        """
        处理用户提示词
        :return: AsyncGenerator[None]
        """
        if not self.user_prompt:
            return

        # 当前会话
        conversation = self.conversations[self.conversation_id]
        # 将正在运行设置为是
        conversation.is_running = True
        # 当前对话
        dialog = Dialog(user_prompt=self.user_prompt)
        # 清空用户提示词
        self.user_prompt = ""
        # 添加对话实例
        conversation.dialogs.update({dialog.id: dialog})
        # 强制更新会话并推送前端
        self.conversations[self.conversation_id] = conversation
        yield

        # 获取数据库状态
        db_state = await self.get_db_state()
        # 获取消息历史列表
        message_history = await db_state.get_message_history(
            conversation_id=self.conversation_id
        )

        # 初始化工具调用唯一标识和片段索引映射字典
        tool_call_ids: dict[str, int] = {}
        # 获取运行流式输出事件
        async for event in run_stream_events(
            task_type=self.task_type,
            user_prompt=dialog.user_prompt,
            message_history=message_history,
        ):
            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.thoughts[index] = Thought(
                                type="thinking", content=content
                            )

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

                            dialog.thoughts[index] = Thought(
                                type="tool_search",
                                content="正在生成检索关键词",
                            )

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

                            dialog.thoughts[index] = Thought(
                                type="capability_load",
                                content="正在生成加载参数",
                            )

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

                            dialog.thoughts[index] = Thought(
                                type="tool_call",
                                content="正在生成调用参数",
                            )

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

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

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

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

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

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

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

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

                        # 能力加载
                        case "capability_load":
                            dialog.thoughts[index].content = f"已加载 {content}"

                        # 工具调用
                        case "tool_call":
                            dialog.thoughts[index].content = f"已调用 {content}"

                case TaskResultEvent(task=task, content=content):
                    # 更新任务
                    conversation.task = task
                    dialog.result_output += content

                # ========== 智能体运行结果事件 ==========
                case AgentRunResultEvent(result=result):
                    # 创建对话记录
                    await db_state.create_dialog_record(
                        conversation_id=self.conversation_id,
                        id=dialog.id,
                        user_prompt=self.user_prompt,
                        thoughts=thoughts_to_dict(dialog.thoughts),
                        result_output=dialog.result_output,
                        usage=usage_to_dict(result.usage),
                    )
                    # 创建结果记录
                    await db_state.create_result_record(
                        conversation_id=self.conversation_id,
                        dialog_id=dialog.id,
                        new_messages=result.new_messages(),
                    )

            # 强制更新会话并推送前端
            self.conversations[self.conversation_id] = conversation
            yield

        # 将正在运行设置为否
        conversation.is_running = False
        self.conversations[self.conversation_id] = conversation
        yield

    @rx.event
    async def generate_prd(self) -> None:
        """
        预订航班
        """
        from application.tasks.book_flight import init_task

        # 当前会话
        conversation = self.conversations[self.conversation_id]
        # 初始化预定航班任务
        conversation.task = init_task()

        self.user_prompt = f"帮我找一班从 {conversation.task.deps.origin} 到 {conversation.task.deps.destination} 在 {conversation.task.deps.date} 的航班"

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