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