Aether/_deprecated_py_src/services/wallet/daily_usage_ledger.py

from __future__ import annotations

from dataclasses import dataclass
from datetime import date, datetime, time, timedelta, timezone
from decimal import Decimal
from zoneinfo import ZoneInfo

from sqlalchemy import func
from sqlalchemy.orm import Session

from src.config import config
from src.core.logger import logger
from src.models.database import Usage, WalletDailyUsageLedger
from src.services.billing.precision import to_money_decimal

APP_TIMEZONE = config.app_timezone


@dataclass(slots=True)
class WalletDailyUsageSnapshot:
    wallet_id: str | None
    billing_date: date
    billing_timezone: str
    total_cost_usd: Decimal
    total_requests: int
    input_tokens: int
    output_tokens: int
    cache_creation_tokens: int
    cache_read_tokens: int
    first_finalized_at: datetime | None
    last_finalized_at: datetime | None
    aggregated_at: datetime
    is_today: bool = False
    id: str | None = None


class WalletDailyUsageLedgerService:
    """钱包每日消费汇总服务。"""

    @staticmethod
    def get_timezone(timezone_name: str | None = None) -> ZoneInfo:
        tz_name = timezone_name or APP_TIMEZONE
        try:
            return ZoneInfo(tz_name)
        except Exception:
            logger.warning("Invalid billing timezone {}, fallback to UTC", tz_name)
            return ZoneInfo("UTC")

    @classmethod
    def get_today_billing_date(cls, timezone_name: str | None = None) -> date:
        tz = cls.get_timezone(timezone_name)
        return datetime.now(tz).date()

    @classmethod
    def get_day_window_utc(
        cls,
        billing_date: date,
        timezone_name: str | None = None,
    ) -> tuple[datetime, datetime]:
        tz = cls.get_timezone(timezone_name)
        local_start = datetime.combine(billing_date, time.min, tzinfo=tz)
        local_end = local_start + timedelta(days=1)
        return local_start.astimezone(timezone.utc), local_end.astimezone(timezone.utc)

    @classmethod
    def aggregate_day(
        cls,
        db: Session,
        billing_date: date,
        *,
        timezone_name: str | None = None,
        commit: bool = True,
    ) -> int:
        tz_name = timezone_name or APP_TIMEZONE
        start_utc, end_utc = cls.get_day_window_utc(billing_date, tz_name)
        now_utc = datetime.now(timezone.utc)

        rows = (
            db.query(
                Usage.wallet_id.label("wallet_id"),
                func.count(Usage.id).label("total_requests"),
                func.coalesce(func.sum(Usage.total_cost_usd), 0).label("total_cost_usd"),
                func.coalesce(func.sum(Usage.input_tokens), 0).label("input_tokens"),
                func.coalesce(func.sum(Usage.output_tokens), 0).label("output_tokens"),
                func.coalesce(func.sum(Usage.cache_creation_input_tokens), 0).label(
                    "cache_creation_tokens"
                ),
                func.coalesce(func.sum(Usage.cache_read_input_tokens), 0).label(
                    "cache_read_tokens"
                ),
                func.min(Usage.finalized_at).label("first_finalized_at"),
                func.max(Usage.finalized_at).label("last_finalized_at"),
            )
            .filter(
                Usage.wallet_id.isnot(None),
                Usage.billing_status == "settled",
                Usage.total_cost_usd > 0,
                Usage.finalized_at >= start_utc,
                Usage.finalized_at < end_utc,
            )
            .group_by(Usage.wallet_id)
            .all()
        )

        existing_ledgers = (
            db.query(WalletDailyUsageLedger)
            .filter(
                WalletDailyUsageLedger.billing_date == billing_date,
                WalletDailyUsageLedger.billing_timezone == tz_name,
            )
            .all()
        )
        existing_map = {ledger.wallet_id: ledger for ledger in existing_ledgers}
        seen_wallet_ids: set[str] = set()

        for row in rows:
            wallet_id = getattr(row, "wallet_id", None)
            if not wallet_id:
                continue
            seen_wallet_ids.add(str(wallet_id))
            ledger = existing_map.get(str(wallet_id))
            if ledger is None:
                ledger = WalletDailyUsageLedger(
                    wallet_id=str(wallet_id),
                    billing_date=billing_date,
                    billing_timezone=tz_name,
                    aggregated_at=now_utc,
                )
                db.add(ledger)

            ledger.total_cost_usd = to_money_decimal(getattr(row, "total_cost_usd", 0) or 0)
            ledger.total_requests = int(getattr(row, "total_requests", 0) or 0)
            ledger.input_tokens = int(getattr(row, "input_tokens", 0) or 0)
            ledger.output_tokens = int(getattr(row, "output_tokens", 0) or 0)
            ledger.cache_creation_tokens = int(getattr(row, "cache_creation_tokens", 0) or 0)
            ledger.cache_read_tokens = int(getattr(row, "cache_read_tokens", 0) or 0)
            ledger.first_finalized_at = getattr(row, "first_finalized_at", None)
            ledger.last_finalized_at = getattr(row, "last_finalized_at", None)
            ledger.aggregated_at = now_utc

        stale_ledgers = [
            ledger for ledger in existing_ledgers if ledger.wallet_id not in seen_wallet_ids
        ]
        for ledger in stale_ledgers:
            db.delete(ledger)

        if commit:
            db.commit()
        return len(seen_wallet_ids)

    @classmethod
    def get_today_snapshot(
        cls,
        db: Session,
        wallet_id: str | None,
        *,
        timezone_name: str | None = None,
    ) -> WalletDailyUsageSnapshot:
        tz_name = timezone_name or APP_TIMEZONE
        billing_date = cls.get_today_billing_date(tz_name)
        start_utc, end_utc = cls.get_day_window_utc(billing_date, tz_name)
        now_utc = datetime.now(timezone.utc)

        if wallet_id:
            row = (
                db.query(
                    func.count(Usage.id).label("total_requests"),
                    func.coalesce(func.sum(Usage.total_cost_usd), 0).label("total_cost_usd"),
                    func.coalesce(func.sum(Usage.input_tokens), 0).label("input_tokens"),
                    func.coalesce(func.sum(Usage.output_tokens), 0).label("output_tokens"),
                    func.coalesce(func.sum(Usage.cache_creation_input_tokens), 0).label(
                        "cache_creation_tokens"
                    ),
                    func.coalesce(func.sum(Usage.cache_read_input_tokens), 0).label(
                        "cache_read_tokens"
                    ),
                    func.min(Usage.finalized_at).label("first_finalized_at"),
                    func.max(Usage.finalized_at).label("last_finalized_at"),
                )
                .filter(
                    Usage.wallet_id == wallet_id,
                    Usage.billing_status == "settled",
                    Usage.total_cost_usd > 0,
                    Usage.finalized_at >= start_utc,
                    Usage.finalized_at < end_utc,
                )
                .first()
            )
        else:
            row = None

        return WalletDailyUsageSnapshot(
            wallet_id=wallet_id,
            billing_date=billing_date,
            billing_timezone=tz_name,
            total_cost_usd=to_money_decimal(getattr(row, "total_cost_usd", 0) or 0),
            total_requests=int(getattr(row, "total_requests", 0) or 0),
            input_tokens=int(getattr(row, "input_tokens", 0) or 0),
            output_tokens=int(getattr(row, "output_tokens", 0) or 0),
            cache_creation_tokens=int(getattr(row, "cache_creation_tokens", 0) or 0),
            cache_read_tokens=int(getattr(row, "cache_read_tokens", 0) or 0),
            first_finalized_at=getattr(row, "first_finalized_at", None),
            last_finalized_at=getattr(row, "last_finalized_at", None),
            aggregated_at=now_utc,
            is_today=True,
        )