magi-code 0.96.1

use crate::auth;
use std::thread::JoinHandle;
use std::time::{Duration, Instant};

use super::MissionControlApp;
use crate::auth::CredentialReadiness;
use crate::config::McPaths;
use crate::providers::codex_usage::{
    AccountCodexUsage, CodexAccountIdentity, CodexUsage, load_codex_usage,
};
use crate::tui::TuiEvent;
use crate::tui::state::MissionControlState;
use crate::tui::usage::{UsageLoadResult, usage_report};

const REFRESH_INTERVAL: Duration = Duration::from_secs(300);
const CONNECTION_CHECK_INTERVAL: Duration = Duration::from_secs(5);

#[derive(Default)]
pub(super) struct CodexQuotaRefresh {
    worker: Option<JoinHandle<QuotaCheck>>,
    account: Option<CodexAccountIdentity>,
    last_attempt: Option<Instant>,
    usage: Option<CodexUsage>,
    last_label_update: Option<Instant>,
    last_check: Option<Instant>,
    modal_requested: bool,
    discard_pending: bool,
}

struct QuotaCheck {
    account: Option<CodexAccountIdentity>,
    result: Option<UsageLoadResult>,
    usage: Option<CodexUsage>,
    attempted_at: Option<Instant>,
}

impl CodexQuotaRefresh {
    pub(super) fn request_modal_refresh(&mut self) {
        self.modal_requested = true;
    }

    pub(super) fn invalidate(&mut self) {
        self.account = None;
        self.last_check = None;
        self.last_attempt = None;
        self.usage = None;
        self.last_label_update = None;
        self.discard_pending = self.worker.is_some();
    }

    pub(super) fn cleanup(&mut self) -> Option<String> {
        let worker = self.worker.take()?;
        // Like decorative startup, this worker owns only cloned paths/data and a
        // wake sender. HTTP has its own timeout; it never owns terminal state.
        match crate::thread_join::join_with_timeout(worker, Duration::from_millis(50)) {
            Ok(Err(_)) => Some("Codex usage worker panicked".into()),
            _ => None,
        }
    }

    fn refresh_due(&self, now: Instant, activity: Option<Instant>) -> bool {
        self.worker.is_none()
            && self.account.is_some()
            && self
                .last_attempt
                .is_none_or(|last| now.duration_since(last) >= REFRESH_INTERVAL)
            && activity.is_some_and(|last| now.duration_since(last) < REFRESH_INTERVAL)
    }
}

impl MissionControlApp {
    pub(super) fn poll_codex_quota(
        &mut self,
        state: &mut MissionControlState,
        now: Instant,
    ) -> bool {
        if state.running_prompt.is_some() {
            state.last_agent_activity = Some(now);
        }
        let mut changed = false;
        if self
            .codex_quota
            .worker
            .as_ref()
            .is_some_and(|worker| worker.is_finished())
        {
            let result = self
                .codex_quota
                .worker
                .take()
                .expect("finished worker")
                .join();
            if self.codex_quota.discard_pending {
                self.codex_quota.discard_pending = false;
            } else {
                let check = result.unwrap_or_else(|_| QuotaCheck {
                    account: self.codex_quota.account,
                    result: Some(UsageLoadResult::Error("Codex usage worker failed".into())),
                    usage: None,
                    attempted_at: Some(now),
                });
                changed = self.apply_codex_quota_check(state, check);
            }
        }
        if self.codex_quota.usage.is_some()
            && self
                .codex_quota
                .last_label_update
                .is_none_or(|last| now.duration_since(last) >= Duration::from_secs(1))
        {
            self.codex_quota.last_label_update = Some(now);
            let label = self
                .codex_quota
                .usage
                .as_ref()
                .and_then(|usage| codex_quota_label(usage, chrono::Utc::now()));
            if label != state.codex_quota_label {
                state.codex_quota_label = label;
                changed = true;
            }
        }
        let refresh_due = self.codex_quota.refresh_due(now, state.last_agent_activity);
        let check_due = self
            .codex_quota
            .last_check
            .is_none_or(|last| now.duration_since(last) >= CONNECTION_CHECK_INTERVAL);
        if self.codex_quota.worker.is_none()
            && (check_due || refresh_due || self.codex_quota.modal_requested)
        {
            let paths = self.config.paths.clone();
            let account = self.codex_quota.account;
            let force = self.codex_quota.modal_requested;
            let sender = self.events.clone();
            self.codex_quota.last_check = Some(now);
            self.codex_quota.worker = Some(std::thread::spawn(move || {
                let check = check_codex_quota(&paths, account, refresh_due || force, force);
                // The join result is authoritative even if this wake is dropped.
                let _ = sender.try_send(TuiEvent::WorkerOutcomeReady);
                check
            }));
        }
        changed
    }

