1  use anyhow::{Context, Result};
  2  use indicatif::{MultiProgress, ProgressBar, ProgressStyle};
  3  use std::collections::HashSet; // Import HashSet
  4  use std::os::unix::fs::MetadataExt; // Import for ino() and dev()
  5  use std::path::{Path, PathBuf};
  6  use std::sync::Arc;
  7  use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
  8  use std::time::{Duration, Instant};
  9  
 10  /// Tracks progress during file transfers
 11  pub struct TransferProgress {
 12      total_files: Arc<AtomicUsize>,
 13      current_file: Arc<AtomicUsize>,
 14      total_bytes: Arc<AtomicU64>,
 15      bytes_transferred: Arc<AtomicU64>,
 16      uncompressed_bytes_processed: Arc<AtomicU64>,
 17      start_time: Instant,
 18      current_filename: Arc<std::sync::Mutex<String>>,
 19      compression_ratio: Arc<std::sync::Mutex<Option<f64>>>,
 20  }
 21  
 22  /// Manages the visual progress display
 23  pub struct ProgressDisplay {
 24      multi_progress: MultiProgress,
 25      pub main_bar: ProgressBar,
 26      pub status_bar: ProgressBar,
 27  }
 28  
 29  impl TransferProgress {
 30      /// Create a new progress tracker
 31      pub fn new() -> Self {
 32          Self {
 33              total_files: Arc::new(AtomicUsize::new(0)),
 34              current_file: Arc::new(AtomicUsize::new(0)),
 35              total_bytes: Arc::new(AtomicU64::new(0)),
 36              bytes_transferred: Arc::new(AtomicU64::new(0)),
 37              uncompressed_bytes_processed: Arc::new(AtomicU64::new(0)),
 38              start_time: Instant::now(),
 39              current_filename: Arc::new(std::sync::Mutex::new(String::new())),
 40              compression_ratio: Arc::new(std::sync::Mutex::new(None)),
 41          }
 42      }
 43  
 44      /// Calculate total size and file count for multiple source paths,
 45      /// correctly handling hard links to avoid double-counting.
 46      pub fn calculate_multiple_sources_info(&self, sources: &[PathBuf]) -> Result<(usize, u64)> {
 47          let mut total_files = 0;
 48          let mut total_bytes = 0;
 49          // A set to store (device_id, inode_id) tuples we've already seen
 50          let mut seen_inodes = HashSet::new();
 51  
 52          for source in sources {
 53              if source.is_file() {
 54                  let metadata = source.metadata().context("Failed to get file metadata")?;
 55                  let inode_key = (metadata.dev(), metadata.ino());
 56  
 57                  // Only count if we haven't seen this inode before
 58                  if seen_inodes.insert(inode_key) {
 59                      total_files += 1;
 60                      total_bytes += metadata.len();
 61                  }
 62              } else if source.is_dir() {
 63                  // Pass the tracking set into the recursive walker
 64                  self.walk_directory(source, &mut total_files, &mut total_bytes, &mut seen_inodes)?;
 65              }
 66          }
 67  
 68          self.total_files.store(total_files, Ordering::Relaxed);
 69          self.total_bytes.store(total_bytes, Ordering::Relaxed);
 70  
 71          Ok((total_files, total_bytes))
 72      }
 73  
 74      /// Recursively walks a directory to calculate size and file count,
 75      /// using a HashSet to avoid counting data from hard links more than once.
 76      fn walk_directory(
 77          &self,
 78          dir: &Path,
 79          files: &mut usize,
 80          bytes: &mut u64,
 81          seen_inodes: &mut HashSet<(u64, u64)>,
 82      ) -> Result<()> {
 83          for entry in std::fs::read_dir(dir).context("Failed to read directory")? {
 84              let entry = entry.context("Failed to read directory entry")?;
 85              let path = entry.path();
 86  
 87              // We use `symlink_metadata` to avoid following symlinks, but get metadata.
 88              let metadata = entry.metadata().context("Failed to get file metadata")?;
 89  
 90              if metadata.is_file() {
 91                  let inode_key = (metadata.dev(), metadata.ino());
 92  
 93                  // Only count the file and its size if it's a new inode.
 94                  if seen_inodes.insert(inode_key) {
 95                      *files += 1;
 96                      *bytes += metadata.len();
 97                  }
 98              } else if metadata.is_dir() {
 99                  // Recurse into subdirectory, passing the set along.
100                  self.walk_directory(&path, files, bytes, seen_inodes)?;
101              }
102          }
103          Ok(())
104      }
105  
106      /// Update current file being processed
107      pub fn set_current_file(&self, filename: &str) {
108          let mut current = self.current_filename.lock().unwrap();
109          *current = filename.to_string();
110          self.current_file.fetch_add(1, Ordering::Relaxed);
111      }
112  
113      /// Update bytes transferred (compressed bytes to the network)
114      pub fn add_bytes_transferred(&self, bytes: u64) {
115          self.bytes_transferred.fetch_add(bytes, Ordering::Relaxed);
116      }
117  
118      /// Update uncompressed bytes processed (original bytes from disk)
119      pub fn add_uncompressed_bytes_processed(&self, bytes: u64) {
120          self.uncompressed_bytes_processed
121              .fetch_add(bytes, Ordering::Relaxed);
122      }
123  
124      /// Set compression ratio
125      pub fn set_compression_ratio(&self, ratio: f64) {
126          let mut ratio_lock = self.compression_ratio.lock().unwrap();
127          *ratio_lock = Some(ratio);
128      }
129  
130      /// Get current progress values
131      pub fn get_progress(&self) -> ProgressSnapshot {
132          let current_file_idx = self.current_file.load(Ordering::Relaxed);
133          let total_files = self.total_files.load(Ordering::Relaxed);
134          let bytes_transferred = self.bytes_transferred.load(Ordering::Relaxed);
135          let total_bytes = self.total_bytes.load(Ordering::Relaxed);
136          let uncompressed_bytes_processed =
137              self.uncompressed_bytes_processed.load(Ordering::Relaxed);
138          let elapsed = self.start_time.elapsed();
139  
140          let current_filename = self.current_filename.lock().unwrap().clone();
141          let compression_ratio = *self.compression_ratio.lock().unwrap();
142  
143          ProgressSnapshot {
144              current_file: current_file_idx,
145              total_files,
146              bytes_transferred,
147              total_bytes,
148              uncompressed_bytes_processed,
149              elapsed,
150              current_filename,
151              compression_ratio,
152          }
153      }
154  }
155  
156  impl ProgressDisplay {
157      /// Create a new progress display
158      pub fn new() -> Result<Self> {
159          let multi_progress = MultiProgress::new();
160  
161          // Enhanced main progress bar with better formatting
162          let main_style = ProgressStyle::with_template(
163              "{spinner:.green} [{elapsed_precise}] {bar:40.cyan/blue} {percent:>3}% {bytes:>9}/{total_bytes:>9} {msg}",
164          )?
165          .progress_chars("█▉▊▋▌▍▎▏  ");
166  
167          // Enhanced status bar with colors and better formatting
168          let status_style = ProgressStyle::with_template("{msg}")?;
169  
170          let main_bar = multi_progress.add(ProgressBar::new(0));
171          main_bar.set_style(main_style);
172  
173          let status_bar = multi_progress.add(ProgressBar::new(0));
174          status_bar.set_style(status_style);
175          status_bar.set_length(1); // Indeterminate
176  
177          Ok(Self {
178              multi_progress,
179              main_bar,
180              status_bar,
181          })
182      }
183  
184      /// Initialize the progress display with source info
185      pub fn initialize(&self, files: usize, bytes: u64, source_name: String) {
186          self.main_bar.set_length(bytes);
187  
188          let formatted_bytes = format_bytes(bytes as f64);
189          let header_msg = format!("Transfer: {} ({}, {})", source_name, files, formatted_bytes);
190          self.main_bar.set_message(header_msg);
191          self.status_bar.set_message("Preparing transfer...");
192      }
193  
194      /// Update the progress display
195      pub fn update(&self, progress: &ProgressSnapshot) {
196          // NOTE: The position is now set in main.rs to use the correct value
197          // self.main_bar.set_position(progress.bytes_transferred);
198  
199          // Calculate speed and ETA based on actual bytes transferred for network speed
200          let speed = if progress.elapsed.as_secs() > 0 {
201              progress.bytes_transferred as f64 / progress.elapsed.as_secs_f64()
202          } else {
203              0.0
204          };
205  
206          // Calculate ETA based on uncompressed bytes remaining
207          let remaining_bytes = progress
208              .total_bytes
209              .saturating_sub(progress.uncompressed_bytes_processed);
210          let eta_seconds = if speed > 0.0 && progress.bytes_transferred > 0 {
211              // Estimate remaining time based on current compression ratio
212              let estimated_total_compressed = (progress.bytes_transferred as f64
213                  / progress.uncompressed_bytes_processed as f64)
214                  * progress.total_bytes as f64;
215              let remaining_compressed_bytes =
216                  estimated_total_compressed - progress.bytes_transferred as f64;
217              remaining_compressed_bytes / speed
218          } else {
219              0.0
220          };
221  
222          let speed_str = format_bytes(speed);
223          let eta_str = format_duration(eta_seconds);
224  
225          let mut status_msg = format!("Speed: {}/s | ETA: {}", speed_str, eta_str);
226  
227          if let Some(ratio) = progress.compression_ratio {
228              status_msg.push_str(&format!(" | Compression: {:.1}%", ratio * 100.0));
229          }
230  
231          self.status_bar.set_message(status_msg);
232      }
233  
234      /// Complete the progress display
235      pub fn finish(&self, progress: &ProgressSnapshot) {
236          let total_time = format_duration(progress.elapsed.as_secs_f64());
237          // Use bytes_transferred for average speed calculation
238          let avg_speed =
239              format_bytes(progress.bytes_transferred as f64 / progress.elapsed.as_secs_f64());
240          let total_bytes_str = format_bytes(progress.bytes_transferred as f64);
241  
242          let completion_msg = format!("Transfer complete: {} in {}", total_bytes_str, total_time);
243  
244          let summary_msg = format!(
245              "Average speed: {}/s | Total time: {}",
246              avg_speed, total_time
247          );
248  
249          self.main_bar.finish_with_message(completion_msg);
250          self.status_bar.finish_with_message(summary_msg);
251      }
252  
253      /// Handle error cleanup
254      pub fn error(&self, error: &str) {
255          self.main_bar.finish_with_message("Transfer failed");
256          let message = format!("Error: {}", error);
257          self.status_bar.finish_with_message(message);
258      }
259  }
260  
261  /// Snapshot of current progress
262  #[derive(Debug, Clone)]
263  pub struct ProgressSnapshot {
264      pub current_file: usize,
265      pub total_files: usize,
266      pub bytes_transferred: u64,
267      pub total_bytes: u64,
268      pub uncompressed_bytes_processed: u64,
269      pub elapsed: Duration,
270      pub current_filename: String,
271      pub compression_ratio: Option<f64>,
272  }
273  
274  /// Format bytes into human-readable format with color codes
275  fn format_bytes(bytes: f64) -> String {
276      const UNITS: &[&str] = &["B", "KB", "MB", "GB", "TB"];
277      let mut size = bytes;
278      let mut unit_index = 0;
279  
280      while size >= 1024.0 && unit_index < UNITS.len() - 1 {
281          size /= 1024.0;
282          unit_index += 1;
283      }
284  
285      if unit_index == 0 {
286          format!("{:.0} {}", size, UNITS[unit_index])
287      } else if size >= 100.0 {
288          format!("{:.0} {}", size, UNITS[unit_index])
289      } else {
290          format!("{:.1} {}", size, UNITS[unit_index])
291      }
292  }
293  
294  /// Format duration into human-readable format
295  fn format_duration(seconds: f64) -> String {
296      if seconds.is_nan() || seconds.is_infinite() || seconds < 0.0 {
297          return "n/a".to_string();
298      }
299      if seconds < 60.0 {
300          format!("{:.0}s", seconds)
301      } else if seconds < 3600.0 {
302          let minutes = (seconds / 60.0).floor();
303          let secs = (seconds % 60.0).round();
304          if secs == 0.0 {
305              format!("{:.0}m", minutes)
306          } else {
307              format!("{:.0}m {:.0}s", minutes, secs)
308          }
309      } else {
310          let hours = (seconds / 3600.0).floor();
311          let minutes = ((seconds % 3600.0) / 60.0).round();
312          if minutes == 0.0 {
313              format!("{:.0}h", hours)
314          } else {
315              format!("{:.0}h {:.0}m", hours, minutes)
316          }
317      }
318  }