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

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 pydantic_ai._uuid import uuid7
from application.states.database import DatabaseState
from application.states.models import (
    Conversation,
    Message,
    MessageType,
    Work,
    WorkType,
    ConversationHistoryItem,
    usage_validate_python,
)


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] = {}  # 键为会话唯一标识，值为会话实例
    # 激活会话唯一标识
    actived_conversation_id: str = ""

    # 显示会话历史，True 表示显示，False 表示隐藏
    is_conversation_history_shown: bool = False
    # 会话历史项显示悬停气泡
    is_conversation_history_item_popover_shown: 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.get_conversations(self.user_id)
        # 若当前用户的会话字典为空则先创建会话记录再添加会话实例
        if not self.conversations:
            self.conversations.update(
                await db_state.create_conversation_record(self.user_id)
            )
        # 将最后一个会话的唯一标识作为激活会话唯一标识
        self.actived_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) -> 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, conversation_id: str, is_shown: bool
    ) -> None:
        """
        设置会话历史项悬停气泡显示 / 隐藏
        :return: None
        """
        self.is_conversation_history_item_popover_shown = (
            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_conversation_record(conversation_id)
        del self.conversations[conversation_id]

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

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

    @rx.event
    def set_actived_conversation(self, conversation_id: str) -> None:
        """
        设置激活会话
        :param conversation_id: 会话唯一标识
        :return: None
        """
        self.actived_conversation_id = conversation_id

    @rx.var
    def message_history(self) -> list[Message]:
        """
        获取当前会话的消息历史
        :return: 当前会话的消息历史
        """
        # 当前会话
        conversation = self.conversations.get(self.actived_conversation_id)
        if not conversation:
            return []
        return list(conversation.messages.values())

    @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.actived_conversation_id = next(reversed(self.conversations.keys()))

    @rx.var
    def user_prompt(self) -> str:
        """
        用户提示词
        :return: 用户提示词
        """
        # 当前会话
        conversation = self.conversations.get(self.actived_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.actived_conversation_id)
        if not conversation:
            return
        conversation.user_prompt = user_prompt.strip()

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

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

    @rx.event
    async def run(self) -> AsyncGenerator[None]:
        """
        运行
        :return: AsyncGenerator[None]
        """
        # 当前会话
        conversation = self.conversations.get(self.actived_conversation_id)
        if not conversation:
            return

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

        # 将正在运行设置为是
        conversation.is_running = True
        # 构建用户提示词消息实例
        message = Message(
            type=MessageType.USER_PROMPT,
            content=(user_prompt := conversation.user_prompt),
        )
        # 添加至消息字典
        conversation.messages[message.id] = message
        # 清空用户提示词
        conversation.user_prompt = ""
        yield  # 通知前端更新渲染

        # 初始化片段索引映射为消息实例唯一标识字典
        index_map_to_message_id: dict[int, str] = {}
        # 初始化工具调用唯一标识集合
        tool_call_ids: set[str] = set()
        try:
            # 获取数据库状态
            db_state = await self.get_db_state()
            # 获取消息历史列表
            message_history = await db_state.get_message_history(
                conversation_id=self.actived_conversation_id
            )
            usage = usage_validate_python(conversation.usage)
            # 匹配工作类型
            match conversation.work:
                # 预定航班
                case WorkType.BOOK_FLIGHT:
                    from application.workshop.book_flight import run_stream_events

                    stream_events = run_stream_events(
                        work=conversation.work,
                        user_prompt=user_prompt,
                        message_history=message_history,
                        usage=usage,
                    )

                # 非结构化对话
                case _:
                    from application.workshop.unstructured_dialogue import (
                        run_stream_events,
                    )
                    # 运行并流式输出事件
                    stream_events = run_stream_events(
                        user_prompt=user_prompt,
                        message_history=message_history,
                        usage=usage,
                    )

            # 获取运行流式输出事件
            async for event in stream_events:
                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

                            # 工具调用分片开始事件
                            case ToolCallPart(
                                tool_name=tool_name, tool_call_id=tool_call_id
                            ):
                                # 构建消息实例
                                message = Message(
                                    type=MessageType.TOOL_CALL,
                                    title=tool_name,
                                    content=",",
                                    is_running=True,
                                )
                                tool_call_ids.add(tool_call_id)
                                # 添加至消息字典
                                dialog.thoughts[index] = Thought(
                                    type="tool_call",
                                    content="正在生成调用参数",
                                )

                            # 文本分片开始事件
                            case TextPart(content=content):
                                # 构建消息实例
                                message = Message(type=MessageType.RESULT_OUTPUT)
                                message.content = content
                                # 添加至消息字典
                                conversation.messages[message.id] = message
                                # 将消息实例唯一标识与片段索引映射
                                index_map_to_message_id[index] = message.id

                    # ========== 增量事件 ==========
                    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

                            # 文本分片增量事件
                            case TextPartDelta(
                                content_delta=content_delta,
                            ):
                                # 增量更新
                                conversation.messages[
                                    index_map_to_message_id[index]
                                ].content += content_delta

                    # ========== 结束事件 ==========
                    case PartEndEvent(
                        index=index,
                        part=part,
                    ):
                        match part:
                            # 思考分片结束事件
                            case ThinkingPart(content=content):
                                # 获取消息实例
                                message = conversation.messages[
                                    index_map_to_message_id[index]
                                ]
                                message.is_running = False
                                message.title = "思考完成"

                    # ========== 函数工具调用事件 ==========
                    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 AgentRunResultEvent(result=result):
                        # 创建对话记录
                        await db_state.create_dialog_record(
                            conversation_id=self.conversation_id,
                            id=dialog.id,
                            user_prompt=dialog.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,
                            _id=dialog.id,
                            new_messages=result.new_messages(),
                        )

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

        except Exception as e:
            ...
        finally:
            # 将正在运行设置为否
            conversation.is_running = False

            self.conversations[self.actived_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_message_collapse(self, message_id: str) -> None:
        """
        展开/折叠消息组件
        """
        # 当前会话
        conversation = self.conversations.get(self.actived_conversation_id)
        if not conversation:
            return
        conversation.messages[message_id].is_expanded ^= True