Aether/_deprecated_py_src/services/usage/quota_scheduler.py

"""
额度周期重置定时任务

支持按天数周期重置额度：
- quota_reset_day: 重置周期（天数），例如7=每周，30=每月
- quota_last_reset_at: 上次重置时间，用于计算下次重置

使用统一的 TaskScheduler 进行调度。
"""

from __future__ import annotations

from datetime import datetime, timezone
from typing import Any

from src.core.enums import ProviderBillingType
from src.core.logger import logger
from src.database import create_session
from src.models.database import Provider
from src.services.system.scheduler import get_scheduler


class QuotaScheduler:
    """额度周期重置调度器"""

    def __init__(self) -> None:
        self.running = False

    async def start(self) -> Any:
        """启动调度器"""
        if self.running:
            logger.warning("Quota scheduler already running")
            return

        self.running = True
        logger.info("Quota scheduler started")

        scheduler = get_scheduler()

        # 每小时检查一次额度重置
        scheduler.add_interval_job(
            self._scheduled_quota_check,
            hours=1,
            job_id="quota_reset_check",
            name="额度周期重置检查",
        )

        # 启动时立即执行一次检查
        await self._check_and_reset_quotas()

    async def stop(self) -> Any:
        """停止调度器"""
        if not self.running:
            return

        self.running = False
        scheduler = get_scheduler()
        scheduler.remove_job("quota_reset_check")
        logger.info("Quota scheduler stopped")

    async def _scheduled_quota_check(self) -> None:
        """额度检查任务（定时调用）"""
        if not self.running:
            return
        await self._check_and_reset_quotas()

    async def _check_and_reset_quotas(self) -> None:
        """检查并重置周期额度"""

        db = create_session()
        try:
            # 获取所有定额类型的提供商
            providers = (
                db.query(Provider)
                .filter(
                    Provider.billing_type == ProviderBillingType.MONTHLY_QUOTA,
                    Provider.is_active == True,
                )
                .all()
            )

            if not providers:
                logger.debug("No quota providers to check")
                return

            now = datetime.now(timezone.utc)
            reset_count = 0

            for provider in providers:
                try:
                    # 如果没有上次重置时间，初始化为当前时间
                    if provider.quota_last_reset_at is None:
                        provider.quota_last_reset_at = now
                        db.commit()
                        logger.info(f"Initialized quota_last_reset_at for provider {provider.name}")
                        continue

                    # 计算距离上次重置的天数
                    days_since_reset = (now - provider.quota_last_reset_at).days

                    # 如果达到或超过重置周期，执行重置
                    if days_since_reset >= provider.quota_reset_day:
                        logger.info(f"Resetting quota for provider {provider.name}")

                        provider.monthly_used_usd = 0.0
                        provider.quota_last_reset_at = now
                        reset_count += 1

                    # 检查是否过期
                    if provider.quota_expires_at and provider.quota_expires_at < now:
                        logger.warning(f"Provider {provider.name} quota expired")
                        # 可以选择禁用过期的提供商
                        # provider.is_active = False

                except Exception as e:
                    logger.exception(f"Error processing provider {provider.name}: {e}")

            if reset_count > 0:
                db.commit()
                logger.info(f"Reset quotas for {reset_count} providers")
        finally:
            db.close()

    async def force_reset(self, provider_id: str | None = None) -> Any:
        """手动强制重置额度"""
        db = create_session()
        try:
            now = datetime.now(timezone.utc)
            if provider_id:
                # 重置指定提供商
                provider = db.query(Provider).filter(Provider.id == provider_id).first()
                if provider and provider.billing_type == ProviderBillingType.MONTHLY_QUOTA:
                    provider.monthly_used_usd = 0.0
                    provider.quota_last_reset_at = now
                    db.commit()
                    logger.info(f"Force reset quota for provider {provider.name}")
            else:
                # 重置所有定额提供商
                providers = (
                    db.query(Provider)
                    .filter(Provider.billing_type == ProviderBillingType.MONTHLY_QUOTA)
                    .all()
                )
                for provider in providers:
                    provider.monthly_used_usd = 0.0
                    provider.quota_last_reset_at = now
                db.commit()
                logger.info(f"Force reset quotas for {len(providers)} providers")
        finally:
            db.close()


# 全局单例
_quota_scheduler = None


def get_quota_scheduler() -> QuotaScheduler:
    """获取全局调度器实例"""
    global _quota_scheduler
    if _quota_scheduler is None:
        _quota_scheduler = QuotaScheduler()
    return _quota_scheduler