1use std::io::{stdout, Error, Write};
6use std::sync::atomic::{AtomicBool, Ordering};
7use std::sync::mpsc::{Sender, TryRecvError};
8use std::sync::Arc;
9
10use log::error;
11
12use irox_time::epoch::UnixTimestamp;
13use irox_time::format::iso8601::ISO8601Duration;
14use irox_time::format::Format;
15use irox_time::Duration;
16
17use crate::{get_human, ProgressPrinter, Task};
18
19pub struct ConsoleProgressBar {
20 width: usize,
21}
22
23impl ConsoleProgressBar {
24 pub fn new(width: usize) -> Self {
25 ConsoleProgressBar { width }
26 }
27
28 pub fn print_infinite_progress(&self, task: &Task) -> Result<(), Error> {
29 let current = task.current_progress_count();
30 if let Some(started) = task.get_started() {
31 let current = current as f64;
32 let elapsed = started.elapsed().as_seconds_f64();
33 let avg_per_sec = current / elapsed;
34 let (avg_per_sec, avg_unit) = get_human!(avg_per_sec);
35 let (current, unit) = get_human!(current);
36
37 let time = (elapsed * 2.) as u64 & 0x7;
38 let state = match time {
39 0 => "\u{25b6}\u{25b9}\u{25b9}\u{25b9}\u{25b9}",
40 1 => "\u{25b9}\u{25b6}\u{25b9}\u{25b9}\u{25b9}",
41 2 => "\u{25b9}\u{25b9}\u{25b6}\u{25b9}\u{25b9}",
42 3 => "\u{25b9}\u{25b9}\u{25b9}\u{25b6}\u{25b9}",
43 4 => "\u{25b9}\u{25b9}\u{25b9}\u{25b9}\u{25b6}",
44 _ => "\u{25b9}\u{25b9}\u{25b9}\u{25b9}\u{25b9}",
45 };
46 let state2 = match time & 0x3 {
47 0 => "\u{25dc}",
48 1 => "\u{25dd}",
49 2 => "\u{25de}",
50 _ => "\u{25df}",
51 };
52
53 let out = format!("| ({current:.02}{unit}) {avg_per_sec:.02}{avg_unit}/s\r");
54 let spaces =
55 " ".repeat(((self.width as i32 - out.len() as i32 - 9).max(1) / 4) as usize);
56 let status = task.current_status().unwrap_or_default();
57 let out = format!("| ({current:.02}{unit}) {spaces}{state2}{spaces} {state} {spaces}{state2}{spaces}{avg_per_sec:.02}{avg_unit}/s {status}\r");
58
59 let mut stdio = stdout();
60 stdio.write_all(out.as_bytes())?;
61 return stdio.flush();
62 }
63 Ok(())
64 }
65
66 pub fn print_progress(&self, task: &Task) -> Result<(), Error> {
67 let pct = task.current_progress_frac();
68 let current = task.current_progress_count();
69 let max = task.max_elements();
70 if max == u64::MAX {
71 return self.print_infinite_progress(task);
72 }
73
74 let rem_str = ISO8601Duration.format(&task.get_remaining_time());
75
76 let w_pct = self.width as f64 * pct;
77 let whole = w_pct.floor() as usize;
78 let part = (8.0 * w_pct.fract()).round() as u8;
79 let mut char = match part {
80 0 => "\u{00A0}",
81 1 => "\u{258F}",
82 2 => "\u{258E}",
83 3 => "\u{258D}",
84 4 => "\u{258C}",
85 5 => "\u{258B}",
86 6 => "\u{258A}",
87 7 => "\u{2589}",
88 _ => "\u{2588}",
89 };
90
91 let mut rem = self.width - whole;
92 if rem == 1 {
94 rem = 0;
95 } else if rem == 0 {
96 char = "";
97 } else {
98 rem -= 1;
99 }
100 let whole = "\u{2588}".repeat(whole);
101 let rem = " ".repeat(rem);
102 let status = task.current_status().unwrap_or_default();
103 let out = format!(
104 "{:>3.0}%|{whole}{char}{rem}| ({current}/{max}) {rem_str} {status}\r",
105 pct * 100.
106 );
107 let mut stdio = stdout();
108 stdio.write_all(out.as_bytes())?;
109 stdio.flush()
110 }
111}
112
113pub struct ConsoleProgressPrinter {
114 thread_handle: Option<std::thread::JoinHandle<()>>,
115 running_flag: Arc<AtomicBool>,
116 sender: Sender<Task>,
117}
118
119impl ConsoleProgressPrinter {
120 pub fn new_update_rate(update_rate: Duration) -> ConsoleProgressPrinter {
121 let running_flag = Arc::new(AtomicBool::new(true));
122 let running = running_flag.clone();
123
124 let (sender, receiver) = std::sync::mpsc::channel();
125
126 #[allow(unused_assignments)]
127 let thread_handle = std::thread::spawn(move || {
128 let mut tasks: Vec<Task> = Vec::new();
129 let mut last_run = UnixTimestamp::now();
130 while running.load(Ordering::Relaxed) {
131 last_run = UnixTimestamp::now();
132 let next_run = last_run + update_rate;
133 let _r = stdout().write_all(&[0x1B, b'[', b'2', b'K']);
134 for task in &tasks {
135 let _res = ConsoleProgressBar::new(60).print_progress(task);
136 }
137
138 match receiver.try_recv() {
139 Ok(task) => tasks.push(task),
140 Err(e) => {
141 if e == TryRecvError::Disconnected {
142 return;
143 }
144 }
145 };
146
147 let delay = next_run - UnixTimestamp::now();
148 std::thread::sleep(delay.into());
149 }
150 let _r = stdout().write_all(&[0x1B, b'[', b'2', b'K']);
151 for task in &tasks {
152 let _res = ConsoleProgressBar::new(40).print_progress(task);
153 }
154 });
155
156 ConsoleProgressPrinter {
157 thread_handle: Some(thread_handle),
158 running_flag,
159 sender,
160 }
161 }
162}
163
164impl Drop for ConsoleProgressPrinter {
165 fn drop(&mut self) {
166 self.running_flag.store(false, Ordering::Relaxed);
167 if let Some(handle) = self.thread_handle.take() {
168 let _res = handle.join();
169 }
170 }
171}
172
173impl ProgressPrinter for ConsoleProgressPrinter {
174 fn track_task_progress(&self, task: &Task) {
175 if let Err(e) = self.sender.send(task.clone()) {
176 error!("Error sending task to printer: {e:?}");
177 }
178 }
179}