Spaces:
Sleeping
Sleeping
| use axum::{ | |
| extract::Path, | |
| Json, | |
| }; | |
| use serde::{Deserialize, Serialize}; | |
| use crate::models::{Account, TokenData, QuotaData}; | |
| use crate::modules; | |
| use crate::error::AppError; | |
| /// API response wrapper | |
| pub struct ApiResponse<T> { | |
| pub success: bool, | |
| pub data: Option<T>, | |
| pub error: Option<String>, | |
| } | |
| impl<T: Serialize> ApiResponse<T> { | |
| pub fn success(data: T) -> Json<Self> { | |
| Json(Self { | |
| success: true, | |
| data: Some(data), | |
| error: None, | |
| }) | |
| } | |
| pub fn error(message: String) -> Json<Self> { | |
| Json(Self { | |
| success: false, | |
| data: None, | |
| error: Some(message), | |
| }) | |
| } | |
| } | |
| /// List all accounts | |
| pub async fn list_accounts() -> Result<Json<ApiResponse<Vec<Account>>>, AppError> { | |
| let accounts = modules::list_accounts().map_err(AppError::Account)?; | |
| Ok(ApiResponse::success(accounts)) | |
| } | |
| /// Add account request | |
| pub struct AddAccountRequest { | |
| pub email: String, | |
| pub refresh_token: String, | |
| } | |
| /// Add new account | |
| pub async fn add_account( | |
| Json(req): Json<AddAccountRequest>, | |
| ) -> Result<Json<ApiResponse<Account>>, AppError> { | |
| modules::logger::log_info(&format!("Adding account: {}", req.email)); | |
| // 1. Use refresh_token to get access_token | |
| let token_res = modules::oauth::refresh_access_token(&req.refresh_token) | |
| .await | |
| .map_err(AppError::OAuth)?; | |
| // 2. Get user info | |
| let user_info = modules::oauth::get_user_info(&token_res.access_token) | |
| .await | |
| .map_err(AppError::OAuth)?; | |
| // 3. Construct TokenData | |
| let token = TokenData::new( | |
| token_res.access_token, | |
| req.refresh_token, | |
| token_res.expires_in, | |
| Some(user_info.email.clone()), | |
| None, | |
| None, | |
| ); | |
| // 4. Add or update account using real email | |
| let mut account = modules::upsert_account( | |
| user_info.email.clone(), | |
| user_info.get_display_name(), | |
| token, | |
| ).map_err(AppError::Account)?; | |
| modules::logger::log_info(&format!("Account added successfully: {}", account.email)); | |
| // 5. Auto refresh quota | |
| let _ = internal_refresh_account_quota(&mut account).await; | |
| Ok(ApiResponse::success(account)) | |
| } | |
| /// Delete account | |
| pub async fn delete_account( | |
| Path(account_id): Path<String>, | |
| ) -> Result<Json<ApiResponse<()>>, AppError> { | |
| modules::logger::log_info(&format!("Deleting account: {}", account_id)); | |
| modules::delete_account(&account_id).map_err(AppError::Account)?; | |
| modules::logger::log_info(&format!("Account deleted: {}", account_id)); | |
| Ok(ApiResponse::success(())) | |
| } | |
| /// Refresh account quota | |
| pub async fn refresh_quota( | |
| Path(account_id): Path<String>, | |
| ) -> Result<Json<ApiResponse<QuotaData>>, AppError> { | |
| modules::logger::log_info(&format!("Refreshing quota for: {}", account_id)); | |
| let mut account = modules::load_account(&account_id).map_err(AppError::Account)?; | |
| let quota = modules::account::fetch_quota_with_retry(&mut account).await?; | |
| // Update account quota | |
| modules::update_account_quota(&account_id, quota.clone()).map_err(AppError::Account)?; | |
| Ok(ApiResponse::success(quota)) | |
| } | |
| /// Refresh stats response | |
| pub struct RefreshStats { | |
| total: usize, | |
| success: usize, | |
| failed: usize, | |
| details: Vec<String>, | |
| } | |
| /// Refresh all account quotas | |
| pub async fn refresh_all_quotas() -> Result<Json<ApiResponse<RefreshStats>>, AppError> { | |
| modules::logger::log_info("Starting batch quota refresh for all accounts"); | |
| let accounts = modules::list_accounts().map_err(AppError::Account)?; | |
| let mut success = 0; | |
| let mut failed = 0; | |
| let mut details = Vec::new(); | |
| // Serial processing to ensure persistence safety | |
| for mut account in accounts { | |
| if let Some(ref q) = account.quota { | |
| if q.is_forbidden { | |
| modules::logger::log_info(&format!("Skipping {} (Forbidden)", account.email)); | |
| continue; | |
| } | |
| } | |
| modules::logger::log_info(&format!("Processing {}", account.email)); | |
| match modules::account::fetch_quota_with_retry(&mut account).await { | |
| Ok(quota) => { | |
| if let Err(e) = modules::update_account_quota(&account.id, quota) { | |
| failed += 1; | |
| let msg = format!("Account {}: Save quota failed - {}", account.email, e); | |
| details.push(msg.clone()); | |
| modules::logger::log_error(&msg); | |
| } else { | |
| success += 1; | |
| modules::logger::log_info("Success"); | |
| } | |
| }, | |
| Err(e) => { | |
| failed += 1; | |
| let msg = format!("Account {}: Fetch quota failed - {}", account.email, e); | |
| details.push(msg.clone()); | |
| modules::logger::log_error(&msg); | |
| } | |
| } | |
| } | |
| modules::logger::log_info(&format!("Batch refresh completed: {} success, {} failed", success, failed)); | |
| Ok(ApiResponse::success(RefreshStats { total: success + failed, success, failed, details })) | |
| } | |
| /// Get current account | |
| pub async fn get_current_account() -> Result<Json<ApiResponse<Option<Account>>>, AppError> { | |
| let account_id = modules::get_current_account_id().map_err(AppError::Account)?; | |
| if let Some(id) = account_id { | |
| let account = modules::load_account(&id).map_err(AppError::Account)?; | |
| Ok(ApiResponse::success(Some(account))) | |
| } else { | |
| Ok(ApiResponse::success(None)) | |
| } | |
| } | |
| /// Set current account | |
| pub async fn set_current_account( | |
| Path(account_id): Path<String>, | |
| ) -> Result<Json<ApiResponse<()>>, AppError> { | |
| modules::logger::log_info(&format!("Setting current account: {}", account_id)); | |
| modules::set_current_account_id(&account_id).map_err(AppError::Account)?; | |
| Ok(ApiResponse::success(())) | |
| } | |
| /// Internal helper: auto refresh quota after adding account | |
| async fn internal_refresh_account_quota(account: &mut Account) -> Result<QuotaData, String> { | |
| modules::logger::log_info(&format!("Auto refreshing quota: {}", account.email)); | |
| match modules::account::fetch_quota_with_retry(account).await { | |
| Ok(quota) => { | |
| let _ = modules::update_account_quota(&account.id, quota.clone()); | |
| Ok(quota) | |
| }, | |
| Err(e) => { | |
| modules::logger::log_warn(&format!("Auto refresh quota failed ({}): {}", account.email, e)); | |
| Err(e.to_string()) | |
| } | |
| } | |
| } | |