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

from pydantic_ai import Agent, ThinkingPartDelta
from pydantic_ai._uuid import uuid7
from pydantic_ai.capabilities import AgentCapability
from pydantic_ai.messages import (
    FunctionToolCallEvent,
    FunctionToolResultEvent,
    LoadCapabilityCallPart,
    ModelMessage,
    PartDeltaEvent,
    PartEndEvent,
    PartStartEvent,
    TextPart,
    TextPartDelta,
    ThinkingPart,
    ToolCallPart,
    ToolSearchCallPart,
)
from pydantic_ai.run import AgentRunResultEvent
from pydantic_ai.models.openai import OpenAIChatModel
from pydantic_ai.providers.openai import OpenAIProvider
import reflex as rx

from application.models import Conversation, EventKind, PartKind, Run, Reasoning, Kind
from application.state.create_conversation_modal import CreateConversationModalState
from application.state.database import DatabaseState


instructions: str = """
# 角色
专业友好AI助手，结构化解答各类问题。

# 输出硬性规则
1. 全文强制标准Markdown，禁止纯文本；不要额外说明排版格式，直接输出内容；
2. 层级使用 `#/##/###`，列表用 `-` 无序列表或数字有序列表；
3. 代码块用 ```语言名``` 包裹；
4. 重点内容标注 **粗体**/*斜体*；
5. 思考、工具日志仅输出文本，适配前端折叠面板，禁止输出HTML标签；
6. 内容分点拆分，排版整洁适配前端Markdown渲染。

# 行文要求
语言通俗，逻辑完整简洁，无多余废话。
"""


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

    # 初始化对话字典
    conversations: Dict[str, Conversation] = {str(uuid7()): Conversation()}
    # 当前对话唯一标识
    conversation_id: str = next(reversed(conversations.keys()))

    # 初始化智能体
    agent: Agent = Agent(
        model=OpenAIChatModel(
            model_name="deepseek-v4-flash",
            provider=OpenAIProvider(
                base_url="https://tokenhub.tencentmaas.com/v1",
                api_key="sk-D9Y1mCe8VlvNqLuSC4mAjqEwxJ2nW4C0h8a7EPn8kg9RLsHq",
            ),
        ),
        instructions=instructions,
        capabilities=None,
        output_type=str,
        retries=1,
    )

    @rx.var
    def get_conversations(self) -> Dict[str, Conversation]:
        """
        获取对话字典，用于前端按照创建倒序渲染对话历史
        :return: 对话字典（倒序）
        """
        return dict(reversed(self.conversations.items()))

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

        # 创建对话
        self.conversations[str(uuid7())] = Conversation(description=description)
        # 将末位对话唯一标识设置为当前对话唯一标识
        self.conversation_id = next(reversed(self.conversations.keys()))

        # 获取创建对话模态窗状态
        create_conversation_modal_state = await self.get_state(
            CreateConversationModalState
        )
        # 关闭创建对话模态窗
        create_conversation_modal_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.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
        """
        if conversation_id not in self.conversations:
            return
        self.conversation_id = conversation_id

    @rx.var
    def get_conversation_description(self) -> str:
        """
        获取当前对话描述，用于前端渲染对话标题
        :return: 当前对话描述
        """
        # 当前对话
        conversation = self.conversations.get(self.conversation_id)
        return conversation.description if conversation else "新对话"

    @rx.var
    def get_runs(self) -> Dict[str, Run]:
        """
        获取运行字典，用于前端渲染运行历史
        :return: 运行字典
        """
        # 当前对话
        conversation = self.conversations.get(self.conversation_id)
        return conversation.runs if conversation else {}

    @rx.var
    def get_run_running_status(self) -> bool:
        """
        获取当前运行运行状态
        :return: 当前运行状态（True 表示正在运行，False 表示非正在运行）
        """
        # 当前对话
        conversation = self.conversations.get(self.conversation_id)
        if not conversation or not conversation.runs:
            return False

        # 末位运行
        run = next(reversed(conversation.runs.values()))
        return run.is_running

    @rx.event
    async def run(self, form_data: dict[str, Any]) -> AsyncGenerator[None]:
        """
        运行
        :param form_data: 表单数据
        :return: AsyncGenerator
        """
        # 当前对话
        conversation = self.conversations.get(self.conversation_id)
        if not conversation:
            return

        # 获取用户提示词
        user_prompt = form_data.get("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.conversation_id
        )

        # 初始化工具调用唯一标识和片段索引映射字典
        tool_call_ids: Dict[str, int] = {}

        # 创建运行
        conversation.runs[str(uuid7())] = Run(user_prompt=user_prompt)
        # 将末位运行唯一标识、运行设置为当前运行唯一标识、运行
        run_id, run = next(reversed(conversation.runs.items()))
        # 将运行状态设置为正在运行
        run.is_running = True
        yield  # 通知前端渲染

        async with self.agent.run_stream_events(
            conversation_id=self.conversation_id,
            user_prompt=user_prompt,
            message_history=message_history,
        ) as events:
            async for event in events:
                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:
                                    run.is_reasoning = True

                                run.reasonings[index] = Reasoning(
                                    kind=Kind.THINKING, content=content
                                )
                                yield

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

                                run.reasonings[index] = Reasoning(
                                    kind=Kind.TOOL_SEARCH, content="正在检索"
                                )
                                yield

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

                                run.reasonings[index] = Reasoning(
                                    kind=Kind.LOAD_CAPABILITY, content="正在加载能力"
                                )
                                yield

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

                                run.reasonings[index] = Reasoning(
                                    kind=Kind.TOOL_CALL, content="正在调用工具"
                                )
                                yield

                            # 文本分片开始事件
                            case TextPart(content=content):
                                run.assistant_content = content
                                yield

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

                            # 文本分片增量事件
                            case TextPartDelta(
                                content_delta=content_delta,
                            ):
                                run.assistant_content += content_delta
                                yield

                    # ========== 结束事件 ==========
                    case PartEndEvent(
                        index=index,
                        part=part,
                        next_part_kind=next_part_kind,
                    ):
                        match part:
                            # 思考分片结束事件
                            case ThinkingPart(part_kind=part_kind, content=content):
                                # 若下一分片种类为文本则将推理状态设置为非正在推理
                                if next_part_kind == Kind.TEXT:
                                    run.is_reasoning = False
                                    yield

                            # 加载能力分片结束事件
                            case LoadCapabilityCallPart(
                                tool_kind=tool_kind, tool_name=tool_name
                            ):
                                yield Event(
                                    event_kind=EventKind(event_kind),
                                    part_index=index,
                                    part_kind=PartKind(tool_kind),
                                    next_part_kind=PartKind(next_part_kind),
                                    tool_name=tool_name,
                                    run_id=current_run_id,
                                )
                            # 调用工具分片结束事件
                            case ToolCallPart(part_kind=part_kind, tool_name=tool_name):
                                yield Event(
                                    event_kind=EventKind(event_kind),
                                    part_index=index,
                                    part_kind=PartKind(part_kind),
                                    next_part_kind=PartKind(next_part_kind),
                                    tool_name=tool_name,
                                    run_id=current_run_id,
                                )

                    # ========== 函数工具调用事件 ==========
                    case FunctionToolCallEvent(
                        part=part,
                    ):
                        match part:
                            # 工具检索分片函数工具调用事件
                            case ToolSearchCallPart(tool_call_id=tool_call_id):
                                index = tool_call_ids[tool_call_id]
                                run.reasonings[index].content = "检索完成"

                    # ========== 函数工具结果事件 ==========
                    case FunctionToolResultEvent(
                        content=content,
                        part=part,
                    ):
                        match part:
                            # 工具检索分片函数工具结果事件
                            case ToolSearchCallPart(tool_call_id=tool_call_id):
                                index = tool_call_ids[tool_call_id]
                                run.reasonings[index].content += content or ""
                                yield

                    # ========== 智能体运行结果事件 ==========
                    case AgentRunResultEvent(result=result):
                        # 保存新增消息
                        await database_state.save_new_messages(
                            conversation_id=self.conversation_id,
                            run_id=run_id,
                            new_messages=result.new_messages(),
                        )
                        # 将运行状态设置为非正在运行
                        run.is_running = False
                        yield