From 2166a4aaf2d4213a477c0603dbba9fb860c7c4e9 Mon Sep 17 00:00:00 2001
From: Harry <HarryYC@users.noreply.github.com>
Date: Tue, 2 Apr 2024 16:01:01 -0700
Subject: [PATCH] Feat: Add progress bar to Youtube download (#481)

* download progress on console

* Feat: add a progress bar to Youtube download

* change download-on-state to interrupted if download failed
---
 enjoy/src/main/youtubedr.ts                   | 81 ++++++++++++-------
 .../videos/youtube-videos-segment.tsx         | 29 ++++++-
 enjoy/src/types/index.d.ts                    |  1 +
 3 files changed, 80 insertions(+), 31 deletions(-)

diff --git a/enjoy/src/main/youtubedr.ts b/enjoy/src/main/youtubedr.ts
index c5e1256e..4dba2072 100644
--- a/enjoy/src/main/youtubedr.ts
+++ b/enjoy/src/main/youtubedr.ts
@@ -1,12 +1,13 @@
 import { app } from "electron";
 import path from "path";
-import { exec } from "child_process";
+import { exec, spawn } from "child_process";
 import fs from "fs-extra";
 import os from "os";
 import log from "@main/logger";
 import snakeCase from "lodash/snakeCase";
 import settings from "@main/settings";
-import url from 'url';
+import url from "url";
+import mainWin from "@main/window";
 
 const __filename = url.fileURLToPath(import.meta.url);
 const __dirname = path.dirname(__filename);
@@ -88,14 +89,15 @@ class Youtubedr {
       quality?: string | number;
       filename?: string;
       directory?: string;
+      webContents?: Electron.WebContents;
     } = {}
   ): Promise<string> {
     const {
       quality,
       filename = this.getYtVideoId(url) + ".mp4",
       directory = app.getPath("downloads"),
+      webContents = mainWin.win.webContents,
     } = options;
-
     const command = [
       this.binFile,
       "download",
@@ -106,37 +108,58 @@ class Youtubedr {
     ].join(" ");
 
     logger.info(`Running command: ${command}`);
+
+    let currentSpeed = "";
+
     return new Promise((resolve, reject) => {
-      exec(
-        command,
-        {
-          timeout: ONE_MINUTE * 15,
-        },
-        (error, stdout, stderr) => {
-          if (error) {
-            logger.error("error", error);
-          }
+      const proc = spawn(this.binFile, [
+        "download",
+        url,
+        `--quality=${quality || "medium"}`,
+        `--filename=${filename}`,
+        `--directory=${directory}`,
+      ]);
 
-          if (stderr) {
-            logger.error("stderr", stderr);
-          }
+      proc.stdout.on("data", (data) => {
+        const output = data.toString();
+        const match = output.match(/iB (\d+) % \[/);
 
-          if (stdout) {
-            logger.debug(stdout);
-          }
-
-          if (fs.existsSync(path.join(directory, filename))) {
-            const stat = fs.statSync(path.join(directory, filename));
-            if (stat.size === 0) {
-              reject(new Error("Youtubedr download failed: empty file"));
-            } else {
-              resolve(path.join(directory, filename));
-            }
-          } else {
-            reject(new Error("Youtubedr download failed: unknown error"));
+        if (match) {
+          let speed = output.match(/\ ] (.*)/);
+          if (speed) {
+            currentSpeed = speed[1];
           }
+          webContents.send("download-on-state", {
+            name: filename,
+            state: "progressing",
+            received: parseInt(match[1]),
+            speed: currentSpeed,
+          });
         }
-      );
+      });
+
+      proc.on("close", (code) => {
+        if (code !== 0) {
+          webContents.send("download-on-state", {
+            name: filename,
+            state: "interrupted",
+          });
+          return reject(
+            new Error(`Youtubedr download failed with code: ${code}`)
+          );
+        }
+
+        if (fs.existsSync(path.join(directory, filename))) {
+          const stat = fs.statSync(path.join(directory, filename));
+          if (stat.size === 0) {
+            reject(new Error("Youtubedr download failed: empty file"));
+          } else {
+            resolve(path.join(directory, filename));
+          }
+        } else {
+          reject(new Error("Youtubedr download failed: unknown error"));
+        }
+      });
     });
   }
 
diff --git a/enjoy/src/renderer/components/videos/youtube-videos-segment.tsx b/enjoy/src/renderer/components/videos/youtube-videos-segment.tsx
index e39bef0f..80050b16 100755
--- a/enjoy/src/renderer/components/videos/youtube-videos-segment.tsx
+++ b/enjoy/src/renderer/components/videos/youtube-videos-segment.tsx
@@ -10,6 +10,7 @@ import {
   DialogTitle,
   DialogContent,
   DialogFooter,
+  Progress,
 } from "@renderer/components/ui";
 import { useNavigate } from "react-router-dom";
 import { LoaderIcon } from "lucide-react";
@@ -22,10 +23,14 @@ export const YoutubeVideosSegment = () => {
     null
   );
   const [submitting, setSubmitting] = useState(false);
+  const [progress, setProgress] = useState(0);
+  const [downloadSpeed, setDownloadSpeed] = useState(null);
 
   const addToLibrary = () => {
     let url = `https://www.youtube.com/watch?v=${selectedVideo?.videoId}`;
     setSubmitting(true);
+    setProgress(0);
+
     EnjoyApp.videos
       .create(url, {
         name: selectedVideo?.title,
@@ -64,6 +69,20 @@ export const YoutubeVideosSegment = () => {
     fetchYoutubeVideos();
   }, []);
 
+  useEffect(() => {
+    EnjoyApp.download.onState((_, downloadState) => {
+      const { state, received, speed } = downloadState;
+      if (state === "progressing") {
+        setProgress(received);
+        setDownloadSpeed(speed);
+      }
+    });
+
+    return () => {
+      EnjoyApp.download.removeAllListeners();
+    };
+  }, [submitting]);
+
   if (!videos?.length) return null;
 
   return (
@@ -101,7 +120,6 @@ export const YoutubeVideosSegment = () => {
           <DialogHeader>
             <DialogTitle>{t("downloadVideo")}</DialogTitle>
           </DialogHeader>
-
           <div className="flex items-center mb-4 bg-muted rounded-lg">
             <div className="aspect-square h-28 overflow-hidden rounded-l-lg">
               <img
@@ -119,7 +137,6 @@ export const YoutubeVideosSegment = () => {
               </div>
             </div>
           </div>
-
           <DialogFooter>
             <Button
               variant="ghost"
@@ -144,6 +161,14 @@ export const YoutubeVideosSegment = () => {
               {t("downloadVideo")}
             </Button>
           </DialogFooter>
+          {submitting && (
+            <div>
+              <Progress value={progress} className="mb-2" />
+              <div className="text-xs line-clamp-1 mb-2 text-right">
+                {downloadSpeed}
+              </div>
+            </div>
+          )}
         </DialogContent>
       </Dialog>
     </>
diff --git a/enjoy/src/types/index.d.ts b/enjoy/src/types/index.d.ts
index b6d7a767..2067ef5c 100644
--- a/enjoy/src/types/index.d.ts
+++ b/enjoy/src/types/index.d.ts
@@ -19,6 +19,7 @@ type DownloadStateType = {
   state: "progressing" | "interrupted" | "completed" | "cancelled";
   received: number;
   total: number;
+  speed?: string;
 };
 
 type NotificationType = {