# -*- coding: utf-8 -*-
"""
数据库状态
"""

from datetime import datetime
from typing import Dict, List, cast, TypeAlias

from pydantic import BaseModel, Field, TypeAdapter
from pydantic_ai import ModelMessage, ModelMessagesTypeAdapter
import reflex as rx
from sqlalchemy import ColumnElement, desc
from sqlmodel import select

from application.models import Conversation, Dialog, RunResult


class ConversationMemory(BaseModel):
    """
    会话内存模型
    """

    description: str = Field(default="新会话", description="会话描述")
    is_running: bool = Field(
        default=False, description="运行状态，True 表示运行中，False 表示运行完成"
    )
    dialogs: Dict[str, DialogMemory] = Field(
        default_factory=dict, description="对话字典"
    )
    created_at: datetime = Field(default_factory=datetime.now, description="创建时间")


class DialogMemory(BaseModel):
    """
    对话内存模型
    """

    question: str = Field(..., description="问题")
    is_thinking: bool = Field(
        default=False, description="思考状态，True 表示思考中，False 表示思考完成"
    )
    is_collapse_expanded: bool = Field(
        default=False, description="思考折叠面板展开状态，True 表示展开，False 表示折叠"
    )
    thought_nodes: list[ThoughtNodeMemory] = Field(
        default_factory=list, description="思考节点列表"
    )
    answer: str = Field(default="", description="回答")


class ThoughtNodeMemory(BaseModel):
    """
    思考节点内存模型
    """

    kind: str
    content: str


# 会话字典类型别名
Conversations: TypeAlias = Dict[str, ConversationMemory]

# 思考节点内存模型类型适配器
ThoughtNodeMemoryTypeAdapter = TypeAdapter(list[ThoughtNodeMemory])


class DatabaseState(rx.State):
    """
    数据库状态
    """

    async def get_conversations(self, user_id: str) -> Conversations:
        """
        获取会话字典
        :param user_id: 用户唯一标识
        :return: 会话字典
        """
        conversations: Conversations = {}
        async with rx.asession() as session:
            result = await session.exec(
                select(Conversation, Dialog)
                .outerjoin(
                    Dialog,
                    cast(
                        ColumnElement[bool],
                        Conversation.id == Dialog.conversation_id,
                    ),
                )
                .where(
                    Conversation.user_id == user_id,
                    Conversation.is_deleted == False,
                )
                .order_by(Conversation.id, Dialog.id)
            )
            for conversation, dialog in result.all():
                conversation_memory = conversations.setdefault(
                    conversation.id,
                    ConversationMemory(
                        description=conversation.description,
                        created_at=conversation.created_at,
                    ),
                )

                if not dialog:
                    continue
                conversation_memory.dialogs[dialog.id] = DialogMemory(
                    question=dialog.question,
                    thought_nodes=ThoughtNodeMemoryTypeAdapter.validate_python(
                        dialog.thought_nodes
                    ),  # 反序列化为思考节点内存实例列表
                    answer=dialog.answer,
                )
        return conversations

    async def create_conversation(
        self, user_id: str, description: str = "新会话"
    ) -> tuple[str, datetime]:
        """
        新建会话
        :param user_id: 用户唯一标识
        :param description: 会话描述
        :return: 新建会话唯一标识、创建时间
        """
        async with rx.asession() as session:
            conversation = Conversation(user_id=user_id, description=description)
            session.add(conversation)
            await session.commit()
            await session.refresh(conversation)
        return conversation.id, conversation.created_at

    async def delete_conversation(self, conversation_id: str) -> None:
        """
        删除指定会话（逻辑删除）
        :param conversation_id: 指定会话唯一标识
        :return: None
        """
        async with rx.asession() as session:
            conversation = await session.get(Conversation, conversation_id)
            if not conversation:
                return
            conversation.is_deleted = True
            await session.commit()

    async def create_dialog(self, conversation_id: str, question: str) -> str:
        """
        新建对话
        :param conversation_id: 会话唯一标识
        :param question: 提问
        :return: 新建对话唯一标识
        """
        async with rx.asession() as session:
            dialog = Dialog(
                conversation_id=conversation_id,
                question=question,
            )
            session.add(dialog)
            await session.commit()
            await session.refresh(dialog)
        return dialog.id

    async def complete_dialog(
        self,
        dialog_id: str,
        thought_nodes: list[ThoughtNodeMemory],
        answer: str,
    ) -> None:
        """
        完成对话
        :param dialog_id: 对话唯一标识
        :param thought_nodes: 思考节点列表
        :param answer: 回答
        :return: None
        """
        async with rx.asession() as session:
            dialog = await session.get(Dialog, dialog_id)
            if not dialog:
                return
            dialog.thought_nodes = ThoughtNodeMemoryTypeAdapter.dump_python(
                thought_nodes
            )
            dialog.answer = answer
            await session.commit()

    async def get_message_history(self, conversation_id: str) -> List[ModelMessage]:
        """
        获取消息历史
        :param conversation_id: 会话唯一标识
        :return: 消息历史
        """
        message_history: List[ModelMessage] = []
        async with rx.asession() as session:
            result = await session.exec(
                select(RunResult)
                .where(RunResult.conversation_id == conversation_id)
                .order_by(desc(RunResult.dialog_id))
            )
            for record in result.all():
                message_history.extend(
                    ModelMessagesTypeAdapter.validate_json(record.new_messages)
                )
        return message_history

    async def save_new_messages(
        self,
        conversation_id: str,
        dialog_id: str,
        new_messages: List[ModelMessage],
    ) -> None:
        """
        保存新增消息
        :param conversation_id: 会话唯一标识
        :param dialog_id: 对话唯一标识
        :param new_messages: 新增消息
        :return: None
        """
        async with rx.asession() as session:
            session.add(
                RunResult(
                    conversation_id=conversation_id,
                    dialog_id=dialog_id,
                    new_messages=ModelMessagesTypeAdapter.dump_json(
                        new_messages
                    ).decode(
                        "utf-8"
                    ),  # 序列化为 JSON 字符串
                )
            )
            await session.commit()