From 45bd435d34e52405d0b4e2ab0f997f2fddead584 Mon Sep 17 00:00:00 2001
From: ThetaDev <t.testboy@gmail.com>
Date: Mon, 8 Aug 2022 17:19:14 +0200
Subject: [PATCH] fix filename issue with downloader

---
 Cargo.toml      |  2 +-
 cli/Cargo.toml  |  1 +
 cli/src/main.rs | 20 ++++++++++++++------
 src/download.rs | 27 ++++++++++++++++++---------
 4 files changed, 34 insertions(+), 16 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 8d3d2d7..26b0545 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -26,7 +26,7 @@ async-trait = "0.1.56"
 chrono = {version = "0.4.19", features = ["serde"]}
 futures = "0.3.21"
 indicatif = "0.17.0"
-slug = "0.1.4"
+filenamify = "0.1.0"
 
 [dev-dependencies]
 env_logger = "0.9.0"
diff --git a/cli/Cargo.toml b/cli/Cargo.toml
index 02d7a3e..8dcf96e 100644
--- a/cli/Cargo.toml
+++ b/cli/Cargo.toml
@@ -11,4 +11,5 @@ tokio = {version = "1.20.0", features = ["rt-multi-thread"]}
 indicatif = "0.17.0"
 futures = "0.3.21"
 anyhow = "1.0"
+env_logger = "0.9.0"
 rusty-tube = {path = "../"}
diff --git a/cli/src/main.rs b/cli/src/main.rs
index 09803b8..5db4d64 100644
--- a/cli/src/main.rs
+++ b/cli/src/main.rs
@@ -1,4 +1,4 @@
-use anyhow::Result;
+use anyhow::{Context, Result};
 use futures::stream::{self, StreamExt};
 use indicatif::{MultiProgress, ProgressBar, ProgressStyle};
 use reqwest::{Client, ClientBuilder};
@@ -7,7 +7,7 @@ use rusty_tube::client::{ClientType, RustyTube};
 async fn download_video(
     video_id: String,
     output_dir: &str,
-    output_fname: &str,
+    output_fname: Option<String>,
     resolution: Option<u32>,
     ffmpeg: &str,
     rt: &RustyTube,
@@ -23,18 +23,26 @@ async fn download_video(
 
     let player_data = rt
         .get_player(video_id.as_str(), ClientType::Desktop)
-        .await?;
+        .await
+        .context(format!(
+            "Failed to fetch player data for video {}",
+            video_id
+        ))?;
 
     rusty_tube::download::download_video(
         &player_data,
         output_dir,
-        // output_fname,
+        output_fname,
         resolution,
         ffmpeg,
         http,
         pb,
     )
-    .await?;
+    .await
+    .context(format!(
+        "Failed to download video '{}' [{}]",
+        player_data.info.title, video_id
+    ))?;
 
     main.inc(1);
     Ok(())
@@ -73,7 +81,7 @@ async fn main() {
             download_video(
                 item.video_id.to_owned(),
                 "tmp",
-                "xyz",
+                None,
                 None,
                 "ffmpeg",
                 &rt,
diff --git a/src/download.rs b/src/download.rs
index ca21dfa..a99feff 100644
--- a/src/download.rs
+++ b/src/download.rs
@@ -1,9 +1,9 @@
-use std::{cmp::Ordering, ffi::OsString, fmt::format, ops::Range, path::PathBuf};
+use std::{cmp::Ordering, ffi::OsString, ops::Range, path::PathBuf};
 
 use anyhow::{anyhow, bail, Result};
 use fancy_regex::Regex;
 use futures::stream::{self, StreamExt};
-use indicatif::{MultiProgress, ProgressBar, ProgressStyle};
+use indicatif::ProgressBar;
 use log::debug;
 use once_cell::sync::Lazy;
 use rand::Rng;
@@ -180,11 +180,18 @@ struct StreamDownload {
 pub async fn download_video(
     player_data: &PlayerData,
     output_dir: &str,
+    output_fname: Option<String>,
     resolution: Option<u32>,
     ffmpeg: &str,
     http: Client,
     pb: ProgressBar,
 ) -> Result<()> {
+    // Download filepath
+    let download_dir = PathBuf::from(output_dir);
+    let title = player_data.info.title.to_owned();
+    let output_fname = output_fname
+        .unwrap_or_else(|| filenamify::filenamify(format!("{} [{}]", title, player_data.info.id)));
+
     // Select streams to download
     let video = match resolution {
         Some(r) => Some(some_or_bail!(
@@ -204,22 +211,22 @@ pub async fn download_video(
         Err(anyhow!("no audio stream"))
     );
 
-    let download_dir = PathBuf::from(output_dir);
-    let title = player_data.info.title.to_owned();
-    let title_fname = slug::slugify(&title); // TODO: slugify
-
     let mut downloads: Vec<StreamDownload> = Vec::new();
 
     video.map(|v| {
         downloads.push(StreamDownload {
-            file: download_dir.join(format!("{}.video{}", title_fname, v.format.extension())),
+            file: download_dir.join(format!("{}.video{}", output_fname, v.format.extension())),
             url: v.url.to_owned(),
             video_codec: Some(v.codec),
             audio_codec: None,
         });
     });
     downloads.push(StreamDownload {
-        file: download_dir.join(format!("{}.audio{}", title_fname, audio.format.extension())),
+        file: download_dir.join(format!(
+            "{}.audio{}",
+            output_fname,
+            audio.format.extension()
+        )),
         url: audio.url.to_owned(),
         video_codec: None,
         audio_codec: Some(audio.codec),
@@ -230,7 +237,7 @@ pub async fn download_video(
 
     pb.set_message(format!("Converting {}", title));
     // let output_file = download_dir.join(format!("{}.mp4", title_fname));
-    convert_streams(&downloads, download_dir.join(title_fname), ffmpeg).await?;
+    convert_streams(&downloads, download_dir.join(output_fname), ffmpeg).await?;
 
     // Delete original files
     stream::iter(&downloads)
@@ -324,6 +331,7 @@ async fn convert_streams<P: Into<PathBuf>>(
     Ok(())
 }
 
+/*
 #[cfg(test)]
 mod tests {
     use crate::client::RustyTube;
@@ -358,3 +366,4 @@ mod tests {
         //     .unwrap();
     }
 }
+*/