1use colored::Colorize;
2use indicatif::{ProgressBar, ProgressStyle};
3use tabled::settings::Style;
4use tabled::{Table, Tabled};
5
6use crate::client::api::CcuInfo;
7use crate::server::storage::StoredServer;
8
9#[derive(Clone, Copy)]
10pub struct Ui {
11 pub quiet: bool,
12}
13
14impl Ui {
15 pub fn new(quiet: bool) -> Self {
16 Self { quiet }
17 }
18
19 pub fn spinner(&self, msg: &str) -> Option<ProgressBar> {
20 if self.quiet {
21 return None;
22 }
23 let pb = ProgressBar::new_spinner();
24 pb.set_style(
25 ProgressStyle::with_template("{spinner:.cyan} {msg}")
26 .unwrap()
27 .tick_strings(&[
28 "\u{280b}", "\u{2819}", "\u{2839}", "\u{2838}", "\u{283c}", "\u{2834}",
29 "\u{2826}", "\u{2827}", "\u{2807}", "\u{280f}",
30 ]),
31 );
32 pb.set_message(msg.to_string());
33 pb.enable_steady_tick(std::time::Duration::from_millis(80));
34 Some(pb)
35 }
36
37 pub fn spinner_done(pb: Option<ProgressBar>, msg: &str) {
38 if let Some(pb) = pb {
39 pb.finish_with_message(msg.to_string());
40 }
41 }
42
43 pub fn spinner_fail(pb: Option<ProgressBar>, msg: &str) {
44 if let Some(pb) = pb {
45 pb.finish_with_message(format!("\u{2717} {msg}"));
46 }
47 }
48
49 pub fn success(&self, msg: &str) {
50 if self.quiet {
51 println!("{msg}");
52 } else {
53 println!("{} {msg}", "\u{2713}".green().bold());
54 }
55 }
56
57 pub fn info(&self, msg: &str) {
58 if self.quiet {
59 println!("{msg}");
60 } else {
61 println!("{} {msg}", "\u{00b7}".dimmed());
62 }
63 }
64
65 pub fn warn(&self, msg: &str) {
66 if self.quiet {
67 eprintln!("warning: {msg}");
68 } else {
69 eprintln!("{} {msg}", "\u{26a0}".yellow().bold());
70 }
71 }
72
73 pub fn error(&self, msg: &str) {
74 if self.quiet {
75 eprintln!("error: {msg}");
76 } else {
77 eprintln!("{} {msg}", "error:".red().bold());
78 }
79 }
80
81 pub fn print_auth_status(&self, email: &str, name: &str) {
82 if self.quiet {
83 println!("signed_in\t{email}\t{name}");
84 } else {
85 println!(
86 "{} Signed in as {} ({})",
87 "\u{2713}".green().bold(),
88 name.bold(),
89 email.dimmed()
90 );
91 }
92 }
93
94 pub fn print_auth_not_signed_in(&self) {
95 if self.quiet {
96 println!("not_signed_in");
97 } else {
98 println!(
99 "{} Not signed in. Run {} to authenticate.",
100 "\u{00b7}".dimmed(),
101 "colab-cli auth login".cyan().bold()
102 );
103 }
104 }
105
106 pub fn print_server_list(&self, servers: &[StoredServer]) {
107 if servers.is_empty() {
108 if self.quiet {
109 println!("no_servers");
110 } else {
111 println!(
112 "{} No servers assigned. Run {} to assign one.",
113 "\u{00b7}".dimmed(),
114 "colab-cli server assign".cyan().bold()
115 );
116 }
117 return;
118 }
119
120 if self.quiet {
121 println!("name\tvariant\taccelerator\tshape\tendpoint\ttoken_expires_at\tassigned_at");
122 for s in servers {
123 println!(
124 "{}\t{}\t{}\t{}\t{}\t{}\t{}",
125 s.label,
126 s.variant,
127 s.accelerator.as_deref().unwrap_or(""),
128 s.shape,
129 s.endpoint,
130 s.token_expires_at
131 .with_timezone(&chrono::Local)
132 .format("%Y-%m-%dT%H:%M:%S%:z"),
133 s.date_assigned
134 .with_timezone(&chrono::Local)
135 .format("%Y-%m-%dT%H:%M:%S%:z"),
136 );
137 }
138 return;
139 }
140
141 #[derive(Tabled)]
142 struct Row {
143 #[tabled(rename = "Name")]
144 name: String,
145 #[tabled(rename = "Type")]
146 variant: String,
147 #[tabled(rename = "Accelerator")]
148 accelerator: String,
149 #[tabled(rename = "Shape")]
150 shape: String,
151 #[tabled(rename = "Endpoint")]
152 endpoint: String,
153 #[tabled(rename = "Token Expires")]
154 token_expires: String,
155 #[tabled(rename = "Assigned")]
156 assigned: String,
157 }
158
159 let now = chrono::Utc::now();
160 let rows: Vec<Row> = servers
161 .iter()
162 .map(|s| {
163 let remaining = s.token_expires_at - now;
164 let expires = if remaining.num_minutes() < 10 {
165 format!("{}m \u{26a0}", remaining.num_minutes())
166 } else {
167 format!("{}m", remaining.num_minutes())
168 };
169 Row {
170 name: s.label.clone(),
171 variant: s.variant.display_name().to_string(),
172 accelerator: s.accelerator.clone().unwrap_or_else(|| "\u{2014}".into()),
173 shape: s.shape.display_name().to_string(),
174 endpoint: truncate(&s.endpoint, 32),
175 token_expires: expires,
176 assigned: s
177 .date_assigned
178 .with_timezone(&chrono::Local)
179 .format("%b %d %H:%M")
180 .to_string(),
181 }
182 })
183 .collect();
184
185 let mut table = Table::new(rows);
186 table.with(Style::sharp());
187 println!("{table}");
188 }
189
190 pub fn print_server_status(&self, s: &StoredServer) {
191 let now = chrono::Utc::now();
192 let remaining = s.token_expires_at - now;
193
194 if self.quiet {
195 println!("name\t{}", s.label);
196 println!("variant\t{}", s.variant);
197 println!("accelerator\t{}", s.accelerator.as_deref().unwrap_or(""));
198 println!("shape\t{}", s.shape);
199 println!("endpoint\t{}", s.endpoint);
200 println!("token_expires_in_minutes\t{}", remaining.num_minutes());
201 println!(
202 "assigned_at\t{}",
203 s.date_assigned
204 .with_timezone(&chrono::Local)
205 .format("%Y-%m-%dT%H:%M:%S%:z")
206 );
207 return;
208 }
209
210 let expiry_str = if remaining.num_minutes() < 10 {
211 format!("{}m {}", remaining.num_minutes(), "(refresh soon)".yellow())
212 } else {
213 format!("{}m", remaining.num_minutes())
214 };
215
216 println!("{}", "Server".bold());
217 kv("Name", &s.label.bold().to_string());
218 kv("Type", &s.variant.display_name().cyan().to_string());
219 kv(
220 "Accelerator",
221 s.accelerator.as_deref().unwrap_or("\u{2014}"),
222 );
223 kv("Shape", s.shape.display_name());
224 kv("Endpoint", &s.endpoint.dimmed().to_string());
225 kv("Token expires", &expiry_str);
226 kv(
227 "Assigned",
228 &s.date_assigned
229 .with_timezone(&chrono::Local)
230 .format("%Y-%m-%d %H:%M")
231 .to_string(),
232 );
233 }
234
235 pub fn print_usage(&self, info: &CcuInfo) {
236 if self.quiet {
237 println!("balance\t{:.2}", info.current_balance);
238 println!("rate_hourly\t{:.2}", info.consumption_rate_hourly);
239 println!("assignments\t{}", info.assignments_count);
240 println!("gpus\t{}", info.eligible_gpus.join(","));
241 println!("tpus\t{}", info.eligible_tpus.join(","));
242 return;
243 }
244
245 println!("{}", "Compute Usage".bold());
246 kv("Balance", &format!("{:.2} CCU", info.current_balance));
247 kv(
248 "Burn rate",
249 &format!("{:.2} CCU/hr", info.consumption_rate_hourly),
250 );
251 kv("Active servers", &info.assignments_count.to_string());
252 }
253
254 pub fn print_accelerators(&self, info: &CcuInfo) {
255 if self.quiet {
256 println!("balance\t{:.2}", info.current_balance);
257 println!("rate_hourly\t{:.2}", info.consumption_rate_hourly);
258 println!("type\tmodel\trate_ccu_per_hour\tstatus");
259 println!("CPU\t-\t{:.2}\tonline", ccu_rate("CPU", ""));
260 for gpu in &info.eligible_gpus {
261 println!("GPU\t{gpu}\t{:.2}\tonline", ccu_rate("GPU", gpu));
262 }
263 for tpu in &info.eligible_tpus {
264 println!("TPU\t{tpu}\t{:.2}\tonline", ccu_rate("TPU", tpu));
265 }
266 return;
267 }
268
269 #[derive(Tabled)]
270 struct AccelRow {
271 #[tabled(rename = "Type")]
272 kind: String,
273 #[tabled(rename = "Model")]
274 model: String,
275 #[tabled(rename = "~CCU/hr")]
276 rate: String,
277 #[tabled(rename = "Status")]
278 status: String,
279 }
280
281 let mut rows = vec![AccelRow {
282 kind: "CPU".to_string(),
283 model: "\u{2014}".to_string(),
284 rate: format!("{:.2}", ccu_rate("CPU", "")),
285 status: "online".green().to_string(),
286 }];
287 for gpu in &info.eligible_gpus {
288 rows.push(AccelRow {
289 kind: "GPU".to_string(),
290 model: gpu.clone(),
291 rate: format!("{:.2}", ccu_rate("GPU", gpu)),
292 status: "online".green().to_string(),
293 });
294 }
295 for tpu in &info.eligible_tpus {
296 rows.push(AccelRow {
297 kind: "TPU".to_string(),
298 model: tpu.clone(),
299 rate: format!("{:.2}", ccu_rate("TPU", tpu)),
300 status: "online".green().to_string(),
301 });
302 }
303
304 let mut table = Table::new(rows);
305 table.with(Style::sharp());
306 println!("{table}");
307
308 println!();
309 println!("{}", "Compute Units".bold());
310 kv("Balance", &format!("{:.2} CCU", info.current_balance));
311 kv(
312 "Usage rate",
313 &format!(
314 "~{:.2} CCU/hr ({} active server{})",
315 info.consumption_rate_hourly,
316 info.assignments_count,
317 if info.assignments_count == 1 { "" } else { "s" }
318 ),
319 );
320 println!();
321 println!(
322 "{}",
323 "Rates are approximate and may vary by region / availability.".dimmed()
324 );
325 }
326
327 pub fn print_system_info(&self, server_name: &str, raw: &str) {
328 let mut sections = std::collections::HashMap::new();
329 let mut current_key: Option<String> = None;
330 let mut buf = String::new();
331 for line in raw.lines() {
332 if let Some(tag) = line.strip_prefix("<<<").and_then(|s| s.strip_suffix(">>>")) {
333 if let Some(k) = current_key.take() {
334 sections.insert(k, buf.trim().to_string());
335 }
336 current_key = Some(tag.to_string());
337 buf.clear();
338 } else {
339 buf.push_str(line);
340 buf.push('\n');
341 }
342 }
343 if let Some(k) = current_key.take() {
344 sections.insert(k, buf.trim().to_string());
345 }
346
347 let get = |k: &str| sections.get(k).cloned().unwrap_or_default();
348
349 if self.quiet {
350 println!("server\t{server_name}");
351 println!("kernel\t{}", get("UNAME"));
352 println!("cpu\t{}", get("CPU").replace('\n', " "));
353 println!("mem\t{}", get("MEM"));
354 println!("disk\t{}", get("DISK"));
355 println!("gpu\t{}", get("GPU").replace('\n', " | "));
356 println!("uptime\t{}", get("UPTIME"));
357 return;
358 }
359
360 println!(
361 "{} {}",
362 "System".bold(),
363 format!("({server_name})").dimmed()
364 );
365 println!();
366
367 if !get("UNAME").is_empty() {
368 kv("Kernel", &get("UNAME"));
369 }
370
371 let cpu = get("CPU");
372 let cpu_lines: Vec<&str> = cpu.lines().collect();
373 if cpu_lines.len() >= 2 {
374 kv(
375 "CPU",
376 &format!("{} ({} cores)", cpu_lines[1].trim(), cpu_lines[0].trim()),
377 );
378 } else if !cpu_lines.is_empty() {
379 kv("CPU", cpu_lines[0]);
380 }
381
382 let mem = get("MEM");
383 let mem_parts: Vec<&str> = mem.split('\t').collect();
384 if mem_parts.len() >= 3 {
385 let used = mem_parts[1].trim();
386 let total = mem_parts[0].trim();
387 kv("Memory", &format!("{used} / {total} used"));
388 }
389
390 let disk = get("DISK");
391 let disk_parts: Vec<&str> = disk.split('\t').collect();
392 if disk_parts.len() >= 4 {
393 kv(
394 "Disk (/)",
395 &format!(
396 "{} / {} used ({})",
397 disk_parts[1].trim(),
398 disk_parts[0].trim(),
399 disk_parts[3].trim()
400 ),
401 );
402 }
403
404 let gpu = get("GPU");
405 if gpu.trim() != "none" && !gpu.is_empty() {
406 for line in gpu.lines() {
407 kv("GPU", line.trim());
408 }
409 } else {
410 kv("GPU", &"none".dimmed().to_string());
411 }
412
413 if !get("UPTIME").is_empty() {
414 kv("Uptime", &get("UPTIME"));
415 }
416 }
417}
418
419fn kv(key: &str, value: &str) {
420 println!(" {:<18} {value}", format!("{key}:").dimmed());
421}
422
423fn truncate(s: &str, max: usize) -> String {
424 if s.len() <= max {
425 s.to_string()
426 } else {
427 format!("{}\u{2026}", &s[..max - 1])
428 }
429}
430
431pub fn ccu_rate(kind: &str, model: &str) -> f64 {
435 match (kind, model) {
436 ("CPU", _) => 0.08,
437 ("GPU", "T4") => 1.76,
438 ("GPU", "L4") => 4.82,
439 ("GPU", "V100") => 4.91,
440 ("GPU", "A100") => 11.77,
441 ("GPU", "H100") => 14.43,
442 ("GPU", _) => 2.00,
443 ("TPU", "v2-8") => 1.96,
444 ("TPU", "v5e-1") => 2.20,
445 ("TPU", _) => 2.00,
446 _ => 0.0,
447 }
448}
449
450#[cfg(test)]
451mod tests {
452 use super::*;
453
454 #[test]
455 fn truncate_short_string_unchanged() {
456 assert_eq!(truncate("abc", 10), "abc");
457 }
458
459 #[test]
460 fn truncate_long_string_ellipsized() {
461 let t = truncate("abcdefghij", 5);
462 assert!(t.ends_with('\u{2026}'));
463 assert_eq!(t.chars().count(), 5);
464 }
465
466 #[test]
467 fn ccu_rate_known_gpus() {
468 assert!((ccu_rate("GPU", "T4") - 1.76).abs() < 1e-9);
469 assert!((ccu_rate("GPU", "A100") - 11.77).abs() < 1e-9);
470 assert!((ccu_rate("GPU", "L4") - 4.82).abs() < 1e-9);
471 }
472
473 #[test]
474 fn ccu_rate_unknown_gpu_has_fallback() {
475 assert_eq!(ccu_rate("GPU", "Quantum9000"), 2.00);
476 }
477
478 #[test]
479 fn ccu_rate_cpu() {
480 assert!(ccu_rate("CPU", "") > 0.0);
481 }
482}