server/src/modules/account.rs

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<Mutex<()>> = 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<PathBuf, String> {
    // 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<PathBuf, String> {
    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<AccountIndex, String> {
    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<Account, String> {
    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<Vec<Account>, 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<String>, token: TokenData) -> Result<Account, String> {
    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<String>, token: TokenData) -> Result<Account, String> {
    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<Option<String>, String> {
    let index = load_account_index()?;
    Ok(index.current_account_id)
}

/// Get current active account info
pub fn get_current_account() -> Result<Option<Account>, 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<Vec<(String, String)>, 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<QuotaData> {
    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<String>)> = 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<String>)> = 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)
}