use walkdir::WalkDir;
use indexmap::IndexMap;
use serde::{Deserialize, Serialize};
use std::{
fs,
sync::atomic,
time::SystemTime,
collections::VecDeque,
path::{Path, PathBuf},
};
pub type FileId = usize;
static NEXT_OUTPUT_FILEID: atomic::AtomicUsize = atomic::AtomicUsize::new(0);
fn get_new_output_file_id() -> usize {
NEXT_OUTPUT_FILEID.fetch_add(1, atomic::Ordering::Relaxed) }
#[derive(Default)]
pub struct FileManager {
pub ignore_empty: bool,
pub output_queue: VecDeque<OutputFile>,
pub input_map: IndexMap<FileId, InputFile>,
pub output_map: IndexMap<FileId, OutputFile>,
pub input_speed_counter: SpeedCounter,
pub output_speed_counter: SpeedCounter,
}
impl FileManager {
pub fn new(ignore_empty: bool) -> Self {
Self { ignore_empty, ..Default::default() }
}
}
impl FileManager {
pub fn add_output_files(&mut self, files: &Vec<PathBuf>) -> color_eyre::Result<()> {
let mut output_files: Vec<OutputFile> = vec![];
for path in files {
if path.is_dir() {
let mut empty_directories: Vec<PathBuf> = vec![];
if !self.ignore_empty {
for entry in WalkDir::new(path)
.into_iter()
.filter_map(Result::ok)
.filter(|e| e.file_type().is_dir())
{
let path = entry.path();
if fs::read_dir(path)
.map(|mut i| i.next().is_none())
.unwrap_or(false)
{
empty_directories.push(path.into());
}
}
}
let directory_files: Vec<PathBuf> = WalkDir::new(path)
.into_iter()
.filter_map(Result::ok)
.filter(|entry| entry.file_type().is_file())
.map(|entry| entry.path().to_path_buf())
.collect();
for p in empty_directories {
let of = OutputFile::new(p, Some(path.clone()), true)?;
output_files.push(of);
}
for p in directory_files {
let of = OutputFile::new(p, Some(path.clone()), false)?;
output_files.push(of);
}
} else {
let of = OutputFile::new(path.clone(), None, false)?;
output_files.push(of);
}
}
self.output_queue.extend(output_files.iter().cloned());
for file in output_files {
self.output_map.insert(file.id, file.clone());
}
Ok(())
}
pub fn get_next_output_file(&mut self) -> Option<OutputFile> {
self.output_queue.pop_front()
}
pub fn get_input_map(&self) -> IndexMap<&FileId, &InputFile> {
self.input_map.iter().collect()
}
pub fn get_output_map_no_dir(&self) -> IndexMap<&FileId, &OutputFile> {
self.output_map
.iter()
.filter(|v| !v.1.meta.is_dir)
.collect()
}
pub fn set_output_finished(&mut self, id: FileId) {
if let Some(output_file) = self.output_map.get_mut(&id) {
output_file.finished = true;
}
}
pub fn get_estimate<P: IOFile>(files: &IndexMap<FileId, P>, speed: f64) -> f64 {
let mut total_size: f64 = 0.0;
for (_i, f) in files {
if !f.get_meta().is_dir && !f.get_finished() {
total_size += (f.get_meta().size as f64) * (1.0 - f.get_progress());
}
}
if total_size > 0.0 {
(total_size * 8.0 / 1_000_000.0) / speed
} else {
0.0
}
}
pub fn get_completion<P: IOFile>(files: &IndexMap<FileId, P>) -> bool {
if !files.is_empty() {
let mut result = true;
for (_i, f) in files {
if !f.get_finished() {
result = false;
}
}
result
} else {
false
}
}
}
pub trait IOFile {
fn get_name(&self) -> Option<&str>;
fn get_progress(&self) -> f64;
fn get_finished(&self) -> bool;
fn get_meta(&self) -> &MetaData;
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct OutputFile {
pub id: FileId,
pub meta: MetaData,
pub progress: f64,
pub finished: bool,
}
impl OutputFile {
fn new(path: PathBuf, base_path: Option<PathBuf>, is_dir: bool) -> color_eyre::Result<Self> {
let meta: MetaData = if is_dir {
MetaData::new(&path, 0, base_path.clone(), true)
} else {
let metadata = fs::metadata(path.clone())?;
MetaData::new(&path, metadata.len() as usize, base_path.clone(), false)
};
Ok(Self {
id: get_new_output_file_id(),
meta,
progress: 0.0,
finished: false,
})
}
}
impl IOFile for OutputFile {
fn get_name(&self) -> Option<&str> {
let name = self.meta.path.file_name();
if let Some(name) = name { name.to_str() } else { None }
}
fn get_progress(&self) -> f64 { self.progress }
fn get_finished(&self) -> bool { self.finished }
fn get_meta(&self) -> &MetaData { &self.meta }
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct InputFile {
pub id: FileId,
pub meta: MetaData,
pub progress: f64,
}
impl InputFile {
pub fn new(id: usize, meta: MetaData) -> Self {
Self { id, meta, progress: 0.0 }
}
}
impl IOFile for InputFile {
fn get_name(&self) -> Option<&str> {
Some(&self.meta.name)
}
fn get_progress(&self) -> f64 {
self.progress
}
fn get_finished(&self) -> bool {
self.progress >= 1.0
}
fn get_meta(&self) -> &MetaData {
&self.meta
}
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct MetaData {
pub is_dir: bool,
pub path: PathBuf,
pub base_path: Option<PathBuf>,
pub name: String,
pub extension: String,
pub size: usize,
pub progress_bytes: usize,
}
impl MetaData {
pub fn new(path: &Path, size: usize, base_path: Option<PathBuf>, is_dir: bool) -> Self {
let p = MetaData::normalize_path(path);
let mut name: String = "".to_string();
let mut extension: String = "".to_string();
if let Some(n) = p.file_name()
&& let Some(s) = n.to_str()
{
name = s.into();
extension = s.into();
};
Self {
is_dir,
base_path,
name,
extension,
size,
progress_bytes: 0,
path: p,
}
}
fn normalize_path(path: &Path) -> PathBuf {
path.to_string_lossy().replace('\\', "/").into()
}
pub fn local_path(&self) -> Option<PathBuf> {
let mut result: Option<PathBuf> = None;
if let Some(path) = self.base_path.clone() {
let cloned_path = path.clone();
let parent = cloned_path.file_name();
if let Some(parent) = parent {
let stripped_path = self.path.strip_prefix(path);
if let Ok(stripped_path) = stripped_path {
let mut os_string = parent.to_os_string();
os_string.push("/");
os_string.push(stripped_path);
result = Some(os_string.into());
}
}
}
result
}
pub fn get_path(&self) -> PathBuf {
if let Some(local_path) = self.local_path() {
local_path
} else {
self.name.clone().into()
}
}
}
#[derive(Clone, Debug)]
pub struct FileProgressReport {
pub file_id: FileId,
pub progress: f64,
}
impl FileProgressReport {
pub fn new(file_id: FileId, progress: f64) -> Self {
Self { file_id, progress }
}
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct SpeedReport {
timestamp: SystemTime,
bytes: u32,
}
impl SpeedReport {
pub fn new(bytes: u32) -> Self {
Self { bytes, timestamp: SystemTime::now() }
}
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct SpeedCounter {
init_flag: bool,
pub speed: f64,
timestamp: SystemTime,
}
impl Default for SpeedCounter {
fn default() -> Self {
Self { init_flag: false, speed: 0.0, timestamp: SystemTime::now() }
}
}
impl SpeedCounter {
pub fn update(&mut self, report: SpeedReport) {
if !self.init_flag {
self.timestamp = report.timestamp;
self. init_flag = true;
} else {
let duration = report.timestamp.duration_since(self.timestamp).unwrap(); let megabits = ((report.bytes * 8) as f64) / 1_000_000.0;
let speed = megabits / duration.as_secs_f64();
self.timestamp = report.timestamp;
self.speed = if self.speed == 0.0 { speed } else { self.speed + (speed - self.speed) * 0.75 } }
}
}