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

from typing import List

import reflex as rx
import sqlalchemy as sa
from sqlmodel import select as sql_select

from application.models import MessageHistory, ModelMessage, ModelMessagesTypeAdapter


class DatabaseState(rx.State):
    """
    数据库状态
    重新初始化：先手动删除 alembic 相关配置和文件夹，再使用 reflex db init 初始化数据库表
    """

    async def get_message_history(self, conversation_id: str) -> List[ModelMessage]:
        """
        获取消息历史
        :param conversation_id: 对话唯一标识
        :return: 消息历史
        """
        message_history = []
        async with rx.asession() as session:
            records = await session.exec(
                sql_select(MessageHistory)
                .where(MessageHistory.conversation_id == conversation_id)
                .order_by(sa.desc("run_id"))
            )
            for record in records.all():
                message_history.extend(
                    ModelMessagesTypeAdapter.validate_json(record.new_message)
                )
        return message_history

    async def save_new_messages(
        self, conversation_id: str, run_id: str, new_messages: List[ModelMessage]
    ) -> None:
        """
        保存新增消息
        :param conversation_id: 对话唯一标识
        :param run_id: 运行唯一标识
        :param new_messages: 新增消息
        :return: None
        """
        async with rx.asession() as session:
            record = MessageHistory.adapt(
                conversation_id=conversation_id,
                run_id=run_id,
                new_message=new_messages,
            )
            session.add(record)
            await session.commit()