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

from typing import List, Dict

import reflex as rx
from sqlalchemy import desc, and_
from sqlmodel import select

from application.models import Conversations, RunResults
from pydantic_ai import ModelMessage, ModelMessagesTypeAdapter


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

    async def get_conversations(self, user_id: str) -> Dict[str, Conversation]:
        """
        获取指定用户的会话字典
        :param user_id: 用户唯一标识
        :return: 指定用户的会话字典
        """
        conversations: Dict[str, Conversation] = {}
        async with rx.asession() as session:
            result = await session.exec(
                select(Conversations, Dialogs, Parts)
                .outerjoin(Dialogs, Conversations.id == Dialogs.conversation_id)  # type: ignore
                .outerjoin(Parts, Dialogs.id == Parts.dialog_id)  # type: ignore
                .where(
                    models.Conversation.user_id == user_id,
                    models.Conversation.is_deleted == False,
                )
                .order_by(
                    models.Conversation.id, models.Dialog.id, models.ReasoningPart.id
                )
            )
            records = result.all()
            for db_conversation, _, _ in records:
                conversations[conversation.id] = models.Conversation(
                    user_id=user_id, description=conversation.description
                )

        return conversations

    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(RunResults)
                .where(RunResults.conversation_id == conversation_id)
                .order_by(desc(RunResults.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:
            record = RunResults.adapt(
                conversation_id=conversation_id,
                dialog_id=dialog_id,
                new_messages=new_messages,
            )
            session.add(record)
            await session.commit()