1use std::collections::HashMap;
5
6use chrono::{DateTime, Utc};
7
8use crate::format::{money, placeholders, substitute, updated_at_hm};
9use crate::pacing::PaceSeverity;
10use crate::pango::{color_span, escape, severity_color};
11use crate::theme::Theme;
12use crate::tooltip::{Line as TooltipLine, render_bordered};
13use crate::usage::MoonshotSnapshot;
14use crate::vendor::{RenderOpts, VendorId, VendorOutcome};
15use crate::waybar::{Class, WaybarOutput};
16
17use super::fetch::FetchOutcome;
18
19pub const DEFAULT_FORMAT: &str = "{km_balance}";
20
21pub fn build_placeholders(snap: &MoonshotSnapshot) -> HashMap<&'static str, String> {
22 placeholders(vec![
23 ("icon", "".to_string()),
24 ("vendor_short", VendorId::Moonshot.short_name().to_string()),
25 ("session_pct", "0".to_string()),
27 ("session_reset", "—".to_string()),
28 ("weekly_pct", "0".to_string()),
29 ("weekly_reset", "—".to_string()),
30 ("plan", "Kimi".to_string()),
31 ("km_balance", money(snap.available, &snap.currency)),
32 ("km_voucher", money(snap.voucher, &snap.currency)),
33 ("km_cash", money(snap.cash, &snap.currency)),
34 ("currency", snap.currency.clone()),
35 ])
36}
37
38pub fn severity(snap: &MoonshotSnapshot) -> PaceSeverity {
42 if snap.available <= 0.0 {
43 return PaceSeverity::Critical;
44 }
45 let (t_critical, t_high, t_mid) = match snap.currency.as_str() {
46 "CNY" => (7.0_f64, 35.0, 140.0),
47 _ => (1.0_f64, 5.0, 20.0),
48 };
49 if snap.available < t_critical {
50 PaceSeverity::Critical
51 } else if snap.available < t_high {
52 PaceSeverity::High
53 } else if snap.available < t_mid {
54 PaceSeverity::Mid
55 } else {
56 PaceSeverity::Low
57 }
58}
59
60pub fn render(
61 outcome: &VendorOutcome,
62 snap: &MoonshotSnapshot,
63 theme: &Theme,
64 opts: &RenderOpts,
65 now: DateTime<Utc>,
66) -> WaybarOutput {
67 let class = Class::from(severity(snap));
68 let format = opts
69 .format
70 .clone()
71 .unwrap_or_else(|| DEFAULT_FORMAT.to_string());
72 let values = build_placeholders(snap);
73
74 let mut text = substitute(&format, &values);
75 if outcome.stale {
76 text.push_str(" ⏸");
77 }
78
79 let wrapper_color = severity_color(severity(snap), theme).to_string();
80 let icon_prefix = match opts.icon.as_deref() {
81 Some(ic) if !ic.is_empty() => format!("{ic} "),
82 _ => String::new(),
83 };
84 let bar_text = color_span(&wrapper_color, &format!("{icon_prefix}{text}"));
85
86 let tooltip = if let Some(fmt) = opts.tooltip_format.as_deref() {
87 substitute(fmt, &values)
88 } else {
89 render_tooltip(outcome, snap, theme, now)
90 };
91
92 WaybarOutput {
93 text: bar_text,
94 tooltip,
95 class,
96 }
97}
98
99fn render_tooltip(
100 outcome: &VendorOutcome,
101 snap: &MoonshotSnapshot,
102 theme: &Theme,
103 now: DateTime<Utc>,
104) -> String {
105 let blue = &theme.blue;
106 let dim = &theme.dim;
107 let fg = &theme.fg;
108 let color = severity_color(severity(snap), theme);
109
110 let mut lines: Vec<TooltipLine> = Vec::new();
111 lines.push(TooltipLine::Center(format!(
112 "<span font_weight='bold' foreground='{blue}'>Kimi (Moonshot)</span>"
113 )));
114 lines.push(TooltipLine::Sep);
115 lines.push(TooltipLine::Body("".into()));
116
117 lines.push(TooltipLine::Body(format!(
118 " <span foreground='{fg}'> Balance</span>"
119 )));
120 lines.push(TooltipLine::Body(format!(
121 " <span font_weight='bold' foreground='{color}'>{bal}</span>",
122 bal = escape(&money(snap.available, &snap.currency))
123 )));
124 lines.push(TooltipLine::Body(format!(
125 " <span foreground='{dim}'> cash {cash} · voucher {voucher}</span>",
126 cash = escape(&money(snap.cash, &snap.currency)),
127 voucher = escape(&money(snap.voucher, &snap.currency))
128 )));
129
130 if snap.available <= 0.0 {
131 lines.push(TooltipLine::Body("".into()));
132 lines.push(TooltipLine::Body(format!(
133 " <span foreground='{}'> out of credit — inference blocked</span>",
134 theme.orange.as_str()
135 )));
136 }
137
138 if let Some((code, msg)) = outcome.last_error.as_ref()
139 && *code != 0
140 {
141 let (icon, ecolor) = if *code >= 500 {
142 ("", theme.red.as_str())
143 } else {
144 ("", theme.orange.as_str())
145 };
146 lines.push(TooltipLine::Body("".into()));
147 lines.push(TooltipLine::Sep);
148 lines.push(TooltipLine::Body(format!(
149 " <span foreground='{ecolor}'> {icon} HTTP {code}</span>"
150 )));
151 lines.push(TooltipLine::Body(format!(
152 " <span foreground='{dim}'>{}</span>",
153 escape(msg)
154 )));
155 }
156
157 let updated = updated_at_hm(now, outcome.cache_age);
158 lines.push(TooltipLine::Body("".into()));
159 lines.push(TooltipLine::Sep);
160 lines.push(TooltipLine::Body(format!(
161 " <span foreground='{dim}'> Updated {updated}</span>"
162 )));
163
164 render_bordered(&lines, theme)
165}
166
167impl From<FetchOutcome> for VendorOutcome {
168 fn from(o: FetchOutcome) -> Self {
169 Self {
170 snapshot: crate::usage::VendorSnapshot::Moonshot(o.snapshot),
171 stale: o.stale,
172 last_error: o.last_error,
173 cache_age: o.cache_age,
174 }
175 }
176}
177
178#[cfg(test)]
179mod tests {
180 use super::*;
181 use crate::usage::MoonshotSnapshot;
182
183 fn sample_snap() -> MoonshotSnapshot {
184 MoonshotSnapshot {
185 available: 49.58,
186 voucher: 46.58,
187 cash: 3.0,
188 currency: "USD".into(),
189 }
190 }
191
192 fn sample_outcome(snap: MoonshotSnapshot) -> VendorOutcome {
193 VendorOutcome {
194 snapshot: crate::usage::VendorSnapshot::Moonshot(snap),
195 stale: false,
196 last_error: None,
197 cache_age: Some(std::time::Duration::from_secs(10)),
198 }
199 }
200
201 fn opts() -> RenderOpts {
202 RenderOpts {
203 format: None,
204 tooltip_format: None,
205 icon: None,
206 pace_tolerance: 5,
207 format_pace_color: false,
208 tooltip_pace_pts: false,
209 }
210 }
211
212 #[test]
213 fn default_render_shows_balance() {
214 let snap = sample_snap();
215 let out = render(
216 &sample_outcome(snap.clone()),
217 &snap,
218 &Theme::default(),
219 &opts(),
220 Utc::now(),
221 );
222 assert!(out.text.contains("$49.58"));
223 }
224
225 #[test]
229 fn a_cash_debt_carries_its_sign_ahead_of_the_symbol() {
230 let mut snap = sample_snap();
231 snap.cash = -5.71;
232 let out = render(
233 &sample_outcome(snap.clone()),
234 &snap,
235 &Theme::default(),
236 &opts(),
237 Utc::now(),
238 );
239 assert!(out.tooltip.contains("cash -$5.71"), "{}", out.tooltip);
240 assert!(!out.tooltip.contains("$-5.71"), "{}", out.tooltip);
241
242 snap.currency = "CNY".into();
243 let out = render(
244 &sample_outcome(snap.clone()),
245 &snap,
246 &Theme::default(),
247 &opts(),
248 Utc::now(),
249 );
250 assert!(out.tooltip.contains("cash -¥5.71"), "{}", out.tooltip);
251 }
252
253 #[test]
254 fn tooltip_includes_cash_and_voucher() {
255 let snap = sample_snap();
256 let out = render(
257 &sample_outcome(snap.clone()),
258 &snap,
259 &Theme::default(),
260 &opts(),
261 Utc::now(),
262 );
263 assert!(out.tooltip.contains("cash $3.00"));
264 assert!(out.tooltip.contains("voucher $46.58"));
265 }
266
267 #[test]
268 fn cny_uses_yuan_symbol() {
269 let mut snap = sample_snap();
270 snap.currency = "CNY".into();
271 let out = render(
272 &sample_outcome(snap.clone()),
273 &snap,
274 &Theme::default(),
275 &opts(),
276 Utc::now(),
277 );
278 assert!(out.text.contains("¥49.58"));
279 }
280
281 #[test]
282 fn zero_balance_is_critical() {
283 let mut snap = sample_snap();
284 snap.available = 0.0;
285 assert_eq!(severity(&snap), PaceSeverity::Critical);
286 }
287}