use crate::error::DatalabError;
use colored::Colorize;
use serde::Serialize;
use std::io::{IsTerminal, Write};
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Instant;
static PROGRESS_ENABLED: AtomicBool = AtomicBool::new(true);
#[derive(Debug, Clone, Serialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ProgressEvent {
Start {
operation: String,
#[serde(skip_serializing_if = "Option::is_none")]
file: Option<String>,
},
Upload {
bytes_sent: u64,
total_bytes: u64,
},
Submit {
request_id: String,
},
Poll {
status: String,
elapsed_secs: f64,
},
CacheHit {
cache_key: String,
},
Complete {
elapsed_secs: f64,
},
Error {
code: String,
message: String,
},
}
#[derive(Clone)]
pub struct Progress {
enabled: bool,
start_time: Instant,
}
impl Progress {
pub fn new(quiet: bool, verbose: bool) -> Self {
let is_tty = std::io::stderr().is_terminal();
let enabled = if quiet {
false
} else if verbose {
true
} else {
is_tty
};
PROGRESS_ENABLED.store(enabled, Ordering::SeqCst);
Self {
enabled,
start_time: Instant::now(),
}
}
#[allow(dead_code)]
pub fn is_enabled(&self) -> bool {
self.enabled
}
pub fn emit(&self, event: ProgressEvent) {
if !self.enabled {
return;
}
if let Ok(json) = serde_json::to_string(&event) {
let _ = writeln!(std::io::stderr(), "{}", json);
}
}
pub fn start(&self, operation: &str, file: Option<&str>) {
self.emit(ProgressEvent::Start {
operation: operation.to_string(),
file: file.map(String::from),
});
}
pub fn upload(&self, bytes_sent: u64, total_bytes: u64) {
self.emit(ProgressEvent::Upload {
bytes_sent,
total_bytes,
});
}
pub fn submit(&self, request_id: &str) {
self.emit(ProgressEvent::Submit {
request_id: request_id.to_string(),
});
}
pub fn poll(&self, status: &str) {
self.emit(ProgressEvent::Poll {
status: status.to_string(),
elapsed_secs: self.start_time.elapsed().as_secs_f64(),
});
}
pub fn cache_hit(&self, cache_key: &str) {
self.emit(ProgressEvent::CacheHit {
cache_key: cache_key.to_string(),
});
}
pub fn complete(&self) {
self.emit(ProgressEvent::Complete {
elapsed_secs: self.start_time.elapsed().as_secs_f64(),
});
}
pub fn error(&self, err: &DatalabError) {
self.emit(ProgressEvent::Error {
code: err.code().to_string(),
message: err.to_string(),
});
}
}
pub struct Output {
is_tty: bool,
use_color: bool,
}
impl Output {
pub fn new() -> Self {
let is_tty = std::io::stderr().is_terminal();
let use_color = is_tty && std::env::var("NO_COLOR").is_err();
Self { is_tty, use_color }
}
pub fn error(&self, err: &DatalabError) {
if self.is_tty {
self.print_colored_error(err);
} else {
eprintln!("{}", err.to_json());
}
}
fn print_colored_error(&self, err: &DatalabError) {
if self.use_color {
eprint!("{}: ", "error".red().bold());
} else {
eprint!("error: ");
}
eprintln!("{}", err);
if let Some(suggestion) = err.suggestion() {
eprintln!();
if self.use_color {
eprintln!("{}: {}", "hint".yellow().bold(), suggestion);
} else {
eprintln!("hint: {}", suggestion);
}
}
if let Some(help_url) = err.help_url() {
if self.use_color {
eprintln!("{}: {}", "help".cyan().bold(), help_url);
} else {
eprintln!("help: {}", help_url);
}
}
}
#[allow(dead_code)]
pub fn info(&self, message: &str) {
if self.is_tty {
if self.use_color {
eprintln!("{}: {}", "info".blue().bold(), message);
} else {
eprintln!("info: {}", message);
}
}
}
#[allow(dead_code)]
pub fn warn(&self, message: &str) {
if self.is_tty {
if self.use_color {
eprintln!("{}: {}", "warning".yellow().bold(), message);
} else {
eprintln!("warning: {}", message);
}
}
}
#[allow(dead_code)]
pub fn success(&self, message: &str) {
if self.is_tty {
if self.use_color {
eprintln!("{}: {}", "success".green().bold(), message);
} else {
eprintln!("success: {}", message);
}
}
}
}
impl Default for Output {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_progress_event_serialization() {
let event = ProgressEvent::Start {
operation: "convert".to_string(),
file: Some("test.pdf".to_string()),
};
let json = serde_json::to_string(&event).unwrap();
assert!(json.contains("\"type\":\"start\""));
assert!(json.contains("\"operation\":\"convert\""));
}
#[test]
fn test_progress_poll_event() {
let event = ProgressEvent::Poll {
status: "processing".to_string(),
elapsed_secs: 1.5,
};
let json = serde_json::to_string(&event).unwrap();
assert!(json.contains("\"type\":\"poll\""));
assert!(json.contains("\"status\":\"processing\""));
}
}