From ee7c8de4d4c724fdc0d8e7d45cbc08b4693f0f59 Mon Sep 17 00:00:00 2001
From: WorldObservationLog <wolc@duck.com>
Date: Sun, 5 May 2024 01:04:24 +0800
Subject: [PATCH] feat: codec alternative

---
 config.toml   |  6 ++++--
 src/config.py |  2 ++
 src/mp4.py    | 14 ++++++++++++--
 src/rip.py    |  4 ++--
 4 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/config.toml b/config.toml
index 129f720..eeb37e7 100644
--- a/config.toml
+++ b/config.toml
@@ -9,9 +9,11 @@ agentPort = 10020
 suMethod = "su -c"
 
 [download]
-atmosConventToM4a = false
+codecAlternative = true
+codecPriority = ["alac", "ec3", "ac3", "aac"]
+atmosConventToM4a = true
 songNameFormat = "{disk}-{tracknum:02d} {title}"
-dirPathFormat = "downloads/{artist}/{album}"
+dirPathFormat = "downloads/{album_artist}/{album}"
 saveLyrics = true
 saveCover = true
 coverFormat = "jpg"
diff --git a/src/config.py b/src/config.py
index 25c78a6..4567eba 100644
--- a/src/config.py
+++ b/src/config.py
@@ -16,6 +16,8 @@ class Device(BaseModel):
 
 
 class Download(BaseModel):
+    codecAlternative: bool
+    codecPriority: list[str]
     atmosConventToM4a: bool
     songNameFormat: str
     dirPathFormat: str
diff --git a/src/mp4.py b/src/mp4.py
index 451ec02..951455e 100644
--- a/src/mp4.py
+++ b/src/mp4.py
@@ -8,6 +8,7 @@ from typing import Tuple
 import m3u8
 import regex
 from bs4 import BeautifulSoup
+from loguru import logger
 
 from src.exceptions import CodecNotFoundException
 from src.metadata import SongMetadata
@@ -22,12 +23,21 @@ async def get_available_codecs(m3u8_url: str) -> Tuple[list[str], list[str]]:
     return codecs, codec_ids
 
 
-async def extract_media(m3u8_url: str, codec: str) -> Tuple[str, list[str], str]:
+async def extract_media(m3u8_url: str, codec: str, song_metadata: SongMetadata,
+                        codec_priority: list[str], alternative_codec: bool = False ) -> Tuple[str, list[str]]:
     parsed_m3u8 = m3u8.load(m3u8_url)
     specifyPlaylist = find_best_codec(parsed_m3u8, codec)
+    if not specifyPlaylist and alternative_codec:
+        logger.warning(f"Codec {codec} of song: {song_metadata.artist} - {song_metadata.title} did not found")
+        for a_codec in codec_priority:
+            specifyPlaylist = find_best_codec(parsed_m3u8, a_codec)
+            if specifyPlaylist:
+                codec = a_codec
+                break
     if not specifyPlaylist:
         raise CodecNotFoundException
     selected_codec = specifyPlaylist.media[0].group_id
+    logger.info(f"Selected codec: {selected_codec} for song: {song_metadata.artist} - {song_metadata.title}")
     stream = m3u8.load(specifyPlaylist.absolute_uri)
     skds = [key.uri for key in stream.keys if regex.match('(skd?://[^"]*)', key.uri)]
     keys = [prefetchKey]
@@ -46,7 +56,7 @@ async def extract_media(m3u8_url: str, codec: str) -> Tuple[str, list[str], str]
     for key in skds:
         if key.endswith(key_suffix) or key.endswith(CodecKeySuffix.KeySuffixDefault):
             keys.append(key)
-    return stream.segment_map[0].absolute_uri, keys, selected_codec
+    return stream.segment_map[0].absolute_uri, keys
 
 
 def extract_song(raw_song: bytes, codec: str) -> SongInfo:
diff --git a/src/rip.py b/src/rip.py
index 1c98266..4fddab5 100644
--- a/src/rip.py
+++ b/src/rip.py
@@ -29,8 +29,8 @@ async def rip_song(song: Song, auth_params: GlobalAuthParams, codec: str, config
         lyrics = await get_song_lyrics(song.id, song.storefront, auth_params.accountAccessToken,
                                        auth_params.dsid, auth_params.accountToken)
         song_metadata.lyrics = lyrics
-    song_uri, keys, selected_codec = await extract_media(song_data.attributes.extendedAssetUrls.enhancedHls, codec)
-    logger.info(f"Selected codec: {selected_codec} for song: {song_metadata.artist} - {song_metadata.title}")
+    song_uri, keys = await extract_media(song_data.attributes.extendedAssetUrls.enhancedHls, codec, song_metadata,
+                                         config.download.codecPriority, config.download.codecAlternative)
     logger.info(f"Downloading song: {song_metadata.artist} - {song_metadata.title}")
     raw_song = await download_song(song_uri)
     song_info = extract_song(raw_song, codec)