1use std::{collections::HashMap, time::Instant, sync::{Arc, Mutex, atomic::{AtomicBool, Ordering}}};
2use ratatui::{
3 backend::CrosstermBackend,
4 layout::{Constraint, Direction, Layout},
5 style::{Color, Modifier, Style},
6 widgets::{Block, Borders, Paragraph, Table, Row, TableState},
7 Terminal,
8};
9use crossterm::{
10 execute, cursor, terminal, ExecutableCommand,
11 terminal::{disable_raw_mode, enable_raw_mode, EnterAlternateScreen, LeaveAlternateScreen},
12 event::{self, Event, KeyCode, KeyEventKind, KeyModifiers},
13};
14use std::io;
15use crate::core::types::FailedOperationInfo;
16use webbrowser;
17
18#[derive(Clone)]
19pub enum OperationStatus {
20 Preprocessing,
21 Ongoing(f64),
22 Completed,
23 Failed(FailedOperationInfo),
24}
25
26pub struct UIState {
27 pub total_files: usize,
28 pub processed_files: usize,
29 pub processed_bytes: u64,
30 pub total_bytes: u64,
31 pub start_time: Instant,
32 pub all_operations: HashMap<String, OperationStatus>,
33 pub album_id: Option<String>,
34 pub file_sizes: HashMap<String, u64>,
35 pub completed_urls: HashMap<String, String>,
36}
37
38impl UIState {
39 pub fn new(total_files: usize, album_id: Option<String>, total_bytes: u64) -> Self {
40 Self {
41 total_files,
42 processed_files: 0,
43 processed_bytes: 0,
44 total_bytes,
45 start_time: Instant::now(),
46 all_operations: HashMap::new(),
47 album_id,
48 file_sizes: HashMap::new(),
49 completed_urls: HashMap::new(),
50 }
51 }
52
53 pub fn add_current_operation(&mut self, name: String, progress: f64, size: u64) {
54 self.all_operations.insert(name.clone(), OperationStatus::Ongoing(progress));
55 self.file_sizes.insert(name, size);
56 }
57
58 pub fn update_progress(&mut self, name: &str, progress: f64) {
59 if let Some(OperationStatus::Ongoing(ref mut p)) = self.all_operations.get_mut(name) {
60 *p = progress;
61 }
62 }
63
64 pub fn remove_current_operation(&mut self, name: &str, url: Option<&str>) {
65 self.all_operations.insert(name.to_string(), OperationStatus::Completed);
66 self.processed_files += 1;
67 if let Some(url) = url {
68 self.completed_urls.insert(name.to_string(), url.to_string());
69 }
70 }
71
72 pub fn add_processed_bytes(&mut self, bytes: u64) {
73 self.processed_bytes += bytes;
74 }
75
76 pub fn add_failed_operation(&mut self, name: String, info: FailedOperationInfo) {
77 self.all_operations.insert(name, OperationStatus::Failed(info));
78 }
79
80 pub fn add_preprocessing(&mut self, name: String, size: u64) {
81 self.file_sizes.insert(name.clone(), size);
82 self.all_operations.insert(name, OperationStatus::Preprocessing);
83 }
84
85 pub fn remove_operation(&mut self, name: &str) {
86 self.all_operations.remove(name);
87 self.file_sizes.remove(name);
88 }
89
90 pub fn add_to_total_files(&mut self, count: usize) {
91 self.total_files += count;
92 }
93}
94
95fn format_size(size: u64) -> String {
96 if size >= 1024 * 1024 * 1024 {
97 format!("{:.1} GB", size as f64 / (1024.0 * 1024.0 * 1024.0))
98 } else if size >= 1024 * 1024 {
99 format!("{:.1} MB", size as f64 / (1024.0 * 1024.0))
100 } else if size >= 1024 {
101 format!("{:.1} KB", size as f64 / 1024.0)
102 } else {
103 format!("{} B", size)
104 }
105}
106
107pub struct UI {
108 terminal: Terminal<CrosstermBackend<io::Stdout>>,
109 table_state: TableState,
110 previous_row_count: usize,
111}
112
113impl UI {
114 pub fn new() -> Result<Self, Box<dyn std::error::Error>> {
115 enable_raw_mode()?;
116 let mut stdout = io::stdout();
117 execute!(stdout, EnterAlternateScreen)?;
118 let backend = CrosstermBackend::new(stdout);
119 let terminal = Terminal::new(backend)?;
120 Ok(Self { terminal, table_state: TableState::default(), previous_row_count: 0 })
121 }
122
123 pub fn draw(&mut self, state: &UIState) -> Result<(), Box<dyn std::error::Error>> {
124 self.terminal.draw(|f| {
125 let size = f.area();
126 let elapsed = state.start_time.elapsed().as_secs_f64();
127 let bytes_per_sec = if elapsed > 0.0 { state.processed_bytes as f64 / elapsed } else { 0.0 };
128 let speed_mb_s = bytes_per_sec / 1_000_000.0;
129
130 let remaining_bytes: u64 = state.total_bytes.saturating_sub(state.processed_bytes);
131
132 let eta_str = if remaining_bytes > 0 && bytes_per_sec > 0.0 {
133 let time_left_seconds = remaining_bytes as f64 / bytes_per_sec;
134 if time_left_seconds < 60.0 {
135 format!(" | ETA: {:.0}s", time_left_seconds)
136 } else if time_left_seconds < 3600.0 {
137 format!(" | ETA: {:.0}m", time_left_seconds / 60.0)
138 } else {
139 format!(" | ETA: {:.1}h", time_left_seconds / 3600.0)
140 }
141 } else {
142 String::new()
143 };
144
145 let header_text = if let Some(album) = &state.album_id {
146 format!("Bunkr Client | Album: {} | Processed: {}/{} | Speed: {:.2} MB/s{}", album, state.processed_files, state.total_files, speed_mb_s, eta_str)
147 } else {
148 format!("Bunkr Client | Processed: {}/{} | Speed: {:.2} MB/s{}", state.processed_files, state.total_files, speed_mb_s, eta_str)
149 };
150 let header = Paragraph::new(header_text)
151 .block(Block::default().borders(Borders::ALL).title("Header"))
152 .style(Style::default().fg(Color::Cyan).add_modifier(Modifier::BOLD));
153
154 let header_height = 3;
155 let chunks = Layout::default()
156 .direction(Direction::Vertical)
157 .constraints([Constraint::Length(header_height), Constraint::Min(0)])
158 .split(size);
159
160 f.render_widget(header, chunks[0]);
161
162 let list_area = chunks[1];
163
164 let mut all_items_vec: Vec<(&String, &OperationStatus)> = state.all_operations.iter().collect();
165 all_items_vec.sort_by(|a, b| a.0.cmp(b.0));
166
167 let current_row_count = all_items_vec.len();
168
169 let rows: Vec<Row> = all_items_vec.iter().map(|(name, status)| {
170 let file_name = std::path::Path::new(name).file_name().unwrap_or(std::ffi::OsStr::new(name)).to_string_lossy();
171 let size = match status {
172 OperationStatus::Failed(info) => info.file_size,
173 _ => *state.file_sizes.get(*name).unwrap_or(&0),
174 };
175 let size_str = format_size(size);
176 let (progress_str, status_str, url_str) = match status {
177 OperationStatus::Preprocessing => ("".to_string(), "Preprocessing".to_string(), "".to_string()),
178 OperationStatus::Ongoing(progress) => (format!("{:.0}%", progress * 100.0), "Ongoing".to_string(), "".to_string()),
179 OperationStatus::Completed => {
180 let url = state.completed_urls.get(*name).cloned().unwrap_or_else(|| "".to_string());
181 ("100%".to_string(), "Completed".to_string(), url)
182 }
183 OperationStatus::Failed(info) => {
184 let status_str_inner = if let Some(code) = info.status_code {
185 format!(" (HTTP {})", code)
186 } else {
187 String::new()
188 };
189 ("".to_string(), format!("Failed{}: {}", status_str_inner, info.error), "".to_string())
190 }
191 };
192 Row::new(vec![file_name.to_string(), size_str, progress_str, status_str, url_str])
193 }).collect();
194
195 let widths = [
196 Constraint::Percentage(25),
197 Constraint::Percentage(12),
198 Constraint::Percentage(12),
199 Constraint::Percentage(25),
200 Constraint::Percentage(26),
201 ];
202
203 let table = Table::new(rows, widths)
204 .block(Block::default().borders(Borders::ALL).title("Operations"))
205 .header(
206 Row::new(vec!["File", "Size", "Progress", "Status", "URL"])
207 .style(Style::default().fg(Color::Cyan).add_modifier(Modifier::BOLD))
208 )
209 .row_highlight_style(Style::default().add_modifier(Modifier::REVERSED));
210
211 if let Some(selected) = self.table_state.selected() {
212 if selected == self.previous_row_count.saturating_sub(1) && current_row_count > self.previous_row_count {
213 self.table_state.select(Some(current_row_count - 1));
214 }
215 }
216
217 self.previous_row_count = current_row_count;
218
219 f.render_stateful_widget(table, list_area, &mut self.table_state);
220 })?;
221 Ok(())
222 }
223
224 pub fn restore(&mut self) -> Result<(), Box<dyn std::error::Error>> {
225 disable_raw_mode()?;
226 execute!(self.terminal.backend_mut(), LeaveAlternateScreen)?;
227 self.terminal.show_cursor()?;
228 Ok(())
229 }
230}
231
232pub fn start_ui(ui_state: Arc<Mutex<UIState>>) -> (std::thread::JoinHandle<()>, Arc<AtomicBool>) {
233 let running = Arc::new(AtomicBool::new(true));
234 let running_clone = running.clone();
235 let ui_state_clone = ui_state.clone();
236 let handle = std::thread::spawn(move || {
237 let mut ui = UI::new().unwrap();
238 while running_clone.load(Ordering::Relaxed) {
239 if event::poll(std::time::Duration::from_millis(0)).unwrap_or(false) {
240 if let Ok(Event::Key(key_event)) = event::read() {
241 if key_event.kind == KeyEventKind::Press || key_event.kind == KeyEventKind::Repeat {
242 match key_event.code {
243 KeyCode::Char('c') if key_event.modifiers.contains(KeyModifiers::CONTROL) => {
244 running_clone.store(false, Ordering::Relaxed);
245 break;
246 }
247 KeyCode::Up => {
248 let selected = ui.table_state.selected().unwrap_or(0);
249 if selected > 0 {
250 ui.table_state.select(Some(selected - 1));
251 }
252 }
253 KeyCode::Down => {
254 let selected = ui.table_state.selected().unwrap_or(0);
255 ui.table_state.select(Some(selected + 1));
256 }
257 KeyCode::Enter => {
258 let state = ui_state_clone.lock().unwrap();
259 let mut all_items_vec: Vec<(&String, &OperationStatus)> = state.all_operations.iter().collect();
260 all_items_vec.sort_by(|a, b| a.0.cmp(b.0));
261 if let Some(selected) = ui.table_state.selected() {
262 if selected < all_items_vec.len() {
263 let (name, status) = &all_items_vec[selected];
264 if let OperationStatus::Completed = status {
265 if let Some(url) = state.completed_urls.get(*name) {
266 let _ = webbrowser::open(url);
267 }
268 }
269 }
270 }
271 }
272 _ => {}
273 }
274 }
275 }
276 }
277 {
278 let state = ui_state_clone.lock().unwrap();
279 ui.draw(&state).unwrap();
280 }
281 std::thread::sleep(std::time::Duration::from_millis(16));
282 }
283 ui.restore().unwrap();
284 });
285 (handle, running)
286}
287
288pub fn stop_ui(handle: std::thread::JoinHandle<()>, running: Arc<AtomicBool>) {
289 running.store(false, std::sync::atomic::Ordering::Relaxed);
291 handle.join().unwrap();
292
293 let mut stdout = io::stdout();
295 stdout.execute(terminal::Clear(terminal::ClearType::All)).unwrap();
296 stdout.execute(cursor::MoveTo(0, 0)).unwrap();
297}