cli: feat: add progress bars #25

warpdl · Apr 15, 2024 · 1d9e5b0 · 1d9e5b0
1 parent c22bca1
commit 1d9e5b0
Show file tree

Hide file tree

Showing 5 changed files with 47 additions and 26 deletions.
diff --git a/cmd/attach.go b/cmd/attach.go
@@ -43,6 +43,6 @@ Save Location`+"\t"+`: %s/
 		d.DownloadDirectory,
 	)
 	fmt.Println(txt)
-	RegisterHandlers(client)
+	RegisterHandlers(client, int64(d.ContentLength))
 	return client.Listen()
 }
diff --git a/cmd/client.go b/cmd/client.go
@@ -1,38 +1,53 @@
 package cmd
 
 import (
-	"fmt"
-	"os"
+	"time"
 
+	"github.com/vbauerster/mpb/v8"
+	cmdCommon "github.com/warpdl/warpdl/cmd/common"
 	"github.com/warpdl/warpdl/common"
 	"github.com/warpdl/warpdl/pkg/warpcli"
 	"github.com/warpdl/warpdl/pkg/warplib"
 )
 
-func downloadStopped(client *warpcli.Client) func(dr *common.DownloadingResponse) error {
+func downloadStopped(client *warpcli.Client, sc *SpeedCounter) func(dr *common.DownloadingResponse) error {
 	return func(dr *common.DownloadingResponse) error {
 		if dr.Hash != warplib.MAIN_HASH {
 			return nil
 		}
-		fmt.Println("Download Stopped: ", dr.DownloadId)
+		sc.bar.Abort(false)
+		// fmt.Println("Download Stopped: ", dr.DownloadId)
 		client.Disconnect()
-		os.Exit(0)
 		return nil
 	}
 }
 
-func downloadProgress(dr *common.DownloadingResponse) error {
-	fmt.Println(dr.Action, dr.DownloadId, dr.Hash, dr.Value)
-	return nil
+func downloadProgress(sc *SpeedCounter) func(dr *common.DownloadingResponse) error {
+	return func(dr *common.DownloadingResponse) error {
+		// fmt.Println(dr.Action, dr.DownloadId, dr.Hash, dr.Value)
+		sc.IncrBy(int(dr.Value))
+		return nil
+	}
 }
 
-func downloadComplete(client *warpcli.Client) func(dr *common.DownloadingResponse) error {
+func downloadComplete(client *warpcli.Client, dbar, cbar *mpb.Bar, sc *SpeedCounter) func(dr *common.DownloadingResponse) error {
 	return func(dr *common.DownloadingResponse) error {
-		fmt.Println("Download Complete: ", dr.Hash)
-		if dr.Hash == warplib.MAIN_HASH {
-			client.Disconnect()
+		// fmt.Println("Download Complete: ", dr.Hash)
+		if dr.Hash != warplib.MAIN_HASH {
+			return nil
+		}
+		defer client.Disconnect()
+		sc.Stop()
+		// fill download bar
+		if dbar.Completed() {
+			return nil
+		}
+		dbar.SetCurrent(dr.Value)
+		// fill compile bar
+		if cbar.Completed() {
+			return nil
 		}
-		os.Exit(0)
+		cbar.SetCurrent(dr.Value)
 		return nil
 	}
 }
@@ -42,34 +57,42 @@ func compileStart(dr *common.DownloadingResponse) error {
 }
 
 func compileComplete(dr *common.DownloadingResponse) error {
-	fmt.Println("Compile Complete: ", dr.Hash)
+	// fmt.Println("Compile Complete: ", dr.Hash)
 	return nil
 }
 
-func compileProgress(dr *common.DownloadingResponse) error {
-	return nil
+func compileProgress(bar *mpb.Bar) func(dr *common.DownloadingResponse) error {
+	return func(dr *common.DownloadingResponse) error {
+		bar.IncrBy(int(dr.Value))
+		return nil
+	}
 }
 
-func RegisterHandlers(client *warpcli.Client) {
+func RegisterHandlers(client *warpcli.Client, contentLength int64) {
+	sc := NewSpeedCounter(4350 * time.Microsecond)
+	p := mpb.New(mpb.WithWidth(64), mpb.WithRefreshRate(time.Millisecond*100))
+	dbar, cbar := cmdCommon.InitBars(p, "", contentLength)
+	sc.SetBar(dbar)
+	sc.Start()
 	client.AddHandler(
 		common.UPDATE_DOWNLOADING,
-		warpcli.NewDownloadingHandler(common.DownloadStopped, downloadStopped(client)),
+		warpcli.NewDownloadingHandler(common.DownloadStopped, downloadStopped(client, sc)),
 	)
 	client.AddHandler(
 		common.UPDATE_DOWNLOADING,
-		warpcli.NewDownloadingHandler(common.DownloadProgress, downloadProgress),
+		warpcli.NewDownloadingHandler(common.DownloadProgress, downloadProgress(sc)),
 	)
 	client.AddHandler(
 		common.UPDATE_DOWNLOADING,
-		warpcli.NewDownloadingHandler(common.DownloadComplete, downloadComplete(client)),
+		warpcli.NewDownloadingHandler(common.DownloadComplete, downloadComplete(client, dbar, cbar, sc)),
 	)
 	client.AddHandler(
 		common.UPDATE_DOWNLOADING,
 		warpcli.NewDownloadingHandler(common.CompileComplete, compileComplete),
 	)
 	client.AddHandler(
 		common.UPDATE_DOWNLOADING,
-		warpcli.NewDownloadingHandler(common.CompileProgress, compileProgress),
+		warpcli.NewDownloadingHandler(common.CompileProgress, compileProgress(cbar)),
 	)
 	client.AddHandler(
 		common.UPDATE_DOWNLOADING,

diff --git a/cmd/download.go b/cmd/download.go
@@ -83,6 +83,6 @@ Max Connections`+"\t"+`: %d
 		txt += fmt.Sprintf("Max Segments\t: %d\n", d.MaxSegments)
 	}
 	fmt.Println(txt)
-	RegisterHandlers(client)
+	RegisterHandlers(client, int64(d.ContentLength))
 	return client.Listen()
 }
diff --git a/cmd/resume.go b/cmd/resume.go
@@ -103,6 +103,6 @@ Max Connections`+"\t"+`: %d
 		txt += fmt.Sprintf("Max Segments\t: %d\n", r.MaxSegments)
 	}
 	fmt.Println(txt)
-	RegisterHandlers(client)
+	RegisterHandlers(client, int64(r.ContentLength))
 	return client.Listen()
 }
diff --git a/pkg/warpcli/dispatcher.go b/pkg/warpcli/dispatcher.go
@@ -25,8 +25,6 @@ func (d *Dispatcher) process(buf []byte) error {
 	if !res.Ok {
 		return errors.New(res.Error)
 	}
-	fmt.Println(res.Update.Type)
-	fmt.Println(string(res.Update.Message))
 	d.mu.RLock()
 	handlers, ok := d.Handlers[res.Update.Type]
 	d.mu.RUnlock()