trade-assistant/backend/app/services/payment.py

import hmac
import hashlib
import json
import logging
from typing import Optional, Dict, Any
from datetime import datetime, timedelta
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select
from app.models.subscription import Subscription
from app.models.user import User
from app.config import settings

logger = logging.getLogger(__name__)

PLANS = {
    "free": {"price": 0, "duration_days": None},
    "pro": {"price": 99, "duration_days": 30},
    "enterprise": {"price": 399, "duration_days": 30},
}


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

    async def get_plans(self) -> Dict[str, Any]:
        return {
            "plans": [
                {
                    "id": "free",
                    "name": "免费版",
                    "price": 0,
                    "features": [
                        "1 个产品",
                        "20 次翻译/天",
                        "5 个客户",
                        "基础回复建议",
                    ],
                },
                {
                    "id": "pro",
                    "name": "Pro 版",
                    "price": 99,
                    "features": [
                        "10 个产品",
                        "无限翻译",
                        "50 个客户",
                        "跟进提醒",
                        "报价单生成",
                    ],
                },
                {
                    "id": "enterprise",
                    "name": "企业版",
                    "price": 399,
                    "features": [
                        "无限产品",
                        "多人协作",
                        "品牌报价单",
                        "专属语料训练",
                        "API 接入",
                    ],
                },
            ]
        }

    async def get_current_subscription(self, user_id: str) -> Dict[str, Any]:
        result = await self.db.execute(
            select(Subscription).where(
                Subscription.user_id == user_id,
                Subscription.status == "active",
            ).order_by(Subscription.created_at.desc()).limit(1)
        )
        sub = result.scalar_one_or_none()

        result = await self.db.execute(
            select(User).where(User.id == user_id)
        )
        user = result.scalar_one_or_none()

        return {
            "plan": user.tier if user else "free",
            "status": sub.status if sub else "active",
            "expires_at": sub.expires_at.isoformat() if sub and sub.expires_at else None,
            "auto_renew": sub.auto_renew if sub else False,
        }

    async def create_order(self, user_id: str, plan: str) -> Dict[str, Any]:
        if plan not in PLANS:
            raise ValueError(f"Invalid plan: {plan}")

        plan_info = PLANS[plan]
        if plan_info["price"] == 0:
            result = await self.db.execute(select(User).where(User.id == user_id))
            user = result.scalar_one_or_none()
            if user:
                user.tier = plan
                await self.db.flush()
            return {"status": "ok", "plan": plan, "amount": 0}

        from app.config import settings
        order_id = f"ORD{datetime.utcnow().strftime('%Y%m%d%H%M%S')}{user_id[-6:]}"

        sub = Subscription(
            user_id=user_id,
            plan=plan,
            status="pending",
            amount=plan_info["price"],
            payment_id=order_id,
        )
        self.db.add(sub)
        await self.db.flush()

        pay_params = {
            "appId": settings.WECHAT_APP_ID or "",
            "timeStamp": str(int(datetime.utcnow().timestamp())),
            "nonceStr": hashlib.md5(order_id.encode()).hexdigest()[:16],
            "package": f"prepay_id={order_id}",
            "signType": "MD5",
        }
        sign_str = "&".join(f"{k}={v}" for k, v in sorted(pay_params.items()))
        sign_str += f"&key={settings.SECRET_KEY}"
        pay_params["paySign"] = hashlib.md5(sign_str.encode()).hexdigest().upper()

        return {
            "status": "pending",
            "order_id": order_id,
            "plan": plan,
            "amount": plan_info["price"],
            "currency": "CNY",
            "pay_params": pay_params,
        }

    async def handle_payment_callback(self, payment_id: str, success: bool) -> bool:
        result = await self.db.execute(
            select(Subscription).where(Subscription.payment_id == payment_id)
        )
        sub = result.scalar_one_or_none()
        if not sub:
            return False

        if success:
            sub.status = "active"
            sub.started_at = datetime.utcnow()
            sub.expires_at = datetime.utcnow() + timedelta(days=PLANS[sub.plan]["duration_days"])

            user_result = await self.db.execute(select(User).where(User.id == sub.user_id))
            user = user_result.scalar_one_or_none()
            if user:
                user.tier = sub.plan
        else:
            sub.status = "failed"

        await self.db.flush()
        return True


payment_service = PaymentService