# -*- coding: utf-8 -*-
"""
全局共享状态
"""

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

from application.models import (
    MessageHistory,
    ModelMessage,
    ModelMessagesTypeAdapter,
)


class GlobalState(rx.State):
    """全局状态"""

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

    async def save_new_message(
        self, chat_id: str, new_message: List[ModelMessage]
    ) -> None:
        """
        保存一轮对话消息
        :param chat_id: 聊天唯一标识
        :param new_message: 一轮对话消息
        :return: None
        """
        async with rx.asession() as session:
            record = MessageHistory.adapt(chat_id=chat_id, new_message=new_message)
            session.add(record)
            await session.commit()