# -*- coding: utf-8 -*-
"""
会话页面状态
"""
from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple, cast

from pydantic_ai._uuid import uuid7
import reflex as rx

from application.models import Conversation, EventKind, PartKind, Run
from application.state.create_conversation import CreateConversationState
from application.state.database import DatabaseState
from application.utils.agent import AIAgent


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

    # 当前会话唯一标识
    current_conversation_id: str = str(uuid7())
    # 会话字典
    conversations: Dict[str, Conversation] = {current_conversation_id: Conversation()}

    # 当前运行唯一标识
    current_run_id: Optional[str] = None

    # 初始化智能体
    agent: AIAgent = AIAgent()

    @rx.var
    def get_conversations(self) -> Dict[str, Conversation]:
        """
        获取会话字典，用于前端渲染会话列表
        :return: 会话字典
        """
        return self.conversations

    @rx.var
    def get_current_conversation_description(self) -> str:
        """
        获取当前会话描述，用于前端渲染导航栏中当前会话描述
        :return: 当前会话描述
        """
        # 若当前会话唯一标识不存在则创建会话
        if self.current_conversation_id not in self.conversations:
            self.conversations[self.current_conversation_id] = Conversation()

        # 当前会话
        current_conversation = self.conversations[self.current_conversation_id]
        return current_conversation.description

    @rx.event
    def set_current_conversation_id(self, conversation_id: str) -> None:
        """
        将指定会话唯一标识设置为当前会话唯一标识
        :param conversation_id: 指定会话唯一标识
        :return: None
        """
        # 若指定会话唯一标识不存在则创建会话
        if conversation_id not in self.conversations:
            self.conversations[conversation_id] = Conversation()

        # 将指定会话唯一标识设置为当前会话唯一标识
        self.current_conversation_id = conversation_id

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

        # 创建会话唯一标识
        self.current_conversation_id = str(uuid7())
        self.conversations[self.current_conversation_id] = Conversation(
            description=description
        )

        # 跨状态读取创建会话模态窗状态
        create_conversation_state = await self.get_state(CreateConversationState)
        create_conversation_state.is_open = False

    @rx.event
    def delete_conversation(self, conversation_id: str) -> None:
        """
        删除会话
        :param conversation_id: 会话唯一标识
        :return: None
        """
        # 若指定会话唯一标识不存在则直接返回
        if conversation_id not in self.conversations:
            return

        del self.conversations[conversation_id]
        # 删除后若会话字典为空则后创建会话
        if not self.conversations:
            self.conversations[str(uuid7())] = Conversation()
        # 删除后若当前会话唯一标识不存在则将第一个会话唯一标识设置为当前会话唯一标识
        if self.current_conversation_id not in self.conversations:
            self.current_conversation_id = next(iter(self.conversations))

    @rx.var
    def get_runs(self) -> Dict[str, Run]:
        """
        获取当前会话运行字典，用于前端渲染运行列表
        :return: 当前会话运行字典
        """
        # 若当前会话唯一标识不存在则创建会话
        if self.current_conversation_id not in self.conversations:
            self.conversations[self.current_conversation_id] = Conversation()

        # 当前会话
        current_conversation = self.conversations[self.current_conversation_id]
        return current_conversation.runs

    @rx.var
    def get_current_run_streaming_status(self) -> bool:
        """
        获取当前运行流式输出状态
        :return: 当前运行流式输出状态（True 表示正在流式输出，False 表示非正在流式输出）
        """
        # 若当前会话唯一标识不存在则创建会话
        if self.current_conversation_id not in self.conversations:
            self.conversations[self.current_conversation_id] = Conversation()

        # 当前会话
        current_conversation = self.conversations[self.current_conversation_id]

        # 若当前运行唯一标识为空或不存在则返回 False
        if (
            not self.current_run_id
            or self.current_run_id not in current_conversation.runs
        ):
            return False

        # 当前运行
        current_run = current_conversation.runs[self.current_run_id]
        return current_run.is_streaming

    @rx.event
    async def run(self, form_data: dict[str, Any]) -> AsyncGenerator[None]:
        """
        运行
        :param form_data: 表单数据
        :return: AsyncGenerator
        """
        # 若当前会话唯一标识不存在则创建会话
        if self.current_conversation_id not in self.conversations:
            self.conversations[self.current_conversation_id] = Conversation()

        # 当前会话
        current_conversation = self.conversations[self.current_conversation_id]

        # 获取用户提示词
        user_prompt = form_data["user_prompt"].strip()
        if not user_prompt:
            return

        # 数据库状态
        database_state = await self.get_state(DatabaseState)
        # 获取当前会话消息历史
        message_history = await database_state.get_message_history(
            conversation_id=self.current_conversation_id
        )

        async for event in self.agent.run(
            conversation_id=self.current_conversation_id,
            user_prompt=user_prompt,
            message_history=message_history,
        ):

            # 若事件为空或当前运行唯一标识为空或当前运行唯一标识与事件运行唯一标识不相同则跳过
            if (
                not event
                or not self.current_run_id
                or self.current_run_id != event.run_id
            ):
                continue

            # ========== 运行开始 ==========
            if event.event_kind == EventKind.RUN_START:
                # 将事件运行唯一标识设置为当前运行唯一标识
                self.current_run_id = event.run_id
                # 创建运行
                current_conversation.runs[self.current_run_id] = Run(
                    user_prompt=user_prompt, is_streaming=True
                )
                yield  # 通知前端渲染
                continue

            # 当前运行
            current_run = current_conversation.runs[self.current_run_id]

            # ========== 推理 ==========
            if (
                event.event_kind == EventKind.PART_START
                and event.part_kind == PartKind.THINKING
            ):
                current_run.thinking += event.event_content
                yield  # 通知前端渲染
                continue

            if part_index := event.part_index:
                if part_index not in current_run.reasonings:
                    current_run.reasonings[part_index] += event.event_content
                    yield  # 通知前端渲染
                    continue

            # ========== 回答 ==========
            if (
                event.part_kind == PartKind.TEXT
                and event.event_content != EventKind.PART_END
            ):
                current_run.answer += event.event_content
                yield  # 通知前端渲染
                continue

            # ========== 运行结束 ==========
            if event.event_kind == EventKind.RUN_END:
                # 保存运行新增消息
                await database_state.save_new_message(
                    conversation_id=self.current_conversation_id,
                    run_id=self.current_run_id,
                    run_new_message=event.run_new_messages,
                )
                # 设置当前运行流式输出状态非正在流式输出
                current_run.is_streaming = False
                yield  # 通知前端渲染
                continue