""" Quota API - Get usage quotas for accounts """ from fastapi import APIRouter, HTTPException, BackgroundTasks from pydantic import BaseModel from typing import Optional, Dict, Any import asyncio from services.quota_service import QuotaService from services.token_service import TokenService router = APIRouter() quota_service = QuotaService() token_service = TokenService() class UsageResponse(BaseModel): currentUsage: int usageLimit: int percentageUsed: float daysRemaining: Optional[int] = None resetDate: Optional[str] = None suspended: bool = False trialUsage: Optional[int] = None trialLimit: Optional[int] = None class QuotaResponse(BaseModel): success: bool accountName: Optional[str] = None usage: Optional[UsageResponse] = None error: Optional[str] = None @router.get("/current", response_model=QuotaResponse) async def get_current_quota(): """Get quota for current active account""" try: info = quota_service.get_current_quota() if not info: return QuotaResponse(success=False, error="No active account") if info.error: return QuotaResponse(success=False, error=info.error) usage = info.usage return QuotaResponse( success=True, accountName=info.account_name, usage=UsageResponse( currentUsage=usage.used if usage else 0, usageLimit=usage.limit if usage else 500, percentageUsed=usage.percent_used if usage else 0, daysRemaining=info.days_until_reset, resetDate=info.reset_date.isoformat() if info.reset_date else None, suspended=usage.suspended if usage else False, trialUsage=usage.trial_used if usage else None, trialLimit=usage.trial_limit if usage else None ) ) except Exception as e: return QuotaResponse(success=False, error=str(e)) @router.get("/{filename}", response_model=QuotaResponse) async def get_account_quota(filename: str): """Get quota for specific account""" token = token_service.get_token(filename) if not token: raise HTTPException(status_code=404, detail="Account not found") try: access_token = token.raw_data.get('accessToken') # Refresh if expired if token.is_expired: new_data = token_service.refresh_token(token) access_token = new_data['accessToken'] info = quota_service.get_quota(access_token, token.auth_method) if info.error: return QuotaResponse(success=False, accountName=token.account_name, error=info.error) usage = info.usage return QuotaResponse( success=True, accountName=token.account_name, usage=UsageResponse( currentUsage=usage.used if usage else 0, usageLimit=usage.limit if usage else 500, percentageUsed=usage.percent_used if usage else 0, daysRemaining=info.days_until_reset, resetDate=info.reset_date.isoformat() if info.reset_date else None, suspended=usage.suspended if usage else False ) ) except Exception as e: return QuotaResponse(success=False, accountName=token.account_name, error=str(e)) @router.get("/all/summary") async def get_all_quotas(): """Get quota summary for all accounts""" tokens = token_service.list_tokens() results = [] for token in tokens: try: access_token = token.raw_data.get('accessToken') if token.is_expired: try: new_data = token_service.refresh_token(token) access_token = new_data['accessToken'] except Exception: results.append({ "accountName": token.account_name, "error": "Token expired and refresh failed" }) continue info = quota_service.get_quota(access_token, token.auth_method) if info.error: results.append({ "accountName": token.account_name, "error": info.error }) else: usage = info.usage results.append({ "accountName": token.account_name, "currentUsage": usage.used if usage else 0, "usageLimit": usage.limit if usage else 500, "percentageUsed": usage.percent_used if usage else 0, "daysRemaining": info.days_until_reset, "suspended": usage.suspended if usage else False }) except Exception as e: results.append({ "accountName": token.account_name, "error": str(e) }) return {"accounts": results, "total": len(results)}