    fn apply_codex_quota_check(
        &mut self,
        state: &mut MissionControlState,
        check: QuotaCheck,
    ) -> bool {
        let connection_changed = self.codex_quota.account != check.account;
        self.codex_quota.account = check.account;
        if connection_changed {
            self.codex_quota.usage = None;
            if check.account.is_none() {
                self.codex_quota.last_attempt = None;
            }
            state.codex_quota_label = None;
        }
        if let Some(attempted_at) = check.attempted_at {
            self.codex_quota.last_attempt = Some(attempted_at);
        }
        let has_result = check.result.is_some();
        if has_result || check.account.is_none() {
            self.codex_quota.usage = check.usage;
            state.codex_quota_label = self
                .codex_quota
                .usage
                .as_ref()
                .and_then(|usage| codex_quota_label(usage, chrono::Utc::now()));
            self.codex_quota.last_label_update = Some(Instant::now());
        }
        if let Some(result) = check.result
            && self.codex_quota.modal_requested
        {
            self.codex_quota.modal_requested = false;
            if let Some(request_id) = state.pending_usage_request_id {
                let mut drain = super::super::DrainResult::default();
                super::super::apply_tui_event_to_state(
                    state,
                    TuiEvent::UsageLoaded { request_id, result },
                    &mut drain,
                );
            }
        }
        connection_changed || has_result
    }
}

fn connected_account(paths: &McPaths) -> Option<CodexAccountIdentity> {
    let store = auth::read_auth_store(paths).ok()?;
    let provider = crate::providers::OPENAI_CODEX_PROVIDER;
    let readiness = auth::classify_provider_auth_record(
        provider,
        store.auth().providers.get(provider),
        chrono::Utc::now().timestamp(),
    );
    if !matches!(
        readiness,
        CredentialReadiness::Ready | CredentialReadiness::Refreshable
    ) {
        return None;
    }
    let auth::AuthProviderRecord::OAuth {
        access, account_id, ..
    } = store.auth().providers.get(provider)?
    else {
        return None;
    };
    CodexAccountIdentity::from_oauth(access, account_id.as_deref())
}

fn check_codex_quota(
    paths: &McPaths,
    previous_account: Option<CodexAccountIdentity>,
    refresh: bool,
    manual: bool,
) -> QuotaCheck {
    check_codex_quota_with_loader(paths, previous_account, refresh, manual, load_codex_usage)
}

fn check_codex_quota_with_loader(
    paths: &McPaths,
    previous_account: Option<CodexAccountIdentity>,
    refresh: bool,
    manual: bool,
    load: impl FnOnce(&McPaths) -> Result<AccountCodexUsage, String>,
) -> QuotaCheck {
    let account = connected_account(paths);
    let mut check = QuotaCheck {
        account,
        result: None,
        usage: None,
        attempted_at: None,
    };
    if account.is_none() {
        if manual {
            check.result = Some(UsageLoadResult::Error(
                "Connect openai-codex with /login first".into(),
            ));
        }
        return check;
    }
    // A credential rotation must not bypass the five-minute activity gate.
    // Initial connection (or an explicit UI reconnect) has no previous account.
    if !refresh && previous_account.is_some() {
        return check;
    }
    check.attempted_at = Some(Instant::now());
    let usage = load(paths);
    // Compare against the account used by the final request, not the token revision
    // before HTTP: this request may have refreshed its own credentials.
    check.account = connected_account(paths);
    if check.account.is_none()
        || usage
            .as_ref()
            .is_ok_and(|usage| Some(usage.account) != check.account)
    {
        check.result = Some(UsageLoadResult::Error(
            "Codex connection changed; retry /usage".into(),
        ));
        return check;
    }
    check.result = Some(match usage {
        Ok(usage) => {
            check.usage = Some(usage.usage.clone());
            UsageLoadResult::Loaded(usage_report(usage.usage))
        }
        Err(error) => UsageLoadResult::Error(error),
    });
    check
}

fn codex_quota_label(usage: &CodexUsage, now: chrono::DateTime<chrono::Utc>) -> Option<String> {
    let weekly = usage
        .plan
        .as_deref()
        .is_some_and(|plan| plan.eq_ignore_ascii_case("pro"));
    let seconds = if weekly { 604_800 } else { 18_000 };
    let window = usage
        .windows
        .iter()
        .find(|window| window.limit_window_seconds == Some(seconds))
        .or_else(|| usage.windows.first())?;
    let weekly = window.limit_window_seconds == Some(604_800);
    let suffix = match window.limit_window_seconds {
        Some(604_800) => "w",
        Some(18_000) => "5h",
        _ => window.label.as_str(),
    };
    let reset = window
        .reset_at
        .and_then(|timestamp| chrono::DateTime::from_timestamp(timestamp, 0))
        .map(|reset| reset.signed_duration_since(now).num_seconds())
        .filter(|&seconds| seconds > 0)
        .map(|seconds| {
            let minutes = seconds.saturating_add(59) / 60;
            if weekly {
                format!(
                    "{}d{}h{}m",
                    minutes / 1_440,
                    minutes / 60 % 24,
                    minutes % 60
                )
            } else if minutes < 60 {
                format!("{minutes}m")
            } else {
                format!("{}h{}m", minutes / 60, minutes % 60)
            }
        });
    Some(format!(
        "Codex {:.0}/100%({})",
        window.used_percent,
        reset.as_deref().unwrap_or(suffix)
    ))
}