Skip to main content

ai_usagebar/deepseek/
vendor.rs

1//! DeepSeek renderer — bar text + bordered Pango tooltip.
2
3use std::collections::HashMap;
4
5use chrono::{DateTime, Utc};
6
7use crate::format::{placeholders, substitute, updated_at_hm};
8use crate::pacing::PaceSeverity;
9use crate::pango::{color_span, escape, severity_color};
10use crate::theme::Theme;
11use crate::tooltip::{Line as TooltipLine, render_bordered};
12use crate::usage::DeepseekSnapshot;
13use crate::vendor::{RenderOpts, VendorOutcome};
14use crate::waybar::{Class, WaybarOutput};
15
16use super::fetch::FetchOutcome;
17
18pub const DEFAULT_FORMAT: &str = "{ds_balance}";
19
20pub fn build_placeholders(snap: &DeepseekSnapshot) -> HashMap<&'static str, String> {
21    let avail = if snap.is_available { "up" } else { "down" };
22    placeholders(vec![
23        ("icon", "󰧑".to_string()),
24        ("vendor_short", "dsk".to_string()),
25        // Cross-vendor aliases — no rate-limit windows for DeepSeek.
26        ("session_pct", "0".to_string()),
27        ("session_reset", "—".to_string()),
28        ("weekly_pct", "0".to_string()),
29        ("weekly_reset", "—".to_string()),
30        ("plan", "DeepSeek".to_string()),
31        ("ds_balance", format_money(snap.balance, &snap.currency)),
32        ("ds_granted", format_money(snap.granted, &snap.currency)),
33        ("ds_topped_up", format_money(snap.topped_up, &snap.currency)),
34        ("ds_available", avail.to_string()),
35        ("currency", snap.currency.clone()),
36    ])
37}
38
39fn format_money(v: f64, currency: &str) -> String {
40    match currency {
41        "USD" => format!("${v:.2}"),
42        "CNY" => format!("¥{v:.2}"),
43        _ => format!("{v:.2} {currency}"),
44    }
45}
46
47pub fn severity(snap: &DeepseekSnapshot) -> PaceSeverity {
48    if !snap.is_available {
49        return PaceSeverity::Critical;
50    }
51    // Thresholds scaled by currency. CNY ≈ 7× USD (rough parity).
52    // critical / high / mid boundaries in each currency unit.
53    let (t_critical, t_high, t_mid) = match snap.currency.as_str() {
54        "CNY" => (7.0_f64, 35.0, 140.0),
55        _ => (1.0_f64, 5.0, 20.0), // USD and unknowns treated as USD-scale
56    };
57    if snap.balance < t_critical {
58        PaceSeverity::Critical
59    } else if snap.balance < t_high {
60        PaceSeverity::High
61    } else if snap.balance < t_mid {
62        PaceSeverity::Mid
63    } else {
64        PaceSeverity::Low
65    }
66}
67
68pub fn render(
69    outcome: &VendorOutcome,
70    snap: &DeepseekSnapshot,
71    theme: &Theme,
72    opts: &RenderOpts,
73    now: DateTime<Utc>,
74) -> WaybarOutput {
75    let class = Class::from(severity(snap));
76    let format = opts
77        .format
78        .clone()
79        .unwrap_or_else(|| DEFAULT_FORMAT.to_string());
80    let values = build_placeholders(snap);
81
82    let mut text = substitute(&format, &values);
83    if outcome.stale {
84        text.push_str(" ⏸");
85    }
86
87    let wrapper_color = severity_color(severity(snap), theme).to_string();
88    let icon_prefix = match opts.icon.as_deref() {
89        Some(ic) if !ic.is_empty() => format!("{ic} "),
90        _ => String::new(),
91    };
92    let bar_text = color_span(&wrapper_color, &format!("{icon_prefix}{text}"));
93
94    let tooltip = if let Some(fmt) = opts.tooltip_format.as_deref() {
95        substitute(fmt, &values)
96    } else {
97        render_tooltip(outcome, snap, theme, now)
98    };
99
100    WaybarOutput {
101        text: bar_text,
102        tooltip,
103        class,
104    }
105}
106
107fn render_tooltip(
108    outcome: &VendorOutcome,
109    snap: &DeepseekSnapshot,
110    theme: &Theme,
111    now: DateTime<Utc>,
112) -> String {
113    let blue = &theme.blue;
114    let dim = &theme.dim;
115    let fg = &theme.fg;
116    let color = severity_color(severity(snap), theme);
117
118    let avail_label = if snap.is_available {
119        "API available"
120    } else {
121        "API unavailable"
122    };
123
124    let mut lines: Vec<TooltipLine> = Vec::new();
125    lines.push(TooltipLine::Center(format!(
126        "<span font_weight='bold' foreground='{blue}'>DeepSeek</span>"
127    )));
128    lines.push(TooltipLine::Sep);
129    lines.push(TooltipLine::Body("".into()));
130
131    lines.push(TooltipLine::Body(format!(
132        " <span foreground='{fg}'>  󰢗  Balance</span>"
133    )));
134    lines.push(TooltipLine::Body(format!(
135        "   <span font_weight='bold' foreground='{color}'>{bal}</span>",
136        bal = escape(&format_money(snap.balance, &snap.currency))
137    )));
138    lines.push(TooltipLine::Body(format!(
139        " <span foreground='{dim}'>     granted {granted} · topped-up {topped}</span>",
140        granted = escape(&format_money(snap.granted, &snap.currency)),
141        topped = escape(&format_money(snap.topped_up, &snap.currency))
142    )));
143
144    lines.push(TooltipLine::Body("".into()));
145    lines.push(TooltipLine::Body(format!(
146        " <span foreground='{dim}'>  󰛴  {avail_label}</span>"
147    )));
148
149    if let Some((code, msg)) = outcome.last_error.as_ref()
150        && *code != 0
151    {
152        let (icon, ecolor) = if *code >= 500 {
153            ("󰅚", theme.red.as_str())
154        } else {
155            ("󰀪", theme.orange.as_str())
156        };
157        lines.push(TooltipLine::Body("".into()));
158        lines.push(TooltipLine::Sep);
159        lines.push(TooltipLine::Body(format!(
160            " <span foreground='{ecolor}'>  {icon}  HTTP {code}</span>"
161        )));
162        lines.push(TooltipLine::Body(format!(
163            "     <span foreground='{dim}'>{}</span>",
164            escape(msg)
165        )));
166    }
167
168    let updated = updated_at_hm(now, outcome.cache_age);
169    lines.push(TooltipLine::Body("".into()));
170    lines.push(TooltipLine::Sep);
171    lines.push(TooltipLine::Body(format!(
172        " <span foreground='{dim}'>  󰅐  Updated {updated}</span>"
173    )));
174
175    render_bordered(&lines, theme)
176}
177
178impl From<FetchOutcome> for VendorOutcome {
179    fn from(o: FetchOutcome) -> Self {
180        Self {
181            snapshot: crate::usage::VendorSnapshot::Deepseek(o.snapshot),
182            stale: o.stale,
183            last_error: o.last_error,
184            cache_age: o.cache_age,
185        }
186    }
187}
188
189#[cfg(test)]
190mod tests {
191    use super::*;
192    use crate::usage::DeepseekSnapshot;
193
194    fn sample_snap() -> DeepseekSnapshot {
195        DeepseekSnapshot {
196            is_available: true,
197            balance: 5.50,
198            granted: 5.00,
199            topped_up: 0.50,
200            currency: "USD".into(),
201        }
202    }
203
204    fn sample_outcome(snap: DeepseekSnapshot) -> VendorOutcome {
205        VendorOutcome {
206            snapshot: crate::usage::VendorSnapshot::Deepseek(snap),
207            stale: false,
208            last_error: None,
209            cache_age: Some(std::time::Duration::from_secs(10)),
210        }
211    }
212
213    fn opts() -> RenderOpts {
214        RenderOpts {
215            format: None,
216            tooltip_format: None,
217            icon: None,
218            pace_tolerance: 5,
219            format_pace_color: false,
220            tooltip_pace_pts: false,
221        }
222    }
223
224    #[test]
225    fn default_render_shows_balance() {
226        let snap = sample_snap();
227        let outcome = sample_outcome(snap.clone());
228        let theme = Theme::default();
229        let out = render(&outcome, &snap, &theme, &opts(), Utc::now());
230        assert!(out.text.contains("$5.50"));
231    }
232
233    #[test]
234    fn tooltip_includes_balance_and_availability() {
235        let snap = sample_snap();
236        let outcome = sample_outcome(snap.clone());
237        let theme = Theme::default();
238        let out = render(&outcome, &snap, &theme, &opts(), Utc::now());
239        assert!(out.tooltip.contains("Balance"));
240        assert!(out.tooltip.contains("$5.50"));
241        assert!(out.tooltip.contains("API available"));
242    }
243
244    #[test]
245    fn unavailable_api_shows_critical_severity() {
246        let mut snap = sample_snap();
247        snap.is_available = false;
248        assert_eq!(severity(&snap), PaceSeverity::Critical);
249    }
250
251    #[test]
252    fn stale_appends_pause() {
253        let snap = sample_snap();
254        let mut outcome = sample_outcome(snap.clone());
255        outcome.stale = true;
256        let theme = Theme::default();
257        let out = render(&outcome, &snap, &theme, &opts(), Utc::now());
258        assert!(out.text.contains("⏸"));
259    }
260
261    #[test]
262    fn cny_format() {
263        let snap = DeepseekSnapshot {
264            is_available: true,
265            balance: 20.0,
266            granted: 20.0,
267            topped_up: 0.0,
268            currency: "CNY".into(),
269        };
270        assert_eq!(format_money(snap.balance, &snap.currency), "¥20.00");
271    }
272}