life-echo/api/routers/quota.py

"""
配额检查 API 路由

「对话轮数」的定义：每条用户发出的消息（Segment 表的记录数）计为 1 轮。
"""
from fastapi import APIRouter, Depends
from pydantic import BaseModel
from typing import Optional

from middleware.auth import get_current_user
from database.models import User
from sqlalchemy.ext.asyncio import AsyncSession
from database import get_async_db
from sqlalchemy import select, func

router = APIRouter(prefix="/api/quota", tags=["quota"])


class QuotaCheckResponse(BaseModel):
    """配额检查响应"""
    has_quota: bool
    remaining_conversations: Optional[int] = None  # 剩余对话轮数
    remaining_chapters: Optional[int] = None
    remaining_words: Optional[int] = None
    # 已用量（前端展示 "已用 X / 共 Y"）
    used_conversations: Optional[int] = None
    used_chapters: Optional[int] = None
    max_conversations: Optional[int] = None
    max_chapters: Optional[int] = None
    message: str


# 计划配额限制（与 plans.py 中 AVAILABLE_PLANS 一致）
# 免费：50 轮对话 + 1 个章节整理；Pro：2000 轮无章节限制；Pro+：10000 轮无章节限制
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
    },
    # 一分钱测试版（仅开发环境 ENABLE_TEST_PLAN=1）：无限对话、无限章节
    "test": {
        "max_conversations": None,
        "max_chapters": None,
        "max_words": None
    },
}


async def get_segment_count(user_id: str, db: AsyncSession) -> int:
    """
    获取用户已消耗的对话轮数（= 该用户所有 Segment 记录数）。
    每条 Segment 对应一次用户发送的消息（文本/语音）。
    """
    from database.models import Segment, Conversation
    stmt = (
        select(func.count(Segment.id))
        .join(Conversation, Segment.conversation_id == Conversation.id)
        .where(Conversation.user_id == user_id)
    )
    result = await db.execute(stmt)
    return result.scalar() or 0


async def get_chapter_count(user_id: str, db: AsyncSession) -> int:
    """获取用户当前章节数量"""
    from database.models import Chapter
    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_create_conversation(
    subscription_type: str,
    conversation_count: int
) -> tuple[bool, str]:
    return check_can_send_message(subscription_type, conversation_count)


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, ""


@router.get("/check", response_model=QuotaCheckResponse)
async def check_quota(
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_async_db)
):
    """
    检查用户配额使用情况
    """
    plan_type = current_user.subscription_type
    quotas = PLAN_QUOTAS.get(plan_type, PLAN_QUOTAS["free"])

    # 统计已用量
    segment_count = await get_segment_count(current_user.id, db)
    chapter_count = await get_chapter_count(current_user.id, db)

    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
    )