video-gen/video-gen-api/app/services/credit/query_service.py

from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime
from decimal import Decimal
from typing import Iterable

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

from app.enums.credit_balance import CreditBalanceStatus
from app.models.credit.balance import UserCreditBalance
from app.services.credit.time_policy import last_usable_at
from app.services.credit.utils import to_credit_decimal, to_float, utc_now


@dataclass(slots=True, frozen=True)
class CreditBalanceSummary:
    available_credits: Decimal
    next_expiring_credits: Decimal
    next_expires_at: datetime | None
    next_last_usable_at: datetime | None

    def to_dict(self) -> dict:
        return {
            "available_credits": to_float(self.available_credits),
            "credits": to_float(self.available_credits),
            "next_expiring_credits": to_float(self.next_expiring_credits),
            "next_expires_at": self.next_expires_at,
            "next_last_usable_at": self.next_last_usable_at,
        }


async def get_available_credits(
    db: AsyncSession,
    user_id: str,
    *,
    request_time: datetime | None = None,
) -> Decimal:
    checked_at = request_time or utc_now()
    result = await db.execute(
        select(func.coalesce(func.sum(UserCreditBalance.unspent_amount), 0)).where(
            UserCreditBalance.user_id == user_id,
            UserCreditBalance.valid_from <= checked_at,
            UserCreditBalance.expires_at > checked_at,
            UserCreditBalance.unspent_amount > 0,
            UserCreditBalance.revoked_at.is_(None),
        )
    )
    return to_credit_decimal(result.scalar_one())


async def get_balance_summary(
    db: AsyncSession,
    user_id: str,
    *,
    request_time: datetime | None = None,
) -> CreditBalanceSummary:
    checked_at = request_time or utc_now()
    available = await get_available_credits(db, user_id, request_time=checked_at)
    expiry_result = await db.execute(
        select(
            UserCreditBalance.expires_at,
            func.sum(UserCreditBalance.unspent_amount).label("amount"),
        )
        .where(
            UserCreditBalance.user_id == user_id,
            UserCreditBalance.valid_from <= checked_at,
            UserCreditBalance.expires_at > checked_at,
            UserCreditBalance.unspent_amount > 0,
            UserCreditBalance.revoked_at.is_(None),
        )
        .group_by(UserCreditBalance.expires_at)
        .order_by(UserCreditBalance.expires_at.asc())
        .limit(1)
    )
    row = expiry_result.first()
    expires_at = row.expires_at if row else None
    expiring = to_credit_decimal(row.amount if row else 0)
    return CreditBalanceSummary(
        available_credits=available,
        next_expiring_credits=expiring,
        next_expires_at=expires_at,
        next_last_usable_at=last_usable_at(expires_at) if expires_at else None,
    )


async def get_user_credit_map(
    db: AsyncSession,
    user_ids: Iterable[str],
    *,
    request_time: datetime | None = None,
) -> dict[str, float]:
    ids = list(dict.fromkeys(str(item) for item in user_ids if item))
    if not ids:
        return {}
    checked_at = request_time or utc_now()
    result = await db.execute(
        select(
            UserCreditBalance.user_id,
            func.coalesce(func.sum(UserCreditBalance.unspent_amount), 0).label("credits"),
        )
        .where(
            UserCreditBalance.user_id.in_(ids),
            UserCreditBalance.valid_from <= checked_at,
            UserCreditBalance.expires_at > checked_at,
            UserCreditBalance.unspent_amount > 0,
            UserCreditBalance.revoked_at.is_(None),
        )
        .group_by(UserCreditBalance.user_id)
    )
    output = {user_id: 0.0 for user_id in ids}
    for row in result:
        output[str(row.user_id)] = to_float(row.credits)
    return output


def attach_credit_snapshot(user: object, credits: Decimal | float | int) -> object:
    # SQLAlchemy Declarative 对象允许附加非映射运行时属性；不会写回 users 表。
    setattr(user, "credits", to_float(to_credit_decimal(credits)))
    return user


def effective_balance_status(balance: UserCreditBalance, *, request_time: datetime | None = None) -> str:
    checked_at = request_time or utc_now()
    if balance.revoked_at is not None or to_credit_decimal(balance.revoked_amount) > 0:
        return CreditBalanceStatus.REVOKED.value
    if balance.expires_at <= checked_at:
        return CreditBalanceStatus.EXPIRED.value
    if balance.valid_from > checked_at:
        return CreditBalanceStatus.SCHEDULED.value
    if to_credit_decimal(balance.unspent_amount) <= 0:
        return CreditBalanceStatus.CONSUMED.value
    return CreditBalanceStatus.ACTIVE.value


def apply_balance_status_filter(stmt, status: str | None, *, request_time: datetime):
    if not status:
        return stmt
    if status == CreditBalanceStatus.REVOKED.value:
        return stmt.where(
            or_(UserCreditBalance.revoked_at.is_not(None), UserCreditBalance.revoked_amount > 0)
        )
    base_not_revoked = and_(
        UserCreditBalance.revoked_at.is_(None),
        UserCreditBalance.revoked_amount <= 0,
    )
    if status == CreditBalanceStatus.EXPIRED.value:
        return stmt.where(base_not_revoked, UserCreditBalance.expires_at <= request_time)
    if status == CreditBalanceStatus.SCHEDULED.value:
        return stmt.where(
            base_not_revoked,
            UserCreditBalance.valid_from > request_time,
            UserCreditBalance.expires_at > request_time,
        )
    if status == CreditBalanceStatus.CONSUMED.value:
        return stmt.where(
            base_not_revoked,
            UserCreditBalance.valid_from <= request_time,
            UserCreditBalance.expires_at > request_time,
            UserCreditBalance.unspent_amount <= 0,
        )
    if status == CreditBalanceStatus.ACTIVE.value:
        return stmt.where(
            base_not_revoked,
            UserCreditBalance.valid_from <= request_time,
            UserCreditBalance.expires_at > request_time,
            UserCreditBalance.unspent_amount > 0,
        )
    return stmt.where(UserCreditBalance.status == status)