Skip to main content

ai_usagebar/supergrok/
vendor.rs

1//! SuperGrok renderer — current subscription credit % + reset countdown.
2
3use std::collections::HashMap;
4
5use chrono::{DateTime, Utc};
6
7use crate::countdown;
8use crate::format::{placeholders, substitute, updated_at_hm};
9use crate::pacing::PaceSeverity;
10use crate::pango::{self, color_span, escape, severity_color, severity_for};
11use crate::theme::Theme;
12use crate::tooltip::{Line as TooltipLine, render_bordered};
13use crate::usage::SuperGrokSnapshot;
14use crate::vendor::{RenderOpts, VendorOutcome};
15use crate::waybar::{Class, WaybarOutput};
16
17use super::fetch::FetchOutcome;
18
19pub const DEFAULT_FORMAT: &str = "{sgk_pct}% · {sgk_reset}";
20
21const DEFAULT_ICON: &str = "󰚩";
22
23pub fn build_placeholders(
24    snap: &SuperGrokSnapshot,
25    now: DateTime<Utc>,
26) -> HashMap<&'static str, String> {
27    let pct = snap.weekly_pct.to_string();
28    let reset = countdown::format(snap.reset_at, now);
29    let prepaid = snap
30        .prepaid_balance
31        .map(|b| {
32            if b < 0.0 {
33                format!("-${:.2}", -b)
34            } else {
35                format!("${b:.2}")
36            }
37        })
38        .unwrap_or_else(|| "—".into());
39
40    placeholders(vec![
41        ("icon", DEFAULT_ICON.to_string()),
42        ("vendor_short", "sgk".to_string()),
43        // Cross-vendor compatibility aliases for the single current pool.
44        ("plan", snap.plan.clone()),
45        ("session_pct", pct.clone()),
46        ("session_reset", reset.clone()),
47        ("weekly_pct", pct.clone()),
48        ("weekly_reset", reset.clone()),
49        // SuperGrok-specific.
50        ("sgk_plan", snap.plan.clone()),
51        ("sgk_pct", pct),
52        ("sgk_reset", reset),
53        ("sgk_period", snap.period.label().to_string()),
54        ("sgk_prepaid", prepaid),
55    ])
56}
57
58pub fn severity(snap: &SuperGrokSnapshot) -> PaceSeverity {
59    severity_for(snap.weekly_pct)
60}
61
62pub fn render(
63    outcome: &VendorOutcome,
64    snap: &SuperGrokSnapshot,
65    theme: &Theme,
66    opts: &RenderOpts,
67    now: DateTime<Utc>,
68) -> WaybarOutput {
69    let class = Class::from(severity(snap));
70    let format = opts
71        .format
72        .clone()
73        .unwrap_or_else(|| DEFAULT_FORMAT.to_string());
74    let mut values = build_placeholders(snap, now);
75    for key in ["plan", "sgk_plan"] {
76        if let Some(value) = values.get_mut(key) {
77            *value = escape(value);
78        }
79    }
80
81    let mut text = substitute(&format, &values);
82    if outcome.stale {
83        text.push_str(" ⏸");
84    }
85
86    let wrapper_color = severity_color(severity(snap), theme).to_string();
87    let icon_prefix = match opts.icon.as_deref() {
88        Some(ic) if !ic.is_empty() => format!("{ic} "),
89        _ => String::new(),
90    };
91    let bar_text = color_span(&wrapper_color, &format!("{icon_prefix}{text}"));
92
93    let tooltip = if let Some(fmt) = opts.tooltip_format.as_deref() {
94        substitute(fmt, &values)
95    } else {
96        render_tooltip(outcome, snap, theme, now)
97    };
98
99    WaybarOutput {
100        text: bar_text,
101        tooltip,
102        class,
103    }
104}
105
106/// One usage-percent row in the same shape as `tooltip::push_window`:
107/// label, progress bar + bold %, then optional dim reset line.
108fn push_pct_row(
109    lines: &mut Vec<TooltipLine>,
110    theme: &Theme,
111    label: &str,
112    pct: i32,
113    reset_at: Option<DateTime<Utc>>,
114    now: DateTime<Utc>,
115) {
116    let fg = &theme.fg;
117    let dim = &theme.dim;
118    let color = severity_color(severity_for(pct), theme);
119    let bar = pango::progress_bar(pct, color, theme, None);
120    lines.push(TooltipLine::Body(format!(
121        " <span foreground='{fg}'>{label}</span>"
122    )));
123    lines.push(TooltipLine::Body(format!(
124        "   {bar}  <span font_weight='bold' foreground='{color}'>{pct}%</span>"
125    )));
126    if reset_at.is_some() {
127        lines.push(TooltipLine::Body(format!(
128            " <span foreground='{dim}'>  ⏱  Resets in {}</span>",
129            escape(&countdown::format(reset_at, now))
130        )));
131    }
132}
133
134fn render_tooltip(
135    outcome: &VendorOutcome,
136    snap: &SuperGrokSnapshot,
137    theme: &Theme,
138    now: DateTime<Utc>,
139) -> String {
140    let blue = &theme.blue;
141    let dim = &theme.dim;
142
143    let mut lines: Vec<TooltipLine> = Vec::new();
144    lines.push(TooltipLine::Center(format!(
145        "<span font_weight='bold' foreground='{blue}'>{}</span>",
146        escape(&snap.plan)
147    )));
148    lines.push(TooltipLine::Sep);
149    lines.push(TooltipLine::Body("".into()));
150
151    let period_label = format!("  󰔟  {} Build credits", snap.period.label());
152    push_pct_row(
153        &mut lines,
154        theme,
155        &period_label,
156        snap.weekly_pct,
157        snap.reset_at,
158        now,
159    );
160
161    if let Some(bal) = snap.prepaid_balance {
162        let bal_s = if bal < 0.0 {
163            format!("-${:.2}", -bal)
164        } else {
165            format!("${bal:.2}")
166        };
167        lines.push(TooltipLine::Body("".into()));
168        lines.push(TooltipLine::Body(format!(
169            " <span foreground='{dim}'>  󰢗  Prepaid API  {}</span>",
170            escape(&bal_s)
171        )));
172    }
173
174    if let Some((code, msg)) = outcome.last_error.as_ref() {
175        let (icon, ecolor) = if *code >= 500 {
176            ("󰅚", theme.red.as_str())
177        } else {
178            ("󰀪", theme.orange.as_str())
179        };
180        let label = if *code == 0 {
181            "Refresh error".to_string()
182        } else {
183            format!("HTTP {code}")
184        };
185        lines.push(TooltipLine::Body("".into()));
186        lines.push(TooltipLine::Sep);
187        lines.push(TooltipLine::Body(format!(
188            " <span foreground='{ecolor}'>  {icon}  {label}</span>"
189        )));
190        lines.push(TooltipLine::Body(format!(
191            "     <span foreground='{dim}'>{}</span>",
192            escape(msg)
193        )));
194    }
195
196    let updated = updated_at_hm(now, outcome.cache_age);
197    lines.push(TooltipLine::Body("".into()));
198    lines.push(TooltipLine::Sep);
199    lines.push(TooltipLine::Body(format!(
200        " <span foreground='{dim}'>  󰅐  Updated {updated}</span>"
201    )));
202
203    render_bordered(&lines, theme)
204}
205
206impl From<FetchOutcome> for VendorOutcome {
207    fn from(o: FetchOutcome) -> Self {
208        Self {
209            snapshot: crate::usage::VendorSnapshot::SuperGrok(o.snapshot),
210            stale: o.stale,
211            last_error: o.last_error,
212            cache_age: o.cache_age,
213        }
214    }
215}
216
217#[cfg(test)]
218mod tests {
219    use super::*;
220    use crate::usage::SuperGrokPeriod;
221    use chrono::TimeZone;
222
223    fn now() -> DateTime<Utc> {
224        Utc.with_ymd_and_hms(2026, 8, 5, 12, 0, 0).unwrap()
225    }
226
227    fn sample_snap() -> SuperGrokSnapshot {
228        SuperGrokSnapshot {
229            plan: "SuperGrok".into(),
230            account: "user-1".into(),
231            weekly_pct: 34,
232            period: SuperGrokPeriod::Weekly,
233            reset_at: Some(now() + chrono::Duration::hours(20)),
234            prepaid_balance: Some(0.0),
235        }
236    }
237
238    fn sample_outcome(snap: SuperGrokSnapshot) -> VendorOutcome {
239        VendorOutcome {
240            snapshot: crate::usage::VendorSnapshot::SuperGrok(snap),
241            stale: false,
242            last_error: None,
243            cache_age: Some(std::time::Duration::from_secs(10)),
244        }
245    }
246
247    fn opts() -> RenderOpts {
248        RenderOpts {
249            format: None,
250            tooltip_format: None,
251            icon: None,
252            pace_tolerance: 5,
253            format_pace_color: false,
254            tooltip_pace_pts: false,
255        }
256    }
257
258    #[test]
259    fn renders_weekly_pct_and_reset() {
260        let snap = sample_snap();
261        let o = sample_outcome(snap.clone());
262        let out = render(&o, &snap, &Theme::default(), &opts(), now());
263        assert!(out.text.contains("34%"));
264        assert!(out.tooltip.contains("Build credits"));
265        assert!(out.tooltip.contains("Weekly"));
266        assert!(out.tooltip.contains("SuperGrok"));
267        // Usage-% vendors (Anthropic / OpenAI) draw a filled progress bar in
268        // the tooltip; SuperGrok must match that shape rather than bare %.
269        assert!(
270            out.tooltip.contains('█') || out.tooltip.contains('░'),
271            "tooltip missing progress bar cells: {}",
272            out.tooltip
273        );
274        assert!(out.tooltip.contains("Resets in"));
275    }
276
277    #[test]
278    fn high_usage_is_critical() {
279        let mut snap = sample_snap();
280        snap.weekly_pct = 95;
281        assert_eq!(severity(&snap), PaceSeverity::Critical);
282    }
283
284    #[test]
285    fn placeholders_include_generic_aliases() {
286        let snap = sample_snap();
287        let ph = build_placeholders(&snap, now());
288        assert_eq!(ph.get("vendor_short").map(String::as_str), Some("sgk"));
289        assert_eq!(ph.get("weekly_pct").map(String::as_str), Some("34"));
290        assert_eq!(ph.get("session_pct").map(String::as_str), Some("34"));
291        assert_eq!(ph.get("sgk_period").map(String::as_str), Some("Weekly"));
292        assert_eq!(ph.get("sgk_prepaid").map(String::as_str), Some("$0.00"));
293    }
294}