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

from pydantic_ai import DeferredToolRequests, DeferredToolResults, ToolApproved
from pydantic_ai.messages import (
    FunctionToolResultEvent,
    PartDeltaEvent,
    PartEndEvent,
    PartStartEvent,
    TextPart,
    TextPartDelta,
    ThinkingPart,
    ThinkingPartDelta,
    ToolCallPart,
    ToolReturnPart,
)
from pydantic_ai.run import AgentRunResultEvent
import reflex as rx

from application.states.database import DatabaseState
from application.states.models import (
    Approval,
    Conversation,
    ConversationHistoryItem,
    Message,
    MessageHistoryItem,
    MessageType,
    Workflow,
    WorkflowType,
)
from application.workshop.events import AgentRuntimeEvent


def format_conversation_created_at(created_at: datetime) -> str:
    """
    格式化会话创建日期时间
    :param created_at: 创建日期时间
    :return: 格式化后的日期时间字符串
    """
    match (datetime.now().date() - created_at.date()).days:
        case 0:
            formatted_created_at = f"{created_at.strftime('%H:%M')}"
        case 1:
            formatted_created_at = f"昨天 {created_at.strftime('%H:%M')}"
        case _:
            formatted_created_at = created_at.strftime("%Y-%m-%d %H:%M")
    return formatted_created_at


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

    # 当前用户唯一标识
    user_id: str = ""
    # 当前用户的会话字典（私有变量）

    # 当前用户的会话字典
    conversations: dict[str, Conversation] = {}  # 键为会话唯一标识，值为会话
    # 当前会话唯一标识
    conversation_id: str = ""

    # 显示会话历史，True 表示显示，False 表示隐藏
    is_conversation_history_shown: bool = False
    # 悬停气泡显示的会话历史项唯一标识（一个会话只允许存在一个悬停气泡会话历史项）
    popover_shown_conversation_history_item_id: str = ""

    # 当前数据库状态（私有变量）
    # 私有变量：reflex 约定以 _ 开头的变量为私有变量，后端不序列化，前端不可使用
    _db_state: DatabaseState | None = None

    async def get_db_state(self) -> DatabaseState:
        """
        获取当前数据库状态
        :return: 当前数据库状态
        """
        # 若数据库状态为 None 则获取数据库状态，否则直接返回当前数据库状态
        if not self._db_state:
            self._db_state = await self.get_state(DatabaseState)
        return self._db_state

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

        # 获取当前数据库状态
        db_state = await self.get_db_state()
        # 查询当前用户的会话字典
        self.conversations = await db_state.read_conversations(self.user_id)
        # 若当前用户的会话字典为空则先创建会话记录再添加会话
        if not self.conversations:
            self.conversations.update(
                await db_state.create_conversation_record(self.user_id)
            )
        # 将最后一个会话的唯一标识作为当前会话唯一标识
        self.conversation_id = next(reversed(self.conversations.keys()))

        # 当前会话
        conversation = self.conversations[self.conversation_id]
        # 最后一个消息
        message = next(reversed(conversation.messages.values()))
        # 若最后一个消息包含审批且审批结果为 None 则将会话正在运行设置为是、消息正在运行设置为是
        if (
            isinstance(approval := message.approval, Approval)
            and approval.result is None
        ):
            conversation.is_running = True
            message.is_running = True

    @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) -> list[ConversationHistoryItem]:
        """
        获取会话历史项列表（按照会话创建日期时间降序排序）
        :return: 会话历史项列表
        """
        return [
            ConversationHistoryItem(
                id=conversation.id,
                description=conversation.description,
                created_at=format_conversation_created_at(conversation.created_at),
            )
            for conversation in reversed(self.conversations.values())
        ]

    @rx.event
    def set_conversation_history_item_popover_shown(
        self, item_id: str, is_shown: bool
    ) -> None:
        """
        设置会话历史项悬停气泡显示 / 隐藏
        :return: None
        """
        self.popover_shown_conversation_history_item_id = item_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.update_conversation_record(conversation_id, is_deleted=True)
        del self.conversations[conversation_id]

        # 删除后，若当前用户的会话字典为空则创建会话
        if not self.conversations:
            self.conversations.update(
                await db_state.create_conversation_record(self.user_id)
            )

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

    @rx.event
    def set_conversation(self, conversation_id: str) -> None:
        """
        设置当前会话
        :param conversation_id: 会话唯一标识
        :return: None
        """
        self.conversation_id = conversation_id

    @rx.var
    def message_history(self) -> list[MessageHistoryItem]:
        """
        获取当前会话的消息历史
        :return: 当前会话的消息历史
        """
        # 当前会话
        conversation = self.conversations.get(self.conversation_id)
        if not conversation:
            return []
        message_history: list[MessageHistoryItem] = []
        for message in conversation.messages.values():
            message_history.append(
                MessageHistoryItem(
                    **message.model_dump(),
                    has_approval=isinstance(approval := message.approval, Approval)
                    and approval.result is None,
                    approval_content=getattr(approval, "content", ""),
                )
            )
        return message_history

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

    @rx.var
    def user_prompt(self) -> str:
        """
        用户提示词
        :return: 用户提示词
        """
        # 当前会话
        conversation = self.conversations.get(self.conversation_id)
        if not conversation:
            return ""
        return conversation.user_prompt

    @rx.event
    def set_user_prompt(self, user_prompt: str) -> None:
        """
        设置用户提示词
        :param user_prompt: 用户提示词
        :return: None
        """
        # 当前会话
        conversation = self.conversations.get(self.conversation_id)
        if not conversation:
            return

        if not user_prompt.strip():
            return

        conversation.user_prompt = user_prompt.strip()

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

    @rx.event
    async def run(
        self,
        is_resume_run: bool = False,
        message_id: str | None = None,
    ) -> AsyncGenerator[None]:
        """
        运行
        :param is_resume_run: 恢复运行
        :param message_id: 消息唯一标识
        :return: AsyncGenerator
        """
        # 当前会话
        conversation = self.conversations.get(self.conversation_id)
        if not conversation:
            return

        # 获取当前数据库状态
        db_state = await self.get_db_state()

        user_prompt = None  # 用户提示词
        deferred_tool_results = None  # 延迟工具结果
        # 若为恢复运行
        if is_resume_run:
            if not message_id:
                return

            # 指定消息唯一标识的消息
            message = conversation.messages.get(message_id)
            if (
                not message
                or not isinstance(approval := message.approval, Approval)
                or not approval.content
            ):
                return

            # 将等待流式输出设置为是
            conversation.awaiting_stream = True
            # 更新消息内容
            message.content = f"您输入的是：\n\n{approval.content}\n\n"
            # 更新审批结果
            approval.result = True
            yield  # 通知前端更新渲染

            # 构建延迟工具结果
            deferred_tool_results = DeferredToolResults(
                approvals={approval.tool_call_id: ToolApproved()},  # 审批通过
                metadata={approval.tool_call_id: {"content": approval.content}},
            )

        else:
            # 若用户提示词为空则直接返回
            if not (user_prompt := conversation.user_prompt):
                return

            # 清空用户提示词
            conversation.user_prompt = ""
            # 将正在运行设置为是
            conversation.is_running = True
            # 将等待流式输出设置为是
            conversation.awaiting_stream = True
            # 先创建消息记录再添加消息
            conversation.messages.update(
                await db_state.create_message_record(
                    conversation_id=conversation.id,
                    message=Message(
                        type=MessageType.USER_PROMPT,
                        content=user_prompt,
                    ),
                )
            )
            yield  # 通知前端更新渲染

        # 初始化片段索引映射为消息唯一标识字典
        index_map_to_message_id: dict[int, str] = {}

        # 查询消息历史列表
        message_history = await db_state.read_message_history(
            conversation_id=self.conversation_id
        )
        # 使用量
        usage = conversation.usage

        # 初始化工作流
        match conversation.workflow.type:
            case WorkflowType.TALK:
                from application.workshop.talk import (
                    run_stream_events,
                )

                # 运行并流式输出事件
                stream_events = run_stream_events(
                    user_prompt=user_prompt,
                    message_history=message_history,
                    usage=usage,
                )
            case WorkflowType.BOOK_FLIGHT:
                from application.workshop.book_flight import (
                    run_stream_events,
                )

                # 运行并流式输出事件
                stream_events = run_stream_events(
                    deps=conversation.workflow.deps,
                    user_prompt=user_prompt,
                    message_history=message_history,
                    usage=usage,
                    deferred_tool_results=deferred_tool_results,
                )

        # 消息列表
        messages: list[Message] = []
        # 获取运行流式输出事件
        async for event in stream_events:
            # 若等待流式输出则将等待流式输出设置为否
            if conversation.awaiting_stream:
                conversation.awaiting_stream = False
            match event:
                # ========== 开始事件 ==========
                case PartStartEvent(
                    index=index,
                    part=part,
                ):
                    match part:
                        # 思考分片
                        case ThinkingPart(content=content):
                            # 构建消息
                            message = Message(
                                type=MessageType.THINKING,
                                title="正在思考",
                                content=content,
                                is_running=True,
                            )
                            # 添加至消息字典
                            conversation.messages[message.id] = message
                            # 将消息唯一标识与片段索引映射
                            index_map_to_message_id[index] = message.id
                            yield  # 通知前端更新渲染

                        # 工具调用分片
                        case ToolCallPart(
                            tool_call_id=tool_call_id, tool_name=tool_name
                        ):
                            if (
                                tool_call_id
                                not in conversation.workflow.tool_call_id_map_to_message_id
                            ):
                                # 构建消息
                                message = Message(
                                    type=MessageType.TOOL_CALL,
                                    title=f"正在{tool_name}",
                                    is_running=True,
                                )
                                # 添加至消息字典
                                conversation.messages[message.id] = message
                                # 将消息唯一标识与工具调用唯一标识映射
                                conversation.workflow.tool_call_id_map_to_message_id[
                                    tool_call_id
                                ] = message.id
                                yield  # 通知前端更新渲染

                        # 文本分片
                        case TextPart(content=content):
                            # 构建消息
                            message = Message(type=MessageType.TEXT, content=content)
                            # 添加至消息字典
                            conversation.messages[message.id] = message
                            # 将消息唯一标识与片段索引映射
                            index_map_to_message_id[index] = message.id
                            yield  # 通知前端更新渲染

                # ========== 增量事件 ==========
                case PartDeltaEvent(index=index, delta=delta):
                    match delta:
                        # 思考分片增量
                        case ThinkingPartDelta(
                            content_delta=content_delta,
                        ):
                            # 若内容增量不为空则增量更新
                            if content_delta:
                                conversation.messages[
                                    index_map_to_message_id[index]
                                ].content += content_delta
                                yield  # 通知前端更新渲染

                        # 文本分片增量
                        case TextPartDelta(
                            content_delta=content_delta,
                        ):
                            # 增量更新
                            conversation.messages[
                                index_map_to_message_id[index]
                            ].content += content_delta
                            yield  # 通知前端更新渲染

                # ========== 结束事件 ==========
                case PartEndEvent(
                    index=index,
                    part=part,
                ):
                    match part:
                        # 思考分片
                        case ThinkingPart():
                            # 获取消息
                            message = conversation.messages[
                                index_map_to_message_id[index]
                            ]
                            message.is_running = False
                            message.title = "思考已完成"
                            yield  # 通知前端更新渲染
                            messages.append(message)

                        # 文本分片
                        case TextPart():
                            # 获取消息
                            message = conversation.messages[
                                index_map_to_message_id[index]
                            ]
                            messages.append(message)

                # ========== 工具调用结果事件 ==========
                case FunctionToolResultEvent(part=part):
                    match part:
                        # 工具返回分片
                        case ToolReturnPart(
                            tool_call_id=tool_call_id,
                            tool_name=tool_name,
                            content=content,
                        ):
                            # 获取消息
                            message = conversation.messages[
                                conversation.workflow.tool_call_id_map_to_message_id[
                                    tool_call_id
                                ]
                            ]
                            message.is_running = False
                            message.title = f"{tool_name}已完成"
                            if isinstance(content, str) and content:
                                message.content = content
                            yield  # 通知前端更新渲染
                            messages.append(message)

                # ========== 智能体运行时事件 ==========
                case AgentRuntimeEvent(tool_call_id=tool_call_id, content=content):
                    # 若工具调用唯一标识为空则创建消息并添加至消息字典，否则若消息内容非空则根据工具调用唯一标识获取消息并更新消息内容
                    if not tool_call_id:
                        message = Message(type=MessageType.TEXT, content=content)
                        conversation.messages[message.id] = message
                        yield  # 通知前端更新渲染
                        messages.append(message)

                    else:
                        if content:
                            message = conversation.messages[
                                conversation.workflow.tool_call_id_map_to_message_id[
                                    tool_call_id
                                ]
                            ]
                            message.content = content
                            yield  # 通知前端更新渲染

                # ========== 智能体运行结果事件 ==========
                case AgentRunResultEvent(result=result):
                    if isinstance(output := result.output, DeferredToolRequests):
                        for appr in output.approvals:
                            # 获取消息
                            message = conversation.messages[
                                conversation.workflow.tool_call_id_map_to_message_id[
                                    tool_call_id := appr.tool_call_id
                                ]
                            ]
                            message.content = output.metadata[tool_call_id]["content"]
                            message.approval = Approval(
                                tool_call_id=tool_call_id,
                            )
                            yield  # 通知前端更新渲染
                            messages.append(message)

                    else:
                        # 若工作流类型非对话则设置为对话
                        if conversation.workflow.type != WorkflowType.TALK:
                            conversation.workflow = Workflow()
                        # 将正在运行设置为否
                        conversation.is_running = False
                        yield  # 通知前端更新渲染

                    # 更新会话记录
                    await db_state.update_conversation_record(
                        conversation.id,
                        usage=result.usage,
                        workflow=conversation.workflow,
                    )

                    # 创建运行记录
                    await db_state.create_run_record(
                        conversation_id=conversation.id,
                        new_messages=result.new_messages(),
                    )

                    # 批量创建消息记录
                    await db_state.create_message_records(
                        conversation_id=conversation.id,
                        messages=messages,
                    )

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

    @rx.var
    def awaiting_stream(self) -> bool:
        """
        当前会话正在等待流式输出
        :return: 当前会话正在等待流式输出（True 表示等待流式输出，False 表示已开始流式输出或已完成）
        """
        # 当前会话
        conversation = self.conversations.get(self.conversation_id)
        if not conversation:
            return False
        return conversation.awaiting_stream

    @rx.event
    def toggle_message_history_item_shown(self, item_id: str) -> None:
        """
        展示 / 隐藏消息历史项
        """
        # 当前会话
        conversation = self.conversations.get(self.conversation_id)
        if not conversation:
            return
        conversation.messages[item_id].is_shown ^= True

    @rx.event
    def set_approval_content(self, item_id: str, content: str) -> None:
        """
        设置审批内容
        :param item_id: 消息历史项唯一标识
        :param content: 内容
        """
        # 当前会话
        conversation = self.conversations.get(self.conversation_id)
        if not conversation:
            return

        # 获取消息
        message = conversation.messages.get(item_id)
        if not message or not isinstance(approval := message.approval, Approval):
            return

        if not content.strip():
            return

        approval.content = content

    @rx.event
    async def set_workflow(self, type: WorkflowType) -> None:
        """
        设置工作流
        """
        # 当前会话
        conversation = self.conversations.get(self.conversation_id)
        if not conversation:
            return

        match type:
            case WorkflowType.BOOK_FLIGHT:
                from application.workshop.book_flight import deps

                # 初始化预定航班工作流
                workflow = Workflow(type=type, deps=deps)
                # 设置用户提示词
                self.set_user_prompt(
                    f"帮我查询并预定在 {workflow.deps.date.strftime('%Y-%m-%d')} 从 {workflow.deps.origin_airport_code} 到 {workflow.deps.destination_airport_code} 的航班"
                )

            case _:
                # 初始化对话工作流
                workflow = Workflow()

        # 设置工作流
        conversation.workflow = workflow
        # 获取数据库状态
        db_state = await self.get_db_state()
        # 更新会话记录
        await db_state.update_conversation_record(
            conversation.id,
            workflow=conversation.workflow,
        )