1use std::collections::HashMap;
5
6use chrono::{DateTime, Utc};
7
8use crate::countdown;
9use crate::format::{placeholders, substitute, updated_at_hm};
10use crate::pacing::PaceSeverity;
11use crate::pango::{color_span, escape, severity_color, severity_for};
12use crate::theme::Theme;
13use crate::tooltip::{Line as TooltipLine, WindowRow, push_window_with_row, render_bordered};
14use crate::usage::{GrokbotSnapshot, UsageWindow};
15use crate::vendor::{RenderOpts, VendorId, VendorOutcome};
16use crate::waybar::{Class, WaybarOutput};
17
18use super::fetch::FetchOutcome;
19
20pub const DEFAULT_FORMAT: &str = "{gbt_weekly_pct}%";
21
22pub fn build_placeholders(
23 snap: &GrokbotSnapshot,
24 now: DateTime<Utc>,
25) -> HashMap<&'static str, String> {
26 let reset = countdown::format(snap.reset_at, now);
27 let allowance = |value: String| {
30 if snap.has_included_allowance {
31 value
32 } else {
33 String::new()
34 }
35 };
36 placeholders(vec![
37 ("icon", VendorId::Grokbot.bar_icon().to_string()),
38 ("vendor_short", VendorId::Grokbot.short_name().to_string()),
39 ("plan", snap.plan.clone()),
41 ("weekly_pct", allowance(snap.weekly_pct.to_string())),
42 ("weekly_reset", allowance(reset.clone())),
43 ("gbt_plan", snap.plan.clone()),
45 ("gbt_weekly_pct", allowance(snap.weekly_pct.to_string())),
46 ("gbt_weekly_reset", allowance(reset)),
47 (
48 "gbt_on_demand",
49 if snap.on_demand_enabled { "on" } else { "off" }.to_string(),
50 ),
51 ])
52}
53
54pub fn severity(snap: &GrokbotSnapshot) -> PaceSeverity {
56 if snap.has_included_allowance {
57 severity_for(snap.weekly_pct)
58 } else {
59 PaceSeverity::Low
60 }
61}
62
63pub fn render(
64 outcome: &VendorOutcome,
65 snap: &GrokbotSnapshot,
66 theme: &Theme,
67 opts: &RenderOpts,
68 now: DateTime<Utc>,
69) -> WaybarOutput {
70 let class = Class::from(severity(snap));
71 let format = opts
72 .format
73 .clone()
74 .unwrap_or_else(|| DEFAULT_FORMAT.to_string());
75 let values = build_placeholders(snap, now);
76 let mut pango_values = values.clone();
79 for key in ["plan", "gbt_plan"] {
80 if let Some(value) = pango_values.get_mut(key) {
81 *value = escape(value);
82 }
83 }
84
85 let mut text = substitute(&format, &pango_values);
86 if outcome.stale {
87 text.push_str(" ⏸");
88 }
89
90 let wrapper_color = severity_color(severity(snap), theme).to_string();
91 let icon_prefix = match opts.icon.as_deref() {
92 Some(ic) if !ic.is_empty() => format!("{ic} "),
93 _ => String::new(),
94 };
95 let bar_text = color_span(&wrapper_color, &format!("{icon_prefix}{text}"));
96
97 let tooltip = if let Some(fmt) = opts.tooltip_format.as_deref() {
98 substitute(fmt, &pango_values)
99 } else {
100 render_tooltip(outcome, snap, theme, now)
101 };
102
103 WaybarOutput {
104 text: bar_text,
105 tooltip,
106 class,
107 }
108}
109
110fn render_tooltip(
111 outcome: &VendorOutcome,
112 snap: &GrokbotSnapshot,
113 theme: &Theme,
114 now: DateTime<Utc>,
115) -> String {
116 let blue = &theme.blue;
117 let dim = &theme.dim;
118 let fg = &theme.fg;
119 let plan_color = severity_color(severity(snap), theme);
120
121 let mut lines: Vec<TooltipLine> = Vec::new();
122 lines.push(TooltipLine::Center(format!(
123 "<span font_weight='bold' foreground='{blue}'>Grok Bot</span>"
124 )));
125 lines.push(TooltipLine::Sep);
126 lines.push(TooltipLine::Body("".into()));
127
128 lines.push(TooltipLine::Body(format!(
129 " <span foreground='{fg}'> Plan</span>"
130 )));
131 lines.push(TooltipLine::Body(format!(
132 " <span font_weight='bold' foreground='{plan_color}'>{}</span>",
133 escape(&snap.plan)
134 )));
135
136 lines.push(TooltipLine::Body("".into()));
137 if snap.has_included_allowance {
138 push_window_with_row(
141 &mut lines,
142 " Weekly included usage",
143 &UsageWindow {
144 utilization_pct: snap.weekly_pct,
145 resets_at: snap.reset_at,
146 window_duration: snap.window.unwrap_or_else(chrono::Duration::zero),
147 },
148 theme,
149 now,
150 WindowRow::default(),
151 );
152 if let Some(note) = snap.on_demand_note() {
153 lines.push(TooltipLine::Body(format!(
154 " <span foreground='{dim}'>{note}</span>"
155 )));
156 }
157 } else {
158 lines.push(TooltipLine::Body(format!(
159 " <span foreground='{dim}'> No included allowance on this account</span>"
160 )));
161 }
162
163 if let Some((code, msg)) = outcome.last_error.as_ref() {
164 let (label, icon, ecolor) = match *code {
166 0 => ("Grok Bot error".to_string(), "", theme.red.as_str()),
167 code if code >= 500 => (format!("HTTP {code}"), "", theme.red.as_str()),
168 code => (format!("HTTP {code}"), "", theme.orange.as_str()),
169 };
170 lines.push(TooltipLine::Body("".into()));
171 lines.push(TooltipLine::Sep);
172 lines.push(TooltipLine::Body(format!(
173 " <span foreground='{ecolor}'> {icon} {label}</span>"
174 )));
175 if msg != &label {
176 lines.push(TooltipLine::Body(format!(
177 " <span foreground='{dim}'>{}</span>",
178 escape(msg)
179 )));
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 o.map(crate::usage::VendorSnapshot::Grokbot)
196 }
197}
198
199#[cfg(test)]
200mod tests {
201 use super::*;
202 use chrono::TimeZone;
203
204 fn now() -> DateTime<Utc> {
205 Utc.with_ymd_and_hms(2026, 9, 14, 12, 0, 0).unwrap()
206 }
207
208 fn sample_snap() -> GrokbotSnapshot {
209 GrokbotSnapshot {
210 plan: "Grok Bot Plan".into(),
211 has_included_allowance: true,
212 weekly_pct: 42,
213 has_available_usage: true,
214 on_demand_enabled: false,
215 period_start: Some(now() - chrono::Duration::days(3)),
216 reset_at: Some(now() + chrono::Duration::days(4)),
217 window: Some(chrono::Duration::days(7)),
218 }
219 }
220
221 fn sample_outcome(snap: GrokbotSnapshot) -> VendorOutcome {
222 VendorOutcome {
223 snapshot: crate::usage::VendorSnapshot::Grokbot(snap),
224 stale: false,
225 last_error: None,
226 cache_age: Some(std::time::Duration::from_secs(10)),
227 }
228 }
229
230 fn opts() -> RenderOpts {
231 RenderOpts {
232 format: None,
233 tooltip_format: None,
234 icon: None,
235 pace_tolerance: 5,
236 format_pace_color: false,
237 tooltip_pace_pts: false,
238 }
239 }
240
241 #[test]
242 fn default_render_shows_the_weekly_percent() {
243 let snap = sample_snap();
244 let outcome = sample_outcome(snap.clone());
245 let out = render(&outcome, &snap, &Theme::default(), &opts(), now());
246 assert!(out.text.contains("42%"), "text: {}", out.text);
247 assert!(!out.text.contains("%%"), "text: {}", out.text);
248 }
249
250 #[test]
251 fn placeholder_set_contains_all_keys() {
252 let snap = sample_snap();
253 let values = build_placeholders(&snap, now());
254 for key in [
255 "gbt_plan",
256 "gbt_weekly_pct",
257 "gbt_weekly_reset",
258 "gbt_on_demand",
259 "plan",
260 "weekly_pct",
261 "weekly_reset",
262 "vendor_short",
263 ] {
264 assert!(values.contains_key(key), "missing placeholder {key}");
265 }
266 assert_eq!(values["gbt_weekly_pct"], "42");
267 assert_eq!(values["weekly_pct"], "42");
268 assert_eq!(values["gbt_on_demand"], "off");
269 assert!(!values["gbt_weekly_reset"].is_empty());
270 }
271
272 #[test]
273 fn no_allowance_renders_empty_weekly_placeholders_not_a_zero() {
274 let snap = GrokbotSnapshot {
275 has_included_allowance: false,
276 weekly_pct: 0,
277 ..sample_snap()
278 };
279 let values = build_placeholders(&snap, now());
280 for key in [
281 "gbt_weekly_pct",
282 "gbt_weekly_reset",
283 "weekly_pct",
284 "weekly_reset",
285 ] {
286 assert_eq!(values[key], "", "{key} must render empty");
287 }
288 let outcome = sample_outcome(snap.clone());
292 let out = render(&outcome, &snap, &Theme::default(), &opts(), now());
293 assert!(!out.text.contains('{'), "{}", out.text);
294 assert!(!out.text.contains('}'), "{}", out.text);
295 assert!(out.text.contains('%'), "{}", out.text);
296 assert!(
297 !out.text.contains('0'),
298 "a fabricated 0 must not appear: {}",
299 out.text
300 );
301 assert_eq!(severity(&snap), PaceSeverity::Low);
302 }
303
304 #[test]
305 fn tooltip_draws_the_weekly_bar_and_a_reset_countdown() {
306 let snap = sample_snap();
307 let outcome = sample_outcome(snap.clone());
308 let out = render(&outcome, &snap, &Theme::default(), &opts(), now());
309 assert!(out.tooltip.contains("Grok Bot"), "{}", out.tooltip);
310 assert!(out.tooltip.contains("Grok Bot Plan"), "{}", out.tooltip);
311 assert!(
312 out.tooltip.contains("Weekly included usage"),
313 "{}",
314 out.tooltip
315 );
316 assert!(out.tooltip.contains("42%"), "{}", out.tooltip);
317 assert!(out.tooltip.contains("Resets in"), "{}", out.tooltip);
318 assert!(!out.tooltip.contains("2026-09-18"), "{}", out.tooltip);
320 }
321
322 #[test]
323 fn the_no_allowance_state_is_a_text_line_not_a_meter() {
324 let snap = GrokbotSnapshot {
325 has_included_allowance: false,
326 weekly_pct: 0,
327 ..sample_snap()
328 };
329 let outcome = sample_outcome(snap.clone());
330 let out = render(&outcome, &snap, &Theme::default(), &opts(), now());
331 assert!(
332 out.tooltip.contains("No included allowance"),
333 "{}",
334 out.tooltip
335 );
336 assert!(
337 !out.tooltip.contains("Weekly included usage"),
338 "{}",
339 out.tooltip
340 );
341 }
342
343 #[test]
344 fn the_on_demand_footnote_appears_only_at_an_exhausted_pool_with_on_demand() {
345 let mut snap = sample_snap();
346 snap.weekly_pct = 100;
347 snap.has_available_usage = true;
348 snap.on_demand_enabled = true;
349 let outcome = sample_outcome(snap.clone());
350 let out = render(&outcome, &snap, &Theme::default(), &opts(), now());
351 assert!(out.tooltip.contains("on-demand"), "{}", out.tooltip);
352
353 snap.on_demand_enabled = false;
354 let outcome = sample_outcome(snap.clone());
355 let out = render(&outcome, &snap, &Theme::default(), &opts(), now());
356 assert!(!out.tooltip.contains("on-demand may"), "{}", out.tooltip);
357 }
358
359 #[test]
360 fn plan_is_pango_escaped() {
361 let mut snap = sample_snap();
362 snap.plan = "A&B <beta>".into();
363 let outcome = sample_outcome(snap.clone());
364 let out = render(&outcome, &snap, &Theme::default(), &opts(), now());
365 assert!(
366 out.tooltip.contains("A&B <beta>"),
367 "tooltip: {}",
368 out.tooltip
369 );
370 let mut o = opts();
371 o.tooltip_format = Some("{gbt_plan}".into());
372 let out = render(&outcome, &snap, &Theme::default(), &o, now());
373 assert_eq!(out.tooltip, "A&B <beta>");
374 }
375
376 #[test]
377 fn stale_appends_pause() {
378 let snap = sample_snap();
379 let mut outcome = sample_outcome(snap.clone());
380 outcome.stale = true;
381 let out = render(&outcome, &snap, &Theme::default(), &opts(), now());
382 assert!(out.text.contains("⏸"));
383 }
384
385 #[test]
386 fn fetch_outcome_conversion_preserves_metadata() {
387 let snap = sample_snap();
388 let fetch = FetchOutcome {
389 snapshot: snap,
390 stale: true,
391 last_error: Some((401, "bad".into())),
392 cache_age: Some(std::time::Duration::from_secs(42)),
393 };
394 let vendor: VendorOutcome = fetch.into();
395 assert!(matches!(
396 vendor.snapshot,
397 crate::usage::VendorSnapshot::Grokbot(_)
398 ));
399 assert!(vendor.stale);
400 assert_eq!(vendor.last_error, Some((401, "bad".into())));
401 assert_eq!(vendor.cache_age, Some(std::time::Duration::from_secs(42)));
402 }
403}