life-echo/api/routers/plans.py

"""
订阅计划相关 API 路由
"""
import os
from fastapi import APIRouter, Depends
from pydantic import BaseModel
from typing import List, Optional
from datetime import datetime
from sqlalchemy.ext.asyncio import AsyncSession

from middleware.auth import get_current_user
from database.models import User
from database import get_async_db

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


class PlanResponse(BaseModel):
    """订阅计划响应"""
    id: str
    name: str
    display_name: str
    price: float
    currency: str
    features: List[str]
    max_conversations: Optional[int] = None  # None表示无限制
    max_chapters: Optional[int] = None
    max_words: Optional[int] = None
    is_popular: bool = False


class CurrentPlanResponse(BaseModel):
    """当前订阅计划响应"""
    plan_id: str
    plan_name: str
    subscription_type: str
    expires_at: Optional[str] = None  # 过期时间，None表示永久
    features: List[str]
    usage: dict  # 使用情况统计


# 预定义的订阅计划
# 免费版：500 轮对话 + 无章节限制
# Pro：88 元，2000 轮对话，无章节限制
# Pro+：288 元，10000 轮对话，无章节限制
AVAILABLE_PLANS = [
    PlanResponse(
        id="free",
        name="free",
        display_name="免费体验版",
        price=0.0,
        currency="CNY",
        features=[
            "500 轮对话",
            "无章节限制",
            "完整回忆录生成流程"
        ],
        max_conversations=500,
        max_chapters=None,
        max_words=None,
        is_popular=False
    ),
    PlanResponse(
        id="pro",
        name="pro",
        display_name="Pro 版",
        price=88.0,
        currency="CNY",
        features=[
            "2000 轮对话",
            "无章节限制",
            "完整回忆录生成"
        ],
        max_conversations=2000,
        max_chapters=None,
        max_words=None,
        is_popular=True
    ),
    PlanResponse(
        id="pro_plus",
        name="pro_plus",
        display_name="Pro+ 版",
        price=288.0,
        currency="CNY",
        features=[
            "10000 轮对话",
            "无章节限制",
            "完整回忆录生成",
            "长期创作无忧"
        ],
        max_conversations=10000,
        max_chapters=None,
        max_words=None,
        is_popular=False
    )
]

# 一分钱测试套餐：仅当 ENABLE_TEST_PLAN=1 时开放，用于开发环境反复测试支付
ENABLE_TEST_PLAN = os.getenv("ENABLE_TEST_PLAN", "").lower() in ("1", "true", "yes")

TEST_PLAN = PlanResponse(
    id="test",
    name="test",
    display_name="一分钱测试版",
    price=0.01,
    currency="CNY",
    features=[
        "无限对话",
        "无限章节整理",
        "仅用于开发环境测试支付"
    ],
    max_conversations=None,
    max_chapters=None,
    max_words=None,
    is_popular=False,
)


def get_plans_for_api() -> List[PlanResponse]:
    """返回对外暴露的套餐列表（含测试套餐当且仅当 ENABLE_TEST_PLAN 开启）。"""
    if ENABLE_TEST_PLAN:
        return AVAILABLE_PLANS + [TEST_PLAN]
    return list(AVAILABLE_PLANS)


def get_plan_by_type(subscription_type: str) -> Optional[PlanResponse]:
    """根据订阅类型获取计划信息。旧字段 premium 按 pro 展示；test 仅当 ENABLE_TEST_PLAN 时有效。"""
    if subscription_type == "premium":
        subscription_type = "pro"
    if subscription_type == "test":
        return TEST_PLAN if ENABLE_TEST_PLAN else AVAILABLE_PLANS[0]
    for plan in AVAILABLE_PLANS:
        if plan.id == subscription_type:
            return plan
    return AVAILABLE_PLANS[0]  # 默认返回免费版


@router.get("", response_model=List[PlanResponse])
async def get_plans():
    """
    获取所有可用的订阅计划（开发环境 ENABLE_TEST_PLAN=1 时包含「一分钱测试版」）。
    """
    return get_plans_for_api()


@router.get("/current", response_model=CurrentPlanResponse)
async def get_current_plan(
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_async_db)
):
    """
    获取当前用户的订阅计划信息
    """
    plan = get_plan_by_type(current_user.subscription_type)
    
    # 计算使用情况（对话轮数 = Segment 数量）
    from routers.quota import get_segment_count, get_chapter_count

    segment_count = await get_segment_count(current_user.id, db)
    chapter_count = await get_chapter_count(current_user.id, db)

    usage = {
        "conversations": segment_count,  # 已用对话轮数
        "chapters": chapter_count,
        "max_conversations": plan.max_conversations,
        "max_chapters": plan.max_chapters,
    }
    
    expires_at = None
    if current_user.subscription_expires_at:
        expires_at = current_user.subscription_expires_at.isoformat()

    return CurrentPlanResponse(
        plan_id=plan.id,
        plan_name=plan.display_name,
        subscription_type=current_user.subscription_type,
        expires_at=expires_at,
        features=plan.features,
        usage=usage
    )