life-echo/api/app/features/quota/service.py

"""
配额检查业务逻辑。

「对话轮数」的定义：每条用户发出的消息（Segment 表的记录数）计为 1 轮。
"""

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

from app.features.conversation.models import Conversation, Segment
from app.features.memoir.models import Chapter
from app.features.quota.schemas import QuotaCheckResponse

PLAN_QUOTAS = {
    "free": {
        "max_conversations": 50,
        "max_chapters": 1,
        "max_words": None,
    },
    "pro": {
        "max_conversations": 2000,
        "max_chapters": None,
        "max_words": None,
    },
    "pro_plus": {
        "max_conversations": 10000,
        "max_chapters": None,
        "max_words": None,
    },
    "premium": {
        "max_conversations": None,
        "max_chapters": None,
        "max_words": None,
    },
    "test": {
        "max_conversations": None,
        "max_chapters": None,
        "max_words": None,
    },
}


async def get_segment_count(user_id: str, db: AsyncSession) -> int:
    """获取用户已消耗的对话轮数（= 该用户所有 Segment 记录数）。"""
    stmt = (
        select(func.count(Segment.id))
        .join(Conversation, Segment.conversation_id == Conversation.id)
        .where(
            Conversation.user_id == user_id,
            Conversation.deleted_at.is_(None),
        )
    )
    result = await db.execute(stmt)
    return result.scalar() or 0


async def get_chapter_count(user_id: str, db: AsyncSession) -> int:
    """获取用户当前章节数量。"""
    stmt = select(func.count(Chapter.id)).where(Chapter.user_id == user_id)
    result = await db.execute(stmt)
    return result.scalar() or 0


async def get_conversation_count(user_id: str, db: AsyncSession) -> int:
    """别名：实际按 Segment 计数。"""
    return await get_segment_count(user_id, db)


def check_can_send_message(
    subscription_type: str,
    segment_count: int,
) -> tuple[bool, str]:
    """检查用户是否还能发送消息（对话轮数）。返回 (是否允许, 提示信息)。"""
    quotas = PLAN_QUOTAS.get(subscription_type, PLAN_QUOTAS["free"])
    max_conv = quotas.get("max_conversations")
    if max_conv is None:
        return True, ""
    if segment_count >= max_conv:
        return (
            False,
            f"对话轮数已用完（{segment_count}/{max_conv}），请升级 Pro 或 Pro+ 继续使用",
        )
    return True, ""


def check_can_submit_organize(
    subscription_type: str,
    chapter_count: int,
) -> tuple[bool, str]:
    """检查是否可以提交整理任务（生成新章节）。免费版仅允许 1 个章节。"""
    quotas = PLAN_QUOTAS.get(subscription_type, PLAN_QUOTAS["free"])
    max_ch = quotas.get("max_chapters")
    if max_ch is None:
        return True, ""
    if chapter_count >= max_ch:
        return False, "章节数量已达上限（免费版仅支持 1 个章节整理），请升级后继续"
    return True, ""


class QuotaService:
    def __init__(self, db: AsyncSession):
        self._db = db

    async def get_usage(self, user_id: str) -> tuple[int, int]:
        """Return (segment_count, chapter_count)."""
        seg = await get_segment_count(user_id, self._db)
        ch = await get_chapter_count(user_id, self._db)
        return seg, ch

    async def check_can_send_message(
        self, user_id: str, subscription_type: str
    ) -> tuple[bool, str]:
        """检查用户是否还能发送消息（对话轮数）。返回 (是否允许, 提示信息)。"""
        count = await get_segment_count(user_id, self._db)
        return check_can_send_message(subscription_type, count)

    async def check_can_submit_organize(
        self, user_id: str, subscription_type: str
    ) -> tuple[bool, str]:
        """检查是否可以提交整理任务（生成新章节）。返回 (是否允许, 提示信息)。"""
        chapter_count = await get_chapter_count(user_id, self._db)
        return check_can_submit_organize(subscription_type, chapter_count)

    async def check(self, user_id: str, subscription_type: str) -> QuotaCheckResponse:
        """检查用户配额使用情况。"""
        quotas = PLAN_QUOTAS.get(subscription_type, PLAN_QUOTAS["free"])
        segment_count, chapter_count = await self.get_usage(user_id)
        max_conversations = quotas.get("max_conversations")
        max_chapters = quotas.get("max_chapters")
        max_words = quotas.get("max_words")

        remaining_conversations = None
        remaining_chapters = None
        remaining_words = None
        if max_conversations is not None:
            remaining_conversations = max(0, max_conversations - segment_count)
        if max_chapters is not None:
            remaining_chapters = max(0, max_chapters - chapter_count)

        has_quota = True
        message = "配额充足"
        if max_conversations is not None and segment_count >= max_conversations:
            has_quota = False
            message = f"对话轮数已用完（{segment_count}/{max_conversations}），请升级 Pro 或 Pro+ 继续使用"
        elif max_chapters is not None and chapter_count >= max_chapters:
            has_quota = False
            message = "章节数量已达上限（免费版仅支持 1 个章节整理），请升级后继续"

        return QuotaCheckResponse(
            has_quota=has_quota,
            remaining_conversations=remaining_conversations,
            remaining_chapters=remaining_chapters,
            remaining_words=remaining_words,
            used_conversations=segment_count,
            used_chapters=chapter_count,
            max_conversations=max_conversations,
            max_chapters=max_chapters,
            message=message,
        )