# -*- coding: utf-8 -*-
"""
聊天页面状态
"""
from typing import Any, AsyncGenerator, Dict, List
from uuid import uuid4

import reflex as rx
from application.utils.agent import Agent
from application.models import Chat, Dialog, Message, PREFIX_MAPING
from application.state.create_chat_modal import CreateChatModalState
from application.state.global_ import GlobalState


# 绑定的智能体列表（键为聊天唯一标识，值为智能体实例）
# 因 rx.state 需在 Python 进程与 WebSocket 之间传输和存储等操作，非序列化对象会失败，故需将智能体实例单独存储
agents: Dict[str, Agent] = {}


def get_agent(chat_id: str) -> Agent:
    """
    获取绑定的智能体
    :return: 绑定的智能体
    """
    if chat_id not in agents:
        agents[chat_id] = Agent(
            chat_id=chat_id,
            instructions="You are a friendly chatbot",  # 智能体指令
        )
    return agents[chat_id]


class ChatState(rx.State):
    """
    聊天页面状态
    """

    # 当前聊天唯一标识
    current_chat_id: str = uuid4().hex
    # 聊天列表
    chats: Dict[str, Chat] = {current_chat_id: Chat()}

    @rx.var
    def get_chat_ids(self) -> List[str]:
        """
        获取聊天唯一标识列表
        :return: 聊天唯一标识列表
        """
        return list(self.chats)

    @rx.var
    def get_chat_descriptions(self) -> Dict[str, str]:
        """
        获取聊天描述列表
        :return: 聊天描述列表
        """
        return {chat_id: chat.description for chat_id, chat in self.chats.items()}

    @rx.var
    def get_current_chat_status(self) -> bool:
        """
        获取当前聊天流式输出状态
        :return: 当前聊天流式输出状态（True 表示正在流式输出，False 表示非正在流式输出）
        """
        chat = self.chats.get(self.current_chat_id)
        return chat.is_streaming if chat else False

    @rx.var
    def get_current_chat_dialogs(self) -> List[Dialog]:
        """
        获取当前聊天对话列表
        :return: 当前聊天对话列表
        """
        chat = self.chats.get(self.current_chat_id)
        return chat.dialogs if chat else []

    @rx.event
    def switch_chat(self, chat_id: str) -> None:
        """
        切换聊天
        :param chat_id: 聊天唯一标识
        :return: None
        """
        if chat_id not in self.chats:
            return
        self.current_chat_id = chat_id

    @rx.event
    async def create_chat(self, form_data: Dict[str, Any]) -> None:
        """
        新建聊天
        :param form_data: 表单数据
        :return: None
        """
        # 生成聊天唯一标识
        self.current_chat_id = uuid4().hex
        self.chats[self.current_chat_id] = Chat(
            description=form_data["chat_description"].strip() or "新聊天"
        )

        # 跨状态读取新建聊天模态窗状态（异步事件）
        create_chat_modal_state = await self.get_state(CreateChatModalState)
        create_chat_modal_state.is_open = False

    @rx.event
    def delete_chat(self, chat_id: str) -> None:
        """
        删除聊天
        :param chat_id: 聊天唯一标识
        :return: None
        """
        if chat_id not in self.chats:
            return
        del self.chats[chat_id]
        # 删除聊天后若聊天列表为空，则新建聊天
        if not self.chats:
            self.chats[uuid4().hex] = Chat()
        # 删除聊天后若当前聊天唯一标识不存在，则切换到第一个聊天
        if self.current_chat_id not in self.chats:
            self.current_chat_id = next(iter(self.chats))

    @rx.event
    async def process_input(self, form_data: dict[str, Any]) -> AsyncGenerator:
        """
        处理输入，返回流式输出
        :param form_data: 表单数据
        :return: AsyncGenerator
        """
        input_ = form_data["input"].strip()
        if not input_:
            return

        # 当前聊天
        current_chat = self.chats[self.current_chat_id]
        if not current_chat:
            return

        # 新增对话
        current_chat.dialogs.append(Dialog(input_=input_))
        # 当前对话
        current_dialog = current_chat.dialogs[-1]

        # 设置当前对话流式输出状态为正在流式输出
        current_chat.is_streaming = True
        yield  # 通知前端渲染输入消息

        # 全局状态
        global_state = await self.get_state(GlobalState)
        # 获取当前聊天绑定的智能体
        agent = get_agent(chat_id=self.current_chat_id)
        # 获取当前聊天消息历史
        message_history = await global_state.get_message_history(
            chat_id=self.current_chat_id
        )

        # 流式输出消息事件
        async for event in agent.stream_messages_events(
            user_prompt=input_, message_history=message_history
        ):
            # 若消息事件为空则跳过
            if not event:
                continue

            # 匹配前缀
            prefix = next((p for p in PREFIX_MAPING if event.startswith(p)), None)
            # 若未匹配到前缀则跳过
            if not prefix:
                continue

            # 根据前缀匹配消息类型
            type_ = PREFIX_MAPING[prefix]
            # 若当前对话输出消息为空或当前消息类型与新消息类型不相同则新增消息
            if not current_dialog.output or current_dialog.output[-1].type_ != type_:
                current_dialog.output.append(Message(type_=type_))

            # 追加消息内容
            current_dialog.output[-1].content += event.removeprefix(prefix)
            yield  # 通知前端渲染输出消息

        # 保存本轮对话消息
        await global_state.save_new_message(
            chat_id=self.current_chat_id, new_message=agent.new_messages
        )
        # 设置当前对话流式输出状态非正在流式输出
        current_chat.is_streaming = False

    @rx.event
    async def toggle_collapse_panel(self, dialog_id: str, message_id: str) -> None:
        """
        打开/关闭折叠面板
        :return: None
        """
        # 当前聊天
        current_chat = self.chats[self.current_chat_id]
        if not current_chat:
            return
        # 目标对话
        target_dialog = next(
            (d for d in current_chat.dialogs if d.id == dialog_id), None
        )
        if not target_dialog:
            return
        # 目标消息
        target_message = next(
            (m for m in target_dialog.output if m.id == message_id), None
        )
        if not target_message:
            return
        # 切换折叠面板打开状态
        target_message.is_open = not target_message.is_open