hunjing/backend/app/repositories/chat.py

"""聊天数据访问。"""

from sqlalchemy import func, select, update
from sqlalchemy.ext.asyncio import AsyncSession

from app.core.errors import AppError
from app.models.chat import ChatMessage, ChatSession


async def get_session(
    session: AsyncSession, session_id: str, *, owner_scope: str | None = None
) -> ChatSession | None:
    stmt = select(ChatSession).where(ChatSession.id == session_id)
    if owner_scope is not None:
        stmt = stmt.where(ChatSession.owner_scope == owner_scope)
    result = await session.execute(stmt)
    return result.scalar_one_or_none()


async def list_recent_messages(
    session: AsyncSession, session_id: str, *, limit: int = 12
) -> list[ChatMessage]:
    result = await session.execute(
        select(ChatMessage)
        .where(ChatMessage.session_id == session_id, ChatMessage.deleted_at.is_(None))
        .order_by(ChatMessage.created_at.desc())
        .limit(limit)
    )
    return list(reversed(result.scalars().all()))


async def list_sessions(
    session: AsyncSession,
    character_id: str,
    *,
    status: str | None = None,
    owner_scope: str | None = None,
    limit: int = 50,
    offset: int = 0,
) -> list[ChatSession]:
    conditions = [ChatSession.character_id == character_id, ChatSession.status != "deleted"]
    if status is not None:
        conditions.append(ChatSession.status == status)
    if owner_scope is not None:
        conditions.append(ChatSession.owner_scope == owner_scope)
    result = await session.execute(
        select(ChatSession)
        .where(*conditions)
        .order_by(ChatSession.updated_at.desc())
        .offset(offset)
        .limit(limit)
    )
    return list(result.scalars().all())


async def count_sessions(
    session: AsyncSession,
    character_id: str,
    *,
    status: str | None = None,
    owner_scope: str | None = None,
) -> int:
    conditions = [ChatSession.character_id == character_id, ChatSession.status != "deleted"]
    if status is not None:
        conditions.append(ChatSession.status == status)
    if owner_scope is not None:
        conditions.append(ChatSession.owner_scope == owner_scope)
    result = await session.execute(select(func.count()).select_from(ChatSession).where(*conditions))
    return result.scalar_one()


async def list_messages(
    session: AsyncSession, session_id: str, *, limit: int = 100, offset: int = 0
) -> list[ChatMessage]:
    result = await session.execute(
        select(ChatMessage)
        .where(ChatMessage.session_id == session_id, ChatMessage.deleted_at.is_(None))
        .order_by(ChatMessage.created_at.asc())
        .offset(offset)
        .limit(limit)
    )
    return list(result.scalars().all())


async def count_messages(session: AsyncSession, session_id: str) -> int:
    result = await session.execute(
        select(func.count())
        .select_from(ChatMessage)
        .where(
            ChatMessage.session_id == session_id,
            ChatMessage.deleted_at.is_(None),
        )
    )
    return result.scalar_one()


async def count_messages_by_sessions(
    session: AsyncSession, session_ids: list[str]
) -> dict[str, int]:
    """批量统计多个会话的未删除消息数（GROUP BY 聚合，避免列表查询 N+1）。"""
    if not session_ids:
        return {}
    result = await session.execute(
        select(ChatMessage.session_id, func.count())
        .where(
            ChatMessage.session_id.in_(session_ids),
            ChatMessage.deleted_at.is_(None),
        )
        .group_by(ChatMessage.session_id)
    )
    return {str(sid): int(cnt) for sid, cnt in result.all()}


async def last_message_previews_by_sessions(
    session: AsyncSession, session_ids: list[str]
) -> dict[str, str]:
    """批量取每个会话最后一条未删除消息的内容（截断由调用方处理）。

    用 row_number() 窗口函数取每个 session_id 的最新一条，兼容 SQLite（3.25+）
    与 PostgreSQL，避免逐会话查询（N+1）。无消息的会话不在结果中。
    """
    if not session_ids:
        return {}
    latest = (
        select(
            ChatMessage.session_id,
            ChatMessage.content,
            func.row_number()
            .over(
                partition_by=ChatMessage.session_id,
                order_by=ChatMessage.created_at.desc(),
            )
            .label("rn"),
        ).where(
            ChatMessage.session_id.in_(session_ids),
            ChatMessage.deleted_at.is_(None),
        )
    ).subquery()
    result = await session.execute(
        select(latest.c.session_id, latest.c.content).where(latest.c.rn == 1)
    )
    return {str(sid): (content or "") for sid, content in result.all()}


async def count_messages_by_role(session: AsyncSession, session_id: str, *, role: str) -> int:
    """统计会话内指定角色的未删除消息数（用于判断首轮对话）。"""
    result = await session.execute(
        select(func.count())
        .select_from(ChatMessage)
        .where(
            ChatMessage.session_id == session_id,
            ChatMessage.role == role,
            ChatMessage.deleted_at.is_(None),
        )
    )
    return result.scalar_one()


async def update_session_title(session: AsyncSession, session_id: str, title: str) -> None:
    """更新会话标题；仅修改托管对象，提交时机由调用方控制。"""
    chat_session = await session.get(ChatSession, session_id)
    if chat_session is not None:
        chat_session.title = title


async def soft_delete_session_messages(session: AsyncSession, session_id: str) -> None:
    """软删除会话内全部未删除消息；提交时机由调用方控制。"""
    from app.core.time import utc_now

    await session.execute(
        update(ChatMessage)
        .where(
            ChatMessage.session_id == session_id,
            ChatMessage.deleted_at.is_(None),
        )
        .values(deleted_at=utc_now())
    )


async def get_message(session: AsyncSession, message_id: str) -> ChatMessage | None:
    """按 ID 获取单条消息。"""
    return await session.get(ChatMessage, message_id)


async def set_message_active_version(
    session: AsyncSession,
    message_id: str,
    version_idx: int,
    versions_json: str,
) -> None:
    """更新消息的活跃版本索引和 versions_json。"""
    from sqlalchemy import update

    result = await session.execute(
        update(ChatMessage)
        .where(ChatMessage.id == message_id)
        .values(active_version_idx=version_idx, versions_json=versions_json)
    )
    if result.rowcount == 0:
        raise AppError("not_found", "消息不存在。")