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