magi-code 0.80.0

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

use super::MissionControlApp;
use crate::config::{self, CredentialReadiness, 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>,
    last_check: Option<Instant>,
    modal_requested: bool,
    discard_pending: bool,
}

struct QuotaCheck {
    account: Option<CodexAccountIdentity>,
    result: Option<UsageLoadResult>,
    label: Option<String>,
    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.discard_pending = self.worker.is_some();
    }

    pub(super) fn cleanup(&mut self) -> Option<String> {
        let worker = self.worker.take()?;
        let deadline = Instant::now() + Duration::from_millis(50);
        while !worker.is_finished() && Instant::now() < deadline {
            std::thread::sleep(Duration::from_millis(5));
        }
        // Like decorative startup, this worker owns only cloned paths/data and a
        // wake sender. HTTP has its own timeout; it never owns terminal state.
        if worker.is_finished() && worker.join().is_err() {
            return 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())),
                    label: None,
                    attempted_at: Some(now),
                });
                changed = self.apply_codex_quota_check(state, check);
            }
        }
        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 {
            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() {
            state.codex_quota_label = check.label;
        }
        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 = config::read_auth_store(paths).ok()?;
    let provider = crate::providers::OPENAI_CODEX_PROVIDER;
    let readiness = config::classify_provider_auth_record(
        provider,
        store.auth().providers.get(provider),
        chrono::Utc::now().timestamp(),
    );
    if !matches!(
        readiness,
        CredentialReadiness::Ready | CredentialReadiness::Refreshable
    ) {
        return None;
    }
    let config::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,
        label: 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.label = codex_quota_label(&usage.usage);
            UsageLoadResult::Loaded(usage_report(usage.usage))
        }
        Err(error) => UsageLoadResult::Error(error),
    });
    check
}

