use std::fs; use std::path::PathBuf; use uuid::Uuid; use crate::models::{Account, AccountIndex, AccountSummary, TokenData, QuotaData}; use crate::modules; use once_cell::sync::Lazy; use std::sync::Mutex; /// Global account write lock to prevent concurrent index file corruption static ACCOUNT_INDEX_LOCK: Lazy> = Lazy::new(|| Mutex::new(())); const DATA_DIR_LOCAL: &str = ".antigravity_tools"; const DATA_DIR_CLOUD: &str = "/data"; const ACCOUNTS_INDEX: &str = "accounts.json"; const ACCOUNTS_DIR: &str = "accounts"; /// Get data directory path (cloud-compatible) /// Priority: /data (HF persistent) > ./data (current dir) > ~/.antigravity_tools (fallback) pub fn get_data_dir() -> Result { // Check for cloud environment (/data) let cloud_path = PathBuf::from(DATA_DIR_CLOUD); if cloud_path.exists() { modules::logger::log_info(&format!("[get_data_dir] Using cloud path: {:?}", cloud_path)); return Ok(cloud_path); } // Check for ./data (current working directory, used in Docker) let cwd_data = PathBuf::from("./data"); if cwd_data.exists() { let canonical = cwd_data.canonicalize() .map_err(|e| format!("Failed to canonicalize ./data: {}", e))?; modules::logger::log_info(&format!("[get_data_dir] Using cwd data path: {:?}", canonical)); return Ok(canonical); } // Fallback to local development path let home = dirs::home_dir().ok_or("Cannot get user home directory")?; let data_dir = home.join(DATA_DIR_LOCAL); // Ensure directory exists if !data_dir.exists() { fs::create_dir_all(&data_dir) .map_err(|e| format!("Failed to create data directory: {}", e))?; } modules::logger::log_info(&format!("[get_data_dir] Using local path: {:?}", data_dir)); Ok(data_dir) } /// Get accounts directory path pub fn get_accounts_dir() -> Result { let data_dir = get_data_dir()?; let accounts_dir = data_dir.join(ACCOUNTS_DIR); if !accounts_dir.exists() { fs::create_dir_all(&accounts_dir) .map_err(|e| format!("Failed to create accounts directory: {}", e))?; } Ok(accounts_dir) } /// Load account index pub fn load_account_index() -> Result { let data_dir = get_data_dir()?; let index_path = data_dir.join(ACCOUNTS_INDEX); if !index_path.exists() { modules::logger::log_warn("Account index file does not exist"); return Ok(AccountIndex::new()); } let content = fs::read_to_string(&index_path) .map_err(|e| format!("Failed to read account index: {}", e))?; let index: AccountIndex = serde_json::from_str(&content) .map_err(|e| format!("Failed to parse account index: {}", e))?; modules::logger::log_info(&format!("Loaded index with {} accounts", index.accounts.len())); Ok(index) } /// Save account index (atomic write) pub fn save_account_index(index: &AccountIndex) -> Result<(), String> { let data_dir = get_data_dir()?; let index_path = data_dir.join(ACCOUNTS_INDEX); let temp_path = data_dir.join(format!("{}.tmp", ACCOUNTS_INDEX)); modules::logger::log_info(&format!("[save_account_index] Saving to: {:?}", index_path)); let content = serde_json::to_string_pretty(index) .map_err(|e| format!("Failed to serialize account index: {}", e))?; // Write to temp file fs::write(&temp_path, &content) .map_err(|e| format!("Failed to write temp index file: {}", e))?; // Atomic rename fs::rename(&temp_path, &index_path) .map_err(|e| format!("Failed to replace index file: {}", e))?; modules::logger::log_info(&format!("[save_account_index] Saved {} accounts successfully", index.accounts.len())); Ok(()) } /// Load account data pub fn load_account(account_id: &str) -> Result { let accounts_dir = get_accounts_dir()?; let account_path = accounts_dir.join(format!("{}.json", account_id)); if !account_path.exists() { return Err(format!("Account not found: {}", account_id)); } let content = fs::read_to_string(&account_path) .map_err(|e| format!("Failed to read account data: {}", e))?; serde_json::from_str(&content) .map_err(|e| format!("Failed to parse account data: {}", e)) } /// Save account data pub fn save_account(account: &Account) -> Result<(), String> { let accounts_dir = get_accounts_dir()?; let account_path = accounts_dir.join(format!("{}.json", account.id)); let content = serde_json::to_string_pretty(account) .map_err(|e| format!("Failed to serialize account data: {}", e))?; fs::write(&account_path, content) .map_err(|e| format!("Failed to save account data: {}", e)) } /// List all accounts pub fn list_accounts() -> Result, String> { modules::logger::log_info("Listing accounts..."); let mut index = load_account_index()?; let mut accounts = Vec::new(); let mut invalid_ids = Vec::new(); for summary in &index.accounts { match load_account(&summary.id) { Ok(account) => accounts.push(account), Err(e) => { modules::logger::log_error(&format!("Failed to load account {}: {}", summary.id, e)); if e.contains("Account not found") || e.contains("Os { code: 2,") || e.contains("No such file") { invalid_ids.push(summary.id.clone()); } }, } } // Auto-fix index: remove invalid account IDs if !invalid_ids.is_empty() { modules::logger::log_warn(&format!("Found {} invalid account indexes, cleaning up...", invalid_ids.len())); index.accounts.retain(|s| !invalid_ids.contains(&s.id)); if let Some(current_id) = &index.current_account_id { if invalid_ids.contains(current_id) { index.current_account_id = index.accounts.first().map(|s| s.id.clone()); } } if let Err(e) = save_account_index(&index) { modules::logger::log_error(&format!("Failed to clean up index: {}", e)); } else { modules::logger::log_info("Index cleanup completed"); } } Ok(accounts) } /// Add account pub fn add_account(email: String, name: Option, token: TokenData) -> Result { let _lock = ACCOUNT_INDEX_LOCK.lock().map_err(|e| format!("Failed to get lock: {}", e))?; let mut index = load_account_index()?; // Check if already exists if index.accounts.iter().any(|s| s.email == email) { return Err(format!("Account already exists: {}", email)); } // Create new account let account_id = Uuid::new_v4().to_string(); let mut account = Account::new(account_id.clone(), email.clone(), token); account.name = name.clone(); // Save account data save_account(&account)?; // Update index index.accounts.push(AccountSummary { id: account_id.clone(), email: email.clone(), name: name.clone(), created_at: account.created_at, last_used: account.last_used, }); // If first account, set as current if index.current_account_id.is_none() { index.current_account_id = Some(account_id); } save_account_index(&index)?; Ok(account) } /// Add or update account pub fn upsert_account(email: String, name: Option, token: TokenData) -> Result { let _lock = ACCOUNT_INDEX_LOCK.lock().map_err(|e| format!("Failed to get lock: {}", e))?; let mut index = load_account_index()?; // Find account ID if exists let existing_account_id = index.accounts.iter() .find(|s| s.email == email) .map(|s| s.id.clone()); if let Some(account_id) = existing_account_id { // Update existing account match load_account(&account_id) { Ok(mut account) => { account.token = token; account.name = name.clone(); account.update_last_used(); save_account(&account)?; // Sync update name in index if let Some(idx_summary) = index.accounts.iter_mut().find(|s| s.id == account_id) { idx_summary.name = name; save_account_index(&index)?; } return Ok(account); }, Err(e) => { modules::logger::log_warn(&format!("Account {} file missing ({}), recreating...", account_id, e)); // Index exists but file missing, recreate let mut account = Account::new(account_id.clone(), email.clone(), token); account.name = name.clone(); save_account(&account)?; if let Some(idx_summary) = index.accounts.iter_mut().find(|s| s.id == account_id) { idx_summary.name = name; save_account_index(&index)?; } return Ok(account); } } } // Not exists, add new drop(_lock); add_account(email, name, token) } /// Delete account pub fn delete_account(account_id: &str) -> Result<(), String> { let _lock = ACCOUNT_INDEX_LOCK.lock().map_err(|e| format!("Failed to get lock: {}", e))?; let mut index = load_account_index()?; // Remove from index let original_len = index.accounts.len(); index.accounts.retain(|s| s.id != account_id); if index.accounts.len() == original_len { return Err(format!("Account ID not found: {}", account_id)); } // If current account, clear it if index.current_account_id.as_deref() == Some(account_id) { index.current_account_id = index.accounts.first().map(|s| s.id.clone()); } save_account_index(&index)?; // Delete account file let accounts_dir = get_accounts_dir()?; let account_path = accounts_dir.join(format!("{}.json", account_id)); if account_path.exists() { fs::remove_file(&account_path) .map_err(|e| format!("Failed to delete account file: {}", e))?; } Ok(()) } /// Batch delete accounts (atomic index operation) pub fn delete_accounts(account_ids: &[String]) -> Result<(), String> { let _lock = ACCOUNT_INDEX_LOCK.lock().map_err(|e| format!("Failed to get lock: {}", e))?; let mut index = load_account_index()?; let accounts_dir = get_accounts_dir()?; for account_id in account_ids { // Remove from index index.accounts.retain(|s| &s.id != account_id); // If current account, clear it if index.current_account_id.as_deref() == Some(account_id) { index.current_account_id = None; } // Delete account file let account_path = accounts_dir.join(format!("{}.json", account_id)); if account_path.exists() { let _ = fs::remove_file(&account_path); } } // If current account is empty, try to select first as default if index.current_account_id.is_none() { index.current_account_id = index.accounts.first().map(|s| s.id.clone()); } save_account_index(&index) } /// Get current account ID pub fn get_current_account_id() -> Result, String> { let index = load_account_index()?; Ok(index.current_account_id) } /// Get current active account info pub fn get_current_account() -> Result, String> { if let Some(id) = get_current_account_id()? { Ok(Some(load_account(&id)?)) } else { Ok(None) } } /// Set current active account ID pub fn set_current_account_id(account_id: &str) -> Result<(), String> { let _lock = ACCOUNT_INDEX_LOCK.lock().map_err(|e| format!("Failed to get lock: {}", e))?; let mut index = load_account_index()?; index.current_account_id = Some(account_id.to_string()); save_account_index(&index) } /// Update account quota pub fn update_account_quota(account_id: &str, quota: QuotaData) -> Result<(), String> { let mut account = load_account(account_id)?; account.update_quota(quota); save_account(&account) } /// Export all account refresh_tokens #[allow(dead_code)] pub fn export_accounts() -> Result, String> { let accounts = list_accounts()?; let mut exports = Vec::new(); for account in accounts { exports.push((account.email, account.token.refresh_token)); } Ok(exports) } /// Fetch quota with retry mechanism pub async fn fetch_quota_with_retry(account: &mut Account) -> crate::error::AppResult { use crate::modules::oauth; use crate::error::AppError; use reqwest::StatusCode; // 1. Time-based check - ensure token is valid let token = oauth::ensure_fresh_token(&account.token).await.map_err(AppError::OAuth)?; if token.access_token != account.token.access_token { modules::logger::log_info(&format!("Token refreshed for: {}", account.email)); account.token = token.clone(); // Re-fetch user name if missing let name = if account.name.is_none() || account.name.as_ref().map_or(false, |n| n.trim().is_empty()) { match oauth::get_user_info(&token.access_token).await { Ok(user_info) => user_info.get_display_name(), Err(_) => None } } else { account.name.clone() }; account.name = name.clone(); upsert_account(account.email.clone(), name, token.clone()).map_err(AppError::Account)?; } // 0. Fill user name if missing if account.name.is_none() || account.name.as_ref().map_or(false, |n| n.trim().is_empty()) { modules::logger::log_info(&format!("Account {} missing name, fetching...", account.email)); match oauth::get_user_info(&account.token.access_token).await { Ok(user_info) => { let display_name = user_info.get_display_name(); modules::logger::log_info(&format!("Got user name: {:?}", display_name)); account.name = display_name.clone(); if let Err(e) = upsert_account(account.email.clone(), display_name, account.token.clone()) { modules::logger::log_warn(&format!("Failed to save user name: {}", e)); } }, Err(e) => { modules::logger::log_warn(&format!("Failed to get user name: {}", e)); } } } // 2. Try to query quota let result: crate::error::AppResult<(QuotaData, Option)> = modules::fetch_quota(&account.token.access_token, &account.email).await; // Capture possible project_id update and save if let Ok((ref _q, ref project_id)) = result { if project_id.is_some() && *project_id != account.token.project_id { modules::logger::log_info(&format!("Detected project_id update ({}), saving...", account.email)); account.token.project_id = project_id.clone(); if let Err(e) = upsert_account(account.email.clone(), account.name.clone(), account.token.clone()) { modules::logger::log_warn(&format!("Failed to save project_id: {}", e)); } } } // 3. Handle 401 error if let Err(AppError::Network(ref e)) = result { if let Some(status) = e.status() { if status == StatusCode::UNAUTHORIZED { modules::logger::log_warn(&format!("401 Unauthorized for {}, forcing refresh...", account.email)); // Force refresh let token_res = oauth::refresh_access_token(&account.token.refresh_token) .await .map_err(AppError::OAuth)?; let new_token = TokenData::new( token_res.access_token.clone(), account.token.refresh_token.clone(), token_res.expires_in, account.token.email.clone(), account.token.project_id.clone(), None, ); // Re-fetch user name let name = if account.name.is_none() || account.name.as_ref().map_or(false, |n| n.trim().is_empty()) { match oauth::get_user_info(&token_res.access_token).await { Ok(user_info) => user_info.get_display_name(), Err(_) => None } } else { account.name.clone() }; account.token = new_token.clone(); account.name = name.clone(); upsert_account(account.email.clone(), name, new_token.clone()).map_err(AppError::Account)?; // Retry query let retry_result: crate::error::AppResult<(QuotaData, Option)> = modules::fetch_quota(&new_token.access_token, &account.email).await; // Also handle retry project_id save if let Ok((ref _q, ref project_id)) = retry_result { if project_id.is_some() && *project_id != account.token.project_id { modules::logger::log_info(&format!("Detected retry project_id update ({}), saving...", account.email)); account.token.project_id = project_id.clone(); let _ = upsert_account(account.email.clone(), account.name.clone(), account.token.clone()); } } if let Err(AppError::Network(ref e)) = retry_result { if let Some(s) = e.status() { if s == StatusCode::FORBIDDEN { let mut q = QuotaData::new(); q.is_forbidden = true; return Ok(q); } } } return retry_result.map(|(q, _)| q); } } } result.map(|(q, _)| q) }