magi-code 0.97.3

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

use super::MissionControlApp;
use crate::{
    providers::{
        claude_subscription::refresh_cli_subscription_status,
        claude_usage::{ClaudeUsage, connected_account, load_claude_usage, quota_label},
    },
    tui::{TuiEvent, state::MissionControlState},
};

const REFRESH_INTERVAL: Duration = Duration::from_secs(300);
const FAILED_REFRESH_RETRY_INTERVAL: Duration = Duration::from_secs(60);
const CONNECTION_CHECK_INTERVAL: Duration = Duration::from_secs(5);
/// Consecutive failed `claude auth status` probes before treating the account as signed out.
const DISCONNECT_AFTER_MISSED_CHECKS: u8 = 3;

#[derive(Default)]
pub(super) struct ClaudeQuotaRefresh {
    worker: Option<JoinHandle<QuotaCheck>>,
    account: Option<[u8; 32]>,
    missed_connection_checks: u8,
    last_attempt: Option<Instant>,
    last_load_failed: bool,
    usage: Option<ClaudeUsage>,
    last_label_update: Option<Instant>,
    last_check: Option<Instant>,
    manual_refresh: bool,
}

struct QuotaCheck {
    account: Option<[u8; 32]>,
    usage: Option<ClaudeUsage>,
    attempted_at: Option<Instant>,
}

impl ClaudeQuotaRefresh {
    pub(super) fn request_refresh(&mut self) {
        self.manual_refresh = true;
    }

    pub(super) fn cleanup(&mut self) -> Option<String> {
        let worker = self.worker.take()?;
        match crate::thread_join::join_with_timeout(worker, Duration::from_millis(50)) {
            Ok(Err(_)) => Some("Claude usage worker panicked".into()),
            _ => None,
        }
    }
}

impl MissionControlApp {
    pub(super) fn poll_claude_quota(
        &mut self,
        state: &mut MissionControlState,
        now: Instant,
    ) -> bool {
        let quota = &mut self.claude_quota;
        let mut changed = false;
        if quota
            .worker
            .as_ref()
            .is_some_and(|worker| worker.is_finished())
        {
            let check = quota
                .worker
                .take()
                .expect("finished worker")
                .join()
                .unwrap_or(QuotaCheck {
                    account: None,
                    usage: None,
                    attempted_at: Some(now),
                });
            let old_label = state.claude_quota_label.clone();
            // A single failed probe is often transient; keep the shown quota until it persists.
            if check.account.is_none() {
                quota.missed_connection_checks = quota.missed_connection_checks.saturating_add(1);
            } else {
                quota.missed_connection_checks = 0;
            }
            let account_changed = quota.account != check.account
                && (check.account.is_some()
                    || quota.missed_connection_checks >= DISCONNECT_AFTER_MISSED_CHECKS);
            if account_changed {
                quota.account = check.account;
                quota.last_attempt = None;
                quota.last_load_failed = false;
                quota.usage = None;
            }
            if let Some(attempted_at) = check.attempted_at {
                quota.last_attempt = Some(attempted_at);
                quota.last_load_failed = check.usage.is_none();
                // Keep last known quota on failed loads; retry sooner instead of blanking the label.
                if check.usage.is_some() {
                    quota.usage = check.usage;
                }
            }
            state.claude_quota_label = quota
                .usage
                .as_ref()
                .and_then(|usage| quota_label(usage, chrono::Utc::now()));
            quota.last_label_update = Some(now);
            changed = old_label != state.claude_quota_label;
        }
        if let Some(usage) = quota.usage.as_ref()
            && quota
                .last_label_update
                .is_none_or(|last| now.duration_since(last) >= Duration::from_secs(1))
        {
            quota.last_label_update = Some(now);
            let label = quota_label(usage, chrono::Utc::now());
            if label != state.claude_quota_label {
                state.claude_quota_label = label;
                changed = true;
            }
        }
        let refresh_interval = if quota.last_load_failed {
            FAILED_REFRESH_RETRY_INTERVAL
        } else {
            REFRESH_INTERVAL
        };
        let recently_active = state
            .last_agent_activity
            .is_some_and(|last| now.duration_since(last) < REFRESH_INTERVAL);
        // Without any quota shown, keep retrying even while idle.
        let refresh_due = quota.account.is_some()
            && quota
                .last_attempt
                .is_none_or(|last| now.duration_since(last) >= refresh_interval)
            && (quota.usage.is_none() || recently_active);
        let check_due = quota
            .last_check
            .is_none_or(|last| now.duration_since(last) >= CONNECTION_CHECK_INTERVAL);
        if quota.worker.is_none() && (check_due || refresh_due || quota.manual_refresh) {
            let previous = quota.account;
            let force = quota.manual_refresh || refresh_due;
            quota.manual_refresh = false;
            quota.last_check = Some(now);
            let sender = self.events.clone();
            quota.worker = Some(std::thread::spawn(move || {
                // Periodic checks must observe sign-outs and account switches, so they
                // bypass the shared status cache; usage reuses this same probe.
                let status = refresh_cli_subscription_status();
                let account = status.as_ref().and_then(connected_account);
                let mut check = QuotaCheck {
                    account,
                    usage: None,
                    attempted_at: None,
                };
                if let Some(status) = status.as_ref().filter(|_| account.is_some())
                    && (force || previous != account)
                {
                    check.attempted_at = Some(Instant::now());
                    let usage = load_claude_usage(status);
                    // Drop usage only when a different account signed in mid-request;
                    // a transiently failed re-probe keeps the loaded quota.
                    match refresh_cli_subscription_status()
                        .as_ref()
                        .and_then(connected_account)
                    {
                        Some(after) if Some(after) != account => check.account = Some(after),
                        _ => check.usage = usage.ok(),
                    }
                }
                let _ = sender.try_send(TuiEvent::WorkerOutcomeReady);
                check
            }));
        }
        changed
    }
}