Skip to main content

bunkr_client/ui/
ui.rs

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    // Stop the UI
290    running.store(false, std::sync::atomic::Ordering::Relaxed);
291    handle.join().unwrap();
292
293    // Clear the UI and print final results
294    let mut stdout = io::stdout();
295    stdout.execute(terminal::Clear(terminal::ClearType::All)).unwrap();
296    stdout.execute(cursor::MoveTo(0, 0)).unwrap();
297}