fn codex_quota_label(usage: &CodexUsage) -> 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))?;
    let suffix = if weekly { "w" } else { "5h" };
    Some(format!("Codex {:.0}/100%({suffix})", window.used_percent))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::providers::codex_usage::CodexUsageWindow;

    fn test_account() -> Option<CodexAccountIdentity> {
        CodexAccountIdentity::from_oauth("test-token", Some("test-account"))
    }

    fn store_account(paths: &McPaths, account_id: &str, access: &str) {
        config::persist_codex_token(
            paths,
            config::NormalizedToken {
                access: access.into(),
                refresh: Some("test-refresh".into()),
                expires: Some(chrono::Utc::now().timestamp() + 3600),
                account_id: account_id.into(),
            },
        )
        .unwrap();
    }

    #[test]
    fn own_refresh_is_accepted_but_account_switch_or_logout_during_http_is_rejected() {
        for connected in [Some("test-account"), Some("other-account"), None] {
            let temp = tempfile::TempDir::new().unwrap();
            let paths = McPaths::from_root(temp.path().join("mc"));
            store_account(&paths, "test-account", "old-token");
            let check =
                check_codex_quota_with_loader(&paths, test_account(), true, true, |paths| {
                    // The successful request uses this refreshed account.
                    store_account(paths, "test-account", "refreshed-token");
                    let account = connected_account(paths).unwrap();
                    if let Some(connected) = connected {
                        store_account(paths, connected, "another-token");
                    } else {
                        config::write_auth(paths, &config::Auth::default()).unwrap();
                    }
                    Ok(AccountCodexUsage {
                        account,
                        usage: CodexUsage {
                            plan: Some("plus".into()),
                            windows: vec![CodexUsageWindow {
                                label: "5h".into(),
                                limit_window_seconds: Some(18_000),
                                used_percent: 12.0,
                                reset_at: None,
                            }],
                        },
                    })
                });
            if connected == Some("test-account") {
                assert!(matches!(check.result, Some(UsageLoadResult::Loaded(_))));
                assert_eq!(check.label.as_deref(), Some("Codex 12/100%(5h)"));
            } else {
                assert!(matches!(check.result, Some(UsageLoadResult::Error(_))));
                assert!(check.label.is_none());
            }
        }
    }

    #[test]
    fn same_account_rotation_preserves_quota_without_bypassing_activity_gate() {
        let temp = tempfile::TempDir::new().unwrap();
        let (sender, _receiver) = crossbeam_channel::bounded(1);
        let mut app = super::super::tests::test_app(&temp, sender);
        store_account(&app.config.paths, "test-account", "old-token");
        let now = Instant::now();
        app.codex_quota.account = connected_account(&app.config.paths);
        store_account(&app.config.paths, "test-account", "rotated-token");
        let check = check_codex_quota(&app.config.paths, app.codex_quota.account, false, false);
        app.codex_quota.last_attempt = Some(now);
        let mut state = MissionControlState {
            codex_quota_label: Some("Codex 9/100%(5h)".into()),
            ..Default::default()
        };
        assert!(!app.apply_codex_quota_check(&mut state, check));
        assert_eq!(state.codex_quota_label.as_deref(), Some("Codex 9/100%(5h)"));
        assert!(!app.codex_quota.refresh_due(now, Some(now)));
        let later = now + REFRESH_INTERVAL;
        assert!(!app.codex_quota.refresh_due(later, None));
        assert!(app.codex_quota.refresh_due(later, Some(later)));
        store_account(&app.config.paths, "other-account", "switched-token");
        let check = check_codex_quota(&app.config.paths, app.codex_quota.account, false, false);
        assert!(app.apply_codex_quota_check(&mut state, check));
        assert!(state.codex_quota_label.is_none());
    }

    #[test]
    fn diagnostics_do_not_count_as_agent_activity_but_live_output_does() {
        use crate::output::OutputEvent;
        let mut state = MissionControlState::default();
        let mut drain = super::super::super::DrainResult::default();
        super::super::super::apply_tui_event_to_state(
            &mut state,
            TuiEvent::Output(OutputEvent::Diagnostic {
                level: "info".into(),
                message: "background notice".into(),
            }),
            &mut drain,
        );
        assert!(state.last_agent_activity.is_none());
        super::super::super::apply_tui_event_to_state(
            &mut state,
            TuiEvent::Output(OutputEvent::AssistantDelta {
                text: "working".into(),
            }),
            &mut drain,
        );
        assert!(state.last_agent_activity.is_some());
    }

    #[test]
    fn failed_refresh_hides_previous_quota_and_completes_waiting_modal() {
        let temp = tempfile::TempDir::new().unwrap();
        let (sender, _receiver) = crossbeam_channel::bounded(1);
        let mut app = super::super::tests::test_app(&temp, sender);
        let mut state = MissionControlState {
            codex_quota_label: Some("Codex 9/100%(5h)".into()),
            ..Default::default()
        };
        app.open_usage_modal(&mut state);
        app.codex_quota.account = test_account();
        let now = Instant::now();
        assert!(app.apply_codex_quota_check(
            &mut state,
            QuotaCheck {
                account: test_account(),
                result: Some(UsageLoadResult::Error("offline".into())),
                label: None,
                attempted_at: Some(now),
            }
        ));
        assert!(state.codex_quota_label.is_none());
        assert!(state.pending_usage_request_id.is_none());
        assert_eq!(state.status, "failed to load provider usage");
        assert!(!app.codex_quota.modal_requested);
        assert!(!app.codex_quota.refresh_due(now, Some(now)));
    }

    #[test]
    fn disconnected_auth_hides_quota_without_http_and_manual_usage_explains_login() {
        let temp = tempfile::TempDir::new().unwrap();
        let paths = McPaths::from_root(temp.path().join("mc"));
        let check = check_codex_quota(&paths, test_account(), true, false);
        assert!(check.account.is_none());
        assert!(check.attempted_at.is_none());
        let check = check_codex_quota(&paths, None, true, true);
        assert!(matches!(check.result, Some(UsageLoadResult::Error(_))));
        assert!(check.attempted_at.is_none());
    }

    #[test]
    fn in_flight_usage_is_shared_and_retained_without_wake_delivery() {
        let temp = tempfile::TempDir::new().unwrap();
        let (sender, receiver) = crossbeam_channel::bounded(1);
        drop(receiver);
        let mut app = super::super::tests::test_app(&temp, sender);
        let mut state = MissionControlState::default();
        let now = Instant::now();
        app.codex_quota.last_check = Some(now);
        app.codex_quota.worker = Some(std::thread::spawn(move || QuotaCheck {
            account: test_account(),
            result: Some(UsageLoadResult::Loaded(usage_report(CodexUsage {
                plan: None,
                windows: vec![],
            }))),
            label: Some("Codex 12/100%(5h)".into()),
            attempted_at: Some(now),
        }));
        app.open_usage_modal(&mut state);
        app.open_usage_modal(&mut state);
        while !app.codex_quota.worker.as_ref().unwrap().is_finished() {
            assert!(now.elapsed() < Duration::from_secs(2));
            std::thread::yield_now();
        }
        assert!(app.poll_codex_quota(&mut state, now));
        assert_eq!(
            state.codex_quota_label.as_deref(),
            Some("Codex 12/100%(5h)")
        );
        assert!(state.pending_usage_request_id.is_none());
        assert!(app.codex_quota.worker.is_none());
    }

    #[test]
    fn quota_uses_plan_and_duration_not_position_or_label() {
        let mut usage = CodexUsage {
            plan: Some("pro".into()),
            windows: vec![
                CodexUsageWindow {
                    label: "misleading".into(),
                    limit_window_seconds: Some(604_800),
                    used_percent: 72.0,
                    reset_at: None,
                },
                CodexUsageWindow {
                    label: "Weekly".into(),
                    limit_window_seconds: Some(18_000),
                    used_percent: 13.0,
                    reset_at: None,
                },
            ],
        };
        assert_eq!(
            codex_quota_label(&usage).as_deref(),
            Some("Codex 72/100%(w)")
        );
        for plan in [Some("plus"), Some("free"), Some("unknown"), None] {
            usage.plan = plan.map(str::to_owned);
            assert_eq!(
                codex_quota_label(&usage).as_deref(),
                Some("Codex 13/100%(5h)")
            );
        }
        usage.windows.pop();
        assert_eq!(codex_quota_label(&usage), None);
    }

    #[test]
    fn refresh_requires_recent_agent_activity_and_five_minutes_without_inflight_work() {
        let start = Instant::now();
        let mut refresh = CodexQuotaRefresh {
            account: test_account(),
            last_attempt: Some(start),
            ..Default::default()
        };
        let due = start + REFRESH_INTERVAL;
        assert!(!refresh.refresh_due(start + Duration::from_secs(299), Some(start)));
        assert!(!refresh.refresh_due(due, None));
        assert!(!refresh.refresh_due(due, Some(start)));
        assert!(refresh.refresh_due(due, Some(due)));
        refresh.worker = Some(std::thread::spawn(|| QuotaCheck {
            account: None,
            result: None,
            label: None,
            attempted_at: None,
        }));
        assert!(!refresh.refresh_due(due, Some(due)));
        refresh.worker.take().unwrap().join().unwrap();
        refresh.invalidate();
        assert!(!refresh.refresh_due(due, Some(due)));
    }
}