X25020501_card_book/core/routers/chat/services.py

# -*- coding = utf-8 -*-
# @Time : 2025/3/1 下午4:44
# @File : services.py
# @Software : PyCharm
# @Author : xingc
# @Desc :
from typing import Dict

from sqlalchemy import select, and_, func
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import aliased, selectinload

from core.base.exceptions import WebsocketTooManyRequests
from core.db.models import User, Chat, Message, ChatParticipant
from core.routers.auth.schema import LoginUser
from core.utils.helper import format_date


async def websocket_callback(ws, pexpire):
    raise WebsocketTooManyRequests("Too many requests")


async def get_user_active_direct_chats(
        db_session: AsyncSession, current_user: User
) -> Dict[str, int] | None:
    """获取活跃聊天"""
    direct_chats_dict = dict()
    query = (
        select(Chat)
        .where(and_(Chat.is_deleted.is_(False), Chat.users.contains(current_user)))
        .options(selectinload(Chat.users))
    )
    result = await db_session.execute(query)
    direct_chats: list[Chat] = result.scalars().all()

    if direct_chats:
        for direct_chat in direct_chats:
            direct_chats_dict[str(direct_chat.guid)] = direct_chat.id
        return direct_chats_dict

    return None

async def get_total_unread_messages_count(db_session: AsyncSession, login_user: LoginUser) -> int:
    """获取未读消息总数"""
    query = (
        select(func.count(Message.id))
        .join(
            ChatParticipant,
            and_(
                Message.chat_id == ChatParticipant.chat_id,
                ChatParticipant.user_id == login_user.user_id
            )
        )
        .filter(
            Message.id > ChatParticipant.last_read_message_id,
            Message.user_id != login_user.user_id,
            ChatParticipant.is_deleted == False
        )
    )
    result = await db_session.execute(query)
    return result.scalar()

async def get_chat_unread_messages_count(db_session: AsyncSession, login_user: LoginUser, chats):
    """查询未读消息的预览"""
    # 查询消息情况
    read_status_alias = aliased(ChatParticipant)
    query = (
        select(Message.chat_id, func.count().label("message_count"))
        .join(
            read_status_alias,
            and_(read_status_alias.user_id == login_user.user_id, read_status_alias.chat_id == Message.chat_id),
        )
        .where(
            and_(
                Message.user_id != login_user.user_id,
                Message.id > func.coalesce(read_status_alias.last_read_message_id, 0),
                Message.chat_id.in_(chats),
            )
        )
        .group_by(Message.chat_id)
    )

    result = await db_session.execute(query)
    unread_messages_info = result.fetchall()
    return unread_messages_info


async def get_chat_messages_preview(db_session: AsyncSession, chats):
    """查询未读消息的预览"""
    latest_subquery = (
        select(
            Message.chat_id,
            func.max(Message.created_at).label("max_created")
        )
        .where(Message.chat_id.in_(chats))
        .group_by(Message.chat_id)
        .subquery()
    )

    LatestMessage = aliased(Message)

    query = (
        select(LatestMessage)
        .join(latest_subquery,
              (LatestMessage.chat_id == latest_subquery.c.chat_id) &
              (LatestMessage.created_at == latest_subquery.c.max_created))
    )
    result = await db_session.execute(query)
    unread_messages_info = result.scalars().all()
    return unread_messages_info


async def get_unread_messages_per_chat(
        db_session: AsyncSession, chats: list[Chat], login_user: LoginUser
) -> list[dict]:
    """获取聊天未读消息数和消息预览"""
    messages_info_per_chat: dict[int, dict] = {chat.id: {'message_count': 0} for chat in chats}

    # 未读消息数
    unread_messages_count = await get_chat_unread_messages_count(
        db_session, login_user, messages_info_per_chat
    )
    for messages_info in unread_messages_count:
        messages_info_per_chat[messages_info[0]] = {
            'message_count': messages_info[1],
        }

    # 消息预览
    messages_info = await get_chat_messages_preview(db_session, messages_info_per_chat)
    for message_info in messages_info:
        messages_info_per_chat[message_info.chat_id].update({
            'message': message_info
        })

    results = []
    for chat in chats:
        users = [user for user in await chat.awaitable_attrs.users if user.user_id != login_user.user_id]
        messages_info = messages_info_per_chat[chat.id]
        results.append(dict(
            chat_guid=chat.guid,
            created_at=format_date(chat.created_at),
            updated_at=format_date(chat.updated_at),
            recipient_user=users[0],
            last_message=messages_info.get('message'),
            is_deleted=chat.is_deleted,
            unread_messages_count=messages_info.get('message_count'),
        ))

    return results


async def get_chat(chat_guid, login_user: User, db_session: AsyncSession,
                   is_query_messages: bool = False) -> Chat | None:
    subquery = (
        select(ChatParticipant.chat_id)
        .where(
            and_(
                ChatParticipant.user_id == login_user.user_id,
                ChatParticipant.is_deleted == False
            )
        )
    ).scalar_subquery()

    conditions = [
        Chat.id.in_(subquery),
        Chat.guid == chat_guid
    ]
    if is_query_messages is False:
        conditions.append(
            Chat.is_deleted == False
        )

    query = (
        select(Chat)
        .where(and_(
            *conditions
        ))
    )
    result = await db_session.execute(query)
    await db_session.commit()
    return result.scalar_one_or_none()