use super::{ProxyToken, TokenManager};
use crate::proxy::token::loader::OnDiskAccountState;

impl TokenManager {
    pub(super) async fn try_preferred_account_for_request(
        &self,
        tokens_snapshot: &mut Vec<ProxyToken>,
        total: &mut usize,
        target_model: &str,
        quota_protection_enabled: bool,
    ) -> Result<Option<(String, String, String, String, u64)>, String> {
        let preferred_id = self.preferred_account_id.read().await.clone();
        if let Some(ref pref_id) = preferred_id {
            if let Some(preferred_token) = tokens_snapshot
                .iter()
                .find(|t| &t.account_id == pref_id)
                .cloned()
            {
                match crate::proxy::token::loader::get_account_state_on_disk(
                    &preferred_token.account_path,
                )
                .await
                {
                    OnDiskAccountState::Disabled => {
                        tracing::warn!(
                            "🔒  Preferred account {} is disabled on disk, purging and falling back",
                            preferred_token.email
                        );
                        self.remove_account(&preferred_token.account_id);
                        self.prune_preferred_from_snapshot(
                            tokens_snapshot,
                            total,
                            &preferred_token,
                            true,
                            pref_id,
                        )
                        .await?;
                    }
                    OnDiskAccountState::Unknown => {
                        tracing::warn!(
                            "🔒  Preferred account {} state on disk is unavailable, falling back",
                            preferred_token.email
                        );
                        self.prune_preferred_from_snapshot(
                            tokens_snapshot,
                            total,
                            &preferred_token,
                            false,
                            pref_id,
                        )
                        .await?;
                    }
                    OnDiskAccountState::Enabled => {
                        if self
                            .is_preferred_eligible_for_target(
                                &preferred_token,
                                target_model,
                                quota_protection_enabled,
                            )
                            .await
                        {
                            return Ok(Some(
                                self.build_preferred_response(preferred_token.clone()).await,
                            ));
                        }
                    }
                }
            } else {
                tracing::warn!(
                    "🔒  Preferred account {} not found in pool, falling back to round-robin",
                    pref_id
                );
            }
        }

        Ok(None)
    }

    async fn prune_preferred_from_snapshot(
        &self,
        tokens_snapshot: &mut Vec<ProxyToken>,
        total: &mut usize,
        preferred_token: &ProxyToken,
        clear_preferred: bool,
        pref_id: &str,
    ) -> Result<(), String> {
        tokens_snapshot.retain(|t| t.account_id != preferred_token.account_id);
        *total = tokens_snapshot.len();

        if clear_preferred {
            let mut preferred = self.preferred_account_id.write().await;
            if preferred.as_deref() == Some(pref_id) {
                *preferred = None;
            }
        }

        if *total == 0 {
            return Err("Token pool is empty".to_string());
        }
        Ok(())
    }

    async fn is_preferred_eligible_for_target(
        &self,
        preferred_token: &ProxyToken,
        target_model: &str,
        quota_protection_enabled: bool,
    ) -> bool {
        let normalized_target =
            crate::proxy::common::model_mapping::normalize_to_standard_id(target_model)
                .unwrap_or_else(|| target_model.to_string());

        let is_rate_limited = self
            .is_rate_limited(&preferred_token.account_id, Some(&normalized_target))
            .await;
        let is_quota_protected = quota_protection_enabled
            && preferred_token
                .protected_models
                .contains(&normalized_target);

        if !is_rate_limited && !is_quota_protected {
            return true;
        }

        if is_rate_limited {
            tracing::warn!(
                "🔒  Preferred account {} is rate-limited, falling back to round-robin",
                preferred_token.email
            );
        } else {
            tracing::warn!(
                "🔒  Preferred account {} is quota-protected for {}, falling back to round-robin",
                preferred_token.email,
                target_model
            );
        }
        false
    }

    async fn build_preferred_response(
        &self,
        preferred_token: ProxyToken,
    ) -> (String, String, String, String, u64) {
        tracing::info!(
            "🔒  Using preferred account: {} (fixed mode)",
            preferred_token.email
        );
        let mut token = preferred_token;

        let now = chrono::Utc::now().timestamp();
        if crate::modules::auth::oauth::should_refresh_token(
            token.timestamp,
            now,
            Some(&token.account_id),
        ) {
            tracing::debug!(
                "Token for account {} is about to expire, refreshing...",
                token.email
            );
            match crate::modules::auth::oauth::refresh_access_token(
                &token.refresh_token,
                Some(&token.account_id),
            )
            .await
            {
                Ok(token_response) => {
                    self.apply_refreshed_token(&mut token, &token_response, now);
                }
                Err(e) => {
                    tracing::warn!("Preferred account token refresh failed: {}", e);
                }
            }
        }

        let project_id = if let Some(pid) = token
            .project_id
            .as_deref()
            .map(str::trim)
            .filter(|pid| !pid.is_empty())
        {
            pid.to_string()
        } else {
            match crate::proxy::project_resolver::fetch_project_id(
                &token.access_token,
                Some(&token.account_id),
            )
            .await
            {
                Ok(pid) => {
                    let normalized = pid.trim();
                    if normalized.is_empty() {
                        "bamboo-precept-lgxtn".to_string()
                    } else {
                        self.apply_project_id(&mut token, normalized);
                        normalized.to_string()
                    }
                }
                Err(_) => "bamboo-precept-lgxtn".to_string(),
            }
        };

        (
            token.access_token,
            project_id,
            token.email,
            token.account_id,
            0,
        )
    }
}
