From ab9eaaf339b8afbe70c87fc2b2744ec188b190e7 Mon Sep 17 00:00:00 2001
From: Mathis Maquenne <124215603+mathismqn@users.noreply.github.com>
Date: Thu, 10 Oct 2024 00:17:33 +0200
Subject: [PATCH] feat(cli): add quality flag

---
 cmd/download.go          | 44 ++++++++++++++++++++++++++++++++--------
 internal/deezer/media.go | 14 +------------
 internal/deezer/song.go  | 16 +++++++++++++--
 3 files changed, 51 insertions(+), 23 deletions(-)

diff --git a/cmd/download.go b/cmd/download.go
index fc119c6..7e0428c 100644
--- a/cmd/download.go
+++ b/cmd/download.go
@@ -12,15 +12,26 @@ import (
 var downloadCmd = &cobra.Command{
 	Use:   "download [album_id...]",
 	Short: "Download songs from one or more albums",
+	Args:  cobra.MinimumNArgs(1),
 	Run: func(cmd *cobra.Command, args []string) {
-		if len(args) == 0 {
-			cmd.Help()
-			return
-		}
-
 		output, _ := cmd.Flags().GetString("output")
+		quality, _ := cmd.Flags().GetString("quality")
 		nAlbums := len(args)
 
+		if quality == "" {
+			quality = "best"
+		}
+		validQualities := map[string]bool{
+			"mp3_128": true,
+			"mp3_320": true,
+			"flac":    true,
+			"best":    true,
+		}
+		if !validQualities[quality] {
+			fmt.Fprintf(os.Stderr, "invalid quality option: %s\n", quality)
+			os.Exit(1)
+		}
+
 		for i, id := range args {
 			fmt.Printf("[%d/%d] Getting data for album %s...", i+1, nAlbums, id)
 			album, err := deezer.GetAlbumData(id)
@@ -42,7 +53,7 @@ var downloadCmd = &cobra.Command{
 
 			for _, song := range album.Songs.Data {
 				fmt.Printf("\nDownloading %s...", song.Title)
-				media, err := song.GetMediaData()
+				media, err := song.GetMediaData(quality)
 				if err != nil {
 					fmt.Fprintf(os.Stderr, " could not get media data: %v\n", err)
 					if err.Error() == "invalid license token" {
@@ -51,13 +62,29 @@ var downloadCmd = &cobra.Command{
 					continue
 				}
 
+				if len(media.Data) == 0 || len(media.Data[0].Media) == 0 || len(media.Data[0].Media[0].Sources) == 0 {
+					fmt.Fprintf(os.Stderr, " could not get media sources\n")
+					continue
+				}
+
+				url := media.Data[0].Media[0].Sources[0].URL
+				for _, source := range media.Data[0].Media[0].Sources {
+					if source.Provider == "ak" {
+						url = source.URL
+						break
+					}
+				}
 				songTitle := song.Title
 				if song.Version != "" {
 					songTitle = fmt.Sprintf("%s %s", song.Title, song.Version)
 				}
+				ext := "mp3"
+				if media.Data[0].Media[0].Format == "FLAC" {
+					ext = "flac"
+				}
 
-				path := path.Join(output, fmt.Sprintf("%s - %s.flac", song.ArtistName, songTitle))
-				err = media.Download(path, song.ID)
+				path := path.Join(output, fmt.Sprintf("%s - %s.%s", song.ArtistName, songTitle, ext))
+				err = media.Download(url, path, song.ID)
 				if err != nil {
 					fmt.Fprintf(os.Stderr, " could not download song: %v\n", err)
 					continue
@@ -71,4 +98,5 @@ var downloadCmd = &cobra.Command{
 func init() {
 	rootCmd.AddCommand(downloadCmd)
 	downloadCmd.Flags().StringP("output", "o", "", "output directory (default is current directory)")
+	downloadCmd.Flags().StringP("quality", "q", "", "download quality [mp3_128, mp3_320, flac, best] (default is best)")
 }
diff --git a/internal/deezer/media.go b/internal/deezer/media.go
index b8d558d..84e4bbb 100644
--- a/internal/deezer/media.go
+++ b/internal/deezer/media.go
@@ -36,19 +36,7 @@ type Source struct {
 
 const ChunkSize = 2048
 
-func (m *Media) Download(path, songID string) error {
-	if len(m.Data) == 0 || len(m.Data[0].Media) == 0 || len(m.Data[0].Media[0].Sources) == 0 {
-		return fmt.Errorf("no media sources found")
-	}
-
-	url := m.Data[0].Media[0].Sources[0].URL
-	for _, source := range m.Data[0].Media[0].Sources {
-		if source.Provider == "ak" {
-			url = source.URL
-			break
-		}
-	}
-
+func (m *Media) Download(url, path, songID string) error {
 	resp, err := http.Get(url)
 	if err != nil {
 		return err
diff --git a/internal/deezer/song.go b/internal/deezer/song.go
index b3cb04b..1843ac3 100644
--- a/internal/deezer/song.go
+++ b/internal/deezer/song.go
@@ -18,9 +18,21 @@ type Song struct {
 	TrackToken string `json:"TRACK_TOKEN"`
 }
 
-func (s *Song) GetMediaData() (*Media, error) {
-	reqBody := fmt.Sprintf(`{"license_token":"%s","media":[{"type":"FULL","formats":[{"cipher":"BF_CBC_STRIPE","format":"FLAC"}]}],"track_tokens":["%s"]}`, config.Cfg.LicenseToken, s.TrackToken)
+func (s *Song) GetMediaData(quality string) (*Media, error) {
+	var formats string
 
+	switch quality {
+	case "mp3_128":
+		formats = `[{"cipher":"BF_CBC_STRIPE","format":"MP3_128"}]`
+	case "mp3_320":
+		formats = `[{"cipher":"BF_CBC_STRIPE","format":"MP3_320"}]`
+	case "flac":
+		formats = `[{"cipher":"BF_CBC_STRIPE","format":"FLAC"}]`
+	case "best":
+		formats = `[{"cipher":"BF_CBC_STRIPE","format":"FLAC"},{"cipher":"BF_CBC_STRIPE","format":"MP3_320"},{"cipher":"BF_CBC_STRIPE","format":"MP3_128"}]`
+	}
+
+	reqBody := fmt.Sprintf(`{"license_token":"%s","media":[{"type":"FULL","formats":%s}],"track_tokens":["%s"]}`, config.Cfg.LicenseToken, formats, s.TrackToken)
 	resp, err := http.Post("https://media.deezer.com/v1/get_url", "application/json", bytes.NewBuffer([]byte(reqBody)))
 	if err != nil {
 		return nil, err