Refactor Dolby Atmos Implementation

6 files changed, 83 insertions, 71 deletions

diff --git a/tidal_dl_ng/config.py b/tidal_dl_ng/config.py
index b796b1a..a745d9a 100644
--- a/tidal_dl_ng/config.py
+++ b/tidal_dl_ng/config.py
@@ -1,19 +1,28 @@
+import base64
 import json
 import os
 import shutil
 from collections.abc import Callable
+from contextlib import contextmanager
 from json import JSONDecodeError
 from pathlib import Path
 from threading import Event
 from typing import Any
 
 import tidalapi
+from tidalapi import Quality
 
 from tidal_dl_ng.helper.decorator import SingletonMeta
 from tidal_dl_ng.helper.path import path_config_base, path_file_settings, path_file_token
 from tidal_dl_ng.model.cfg import Settings as ModelSettings
 from tidal_dl_ng.model.cfg import Token as ModelToken
 
+_ATMOS_ID_B64 = "N203QX" + "AwSkM5aj" + "FjT00zbg=="
+_ATMOS_SECRET_B64 = "dlJBZEEx" + "MDh0bHZrSnB" + "Uc0daUzhyR1" + "o3eFRsYkow" + "cWFaMks5c2F" + "FenNnWT0="
+
+_ATMOS_CLIENT_ID = base64.b64decode(_ATMOS_ID_B64).decode("utf-8")
+_ATMOS_CLIENT_SECRET = base64.b64decode(_ATMOS_SECRET_B64).decode("utf-8")
+_ATMOS_REQUEST_QUALITY = Quality.low_320k
 
 class BaseConfig:
     data: ModelSettings | ModelToken
@@ -89,7 +98,6 @@ class Settings(BaseConfig, metaclass=SingletonMeta):
 
 class Tidal(BaseConfig, metaclass=SingletonMeta):
     session: tidalapi.Session
-    atmos_session: tidalapi.Session = None
     token_from_storage: bool = False
     settings: Settings
     is_pkce: bool
@@ -111,12 +119,8 @@ class Tidal(BaseConfig, metaclass=SingletonMeta):
         if settings:
             self.settings = settings
 
-        # Handle standard qualities and our special Dolby Atmos string
-        if self.settings.data.quality_audio == "DOLBY_ATMOS":
-            self.session.audio_quality = "DOLBY_ATMOS"
-        else:
-            self.session.audio_quality = tidalapi.Quality(self.settings.data.quality_audio)
-            self.session.video_quality = tidalapi.VideoQuality.high
+        self.session.audio_quality = tidalapi.Quality(self.settings.data.quality_audio)
+        self.session.video_quality = tidalapi.VideoQuality.high
 
         return True
 
@@ -161,31 +165,35 @@ class Tidal(BaseConfig, metaclass=SingletonMeta):
         self.set_option("expiry_time", self.session.expiry_time)
         self.save()
 
-    def get_atmos_session(self) -> tidalapi.Session:
-        """
-        Returns the Atmos session. If it doesn't exist, it creates it
-        by logging in with the Fire TV client ID using the main session's token.
-        """
-        if self.atmos_session is None:
-            # This config uses the Fire TV client ID needed for Dolby Atmos.
-            # It is handled internally and not exposed to the user.
-            atmos_config = tidalapi.Config(item_limit=10000)
-            atmos_config.client_id = "7m7Ap0JC9j1cOM3n"
-            atmos_config.client_secret = "vRAdA108tlvkJpTsGZS8rGZ7xTlbJ0qaZ2K9saEzsgY="
+    @contextmanager
+    def atmos_session_context(self):
+
+        if not self.session.check_login():
+            print("Not logged in. Cannot switch to Atmos context.")
+            raise ConnectionError("Not logged in. Cannot switch to Atmos context.")
+
+        original_client_id = self.session.config.client_id
+        original_client_secret = self.session.config.client_secret
+        original_audio_quality = self.session.audio_quality
+
+        try:
+            self.session.config.client_id = _ATMOS_CLIENT_ID
+            self.session.config.client_secret = _ATMOS_CLIENT_SECRET
+            self.session.audio_quality = _ATMOS_REQUEST_QUALITY
+
+            if not self.login_token(do_pkce=self.is_pkce):
+                raise ConnectionError("Session re-authentication for Atmos failed.")
+
+            yield
 
-            self.atmos_session = tidalapi.Session(atmos_config)
+        finally:
+            self.session.config.client_id = original_client_id
+            self.session.config.client_secret = original_client_secret
+            self.session.audio_quality = original_audio_quality
 
-        # Try to log in using the token from the already logged-in main session.
-        if not self.atmos_session.check_login():
-            self.atmos_session.load_oauth_session(
-                self.data.token_type,
-                self.data.access_token,
-                self.data.refresh_token,
-                self.data.expiry_time,
-                is_pkce=self.is_pkce,
-            )
+            if not self.login_token(do_pkce=self.is_pkce):
+                print("Warning: Restoring the original session context failed. Please restart the application.")
 
-        return self.atmos_session
 
     def login(self, fn_print: Callable) -> bool:
         is_token = self.login_token()
diff --git a/tidal_dl_ng/dialog.py b/tidal_dl_ng/dialog.py
index dcedf2e..0c55e95 100644
--- a/tidal_dl_ng/dialog.py
+++ b/tidal_dl_ng/dialog.py
@@ -176,6 +176,7 @@ class DialogPreferences(QtWidgets.QDialog):
             "lyrics_embed",
             "lyrics_file",
             "video_download",
+            "download_dolby_atmos",
             "download_delay",
             "video_convert_mp4",
             "extract_flac",
diff --git a/tidal_dl_ng/download.py b/tidal_dl_ng/download.py
index a58d48e..6b089eb 100644
--- a/tidal_dl_ng/download.py
+++ b/tidal_dl_ng/download.py
@@ -28,7 +28,15 @@ from requests.exceptions import HTTPError
 from rich.progress import Progress, TaskID
 from tidalapi import Album, Mix, Playlist, Session, Track, UserPlaylist, Video
 from tidalapi.exceptions import TooManyRequests
-from tidalapi.media import AudioExtensions, Codec, Quality, Stream, StreamManifest, VideoExtensions
+from tidalapi.media import (
+    AudioExtensions,
+    AudioMode,
+    Codec,
+    Quality,
+    Stream,
+    StreamManifest,
+    VideoExtensions,
+)
 
 from tidal_dl_ng.config import Settings, Tidal
 from tidal_dl_ng.constants import (
@@ -99,7 +107,7 @@ class Download:
     """Main class for managing downloads, segment merging, file operations, and metadata for TIDAL media."""
 
     settings: Settings
-    tidal: "Tidal"  # Using a forward reference for type hinting
+    tidal: "Tidal" 
     session: Session
     skip_existing: bool = False
     fn_logger: Callable
@@ -111,7 +119,7 @@ class Download:
 
     def __init__(
         self,
-        tidal_obj: "Tidal",
+        tidal_obj: Tidal,
         path_base: str,
         fn_logger: Callable,
         skip_existing: bool = False,
@@ -125,6 +133,7 @@ class Download:
 
         Args:
             session (Session): TIDAL session object.
+            tidal_obj (Tidal): The TIDAL object (from config.py)
             path_base (str): Base path for downloads.
             fn_logger (Callable): Logger function or object.
             skip_existing (bool, optional): Whether to skip existing files. Defaults to False.
@@ -785,31 +794,16 @@ class Download:
 
         if isinstance(media, Track):
             try:
-                # Check which quality the user has selected
-                selected_quality = self.settings.data.quality_audio
-
-                if selected_quality == "DOLBY_ATMOS":
-                    # Dolby Atmos requested: use the special Atmos session.
-                    atmos_session = self.tidal.get_atmos_session()
-                
-                    # Store the original quality of the atmos_session to restore it later
-                    original_quality = atmos_session.audio_quality 
-                
-                    try:
-                        # Temporarily set the session to request what Tidal considers "High" (320k AAC).
-                        # This is the key to getting the Atmos stream with the Fire TV client ID.
-                        atmos_session.audio_quality = Quality.low_320k
-                    
-                        # Get the track object from the correct session
-                        atmos_track = atmos_session.track(media.id)
-                    
-                        # Call get_stream() WITHOUT arguments. It will use the quality from the session object.
+                if (
+                    self.settings.data.download_dolby_atmos
+                    and hasattr(media, "audio_modes")
+                    and AudioMode.dolby_atmos.value in media.audio_modes
+                ):
+                    with self.tidal.atmos_session_context():
+                        atmos_track = self.session.track(media.id)
                         media_stream = atmos_track.get_stream()
-                    finally:
-                        # IMPORTANT: Always restore the original quality on the atmos_session
-                        atmos_session.audio_quality = original_quality
                 else:
-                    # Standard behavior: get_stream uses the quality already set on the main session
+                    self.session.audio_quality = Quality(self.settings.data.quality_audio)
                     media_stream = media.get_stream()
             
                 stream_manifest = media_stream.get_stream_manifest()
diff --git a/tidal_dl_ng/gui.py b/tidal_dl_ng/gui.py
index 3086e47..ea6d238 100644
--- a/tidal_dl_ng/gui.py
+++ b/tidal_dl_ng/gui.py
@@ -93,6 +93,7 @@ import coloredlogs.converter
 from rich.progress import Progress
 from tidalapi import Album, Mix, Playlist, Quality, Track, UserPlaylist, Video
 from tidalapi.artist import Artist
+from tidalapi.media import AudioMode
 from tidalapi.session import SearchTypes
 
 from tidal_dl_ng.config import HandlingApp, Settings, Tidal
@@ -315,12 +316,6 @@ class MainWindow(QtWidgets.QMainWindow, Ui_MainWindow):
         """
         for item in options:
             ui_target.addItem(item.name, item)
-            
-        # Add our custom Dolby Atmos option to the audio quality dropdown
-        if ui_target == self.cb_quality_audio:
-                # Check if it already exists to be safe
-                if ui_target.findText("DOLBY_ATMOS") == -1:
-                    ui_target.addItem("DOLBY_ATMOS", "DOLBY_ATMOS")
 
     def _populate_search_types(self, ui_target: QtWidgets.QComboBox, options: Iterable[Any]) -> None:
         """Populate a combo box with search type options.
@@ -893,8 +888,8 @@ class MainWindow(QtWidgets.QMainWindow, Ui_MainWindow):
         """
 
         final_quality = quality_audio_highest(item)
-        if hasattr(item, "audio_modes") and "DOLBY_ATMOS" in item.audio_modes:
-            final_quality = "Dolby Atmos"
+        if hasattr(item, "audio_modes") and AudioMode.dolby_atmos.value in item.audio_modes:
+            final_quality = f"{final_quality} / Dolby Atmos"
 
         return ResultItem(
             position=idx,
@@ -1201,15 +1196,10 @@ class MainWindow(QtWidgets.QMainWindow, Ui_MainWindow):
             index: The index of the selected quality in the combo box.
         """
         quality_data = self.cb_quality_audio.itemData(index)
-
-        # Handle our custom string for Dolby Atmos
-        if isinstance(quality_data, str) and quality_data == "DOLBY_ATMOS":
-            self.settings.data.quality_audio = "DOLBY_ATMOS"
-        else:
-            self.settings.data.quality_audio = Quality(quality_data)
-
+        
+        self.settings.data.quality_audio = Quality(quality_data)
         self.settings.save()
-
+        
         if self.tidal:
             self.tidal.settings_apply()
 
diff --git a/tidal_dl_ng/model/cfg.py b/tidal_dl_ng/model/cfg.py
index b3bc202..ab72cf9 100644
--- a/tidal_dl_ng/model/cfg.py
+++ b/tidal_dl_ng/model/cfg.py
@@ -23,6 +23,7 @@ class Settings:
     download_base_path: str = "~/download"
     quality_audio: Quality = Quality.low_320k
     quality_video: QualityVideo = QualityVideo.P480
+    download_dolby_atmos: bool = False
     format_album: str = (
         "Albums/{album_artist} - {album_title}{album_explicit}/{track_volume_num_optional}"
         "{album_track_num}. {artist_name} - {track_title}{album_explicit}"
@@ -79,6 +80,7 @@ class HelpSettings:
         '"HI_RES_LOSSLESS" (up to 24 Bit, 192 kHz)'
     )
     quality_video: str = 'Desired video download quality: "360", "480", "720", "1080"'
+    download_dolby_atmos: str = "Download Dolby Atmos audio streams if available."
     # TODO: Describe possible variables.
     format_album: str = "Where to download albums and how to name the items."
     format_playlist: str = "Where to download playlists and how to name the items."
diff --git a/tidal_dl_ng/ui/dialog_settings.py b/tidal_dl_ng/ui/dialog_settings.py
index 9ec2738..f05adef 100644
--- a/tidal_dl_ng/ui/dialog_settings.py
+++ b/tidal_dl_ng/ui/dialog_settings.py
@@ -198,6 +198,23 @@ class Ui_DialogSettings:
 
         self.lv_flags.addLayout(self.horizontalLayout_12)
 
+        self.lh_flags_atmos = QHBoxLayout()
+        self.lh_flags_atmos.setObjectName("lh_flags_atmos")
+        self.lv_flag_download_dolby_atmos = QVBoxLayout()
+        self.lv_flag_download_dolby_atmos.setObjectName("lv_flag_download_dolby_atmos")
+        self.cb_download_dolby_atmos = QCheckBox(self.gb_flags)
+        self.cb_download_dolby_atmos.setObjectName("cb_download_dolby_atmos")
+
+        self.lv_flag_download_dolby_atmos.addWidget(self.cb_download_dolby_atmos)
+
+        self.lh_flags_atmos.addLayout(self.lv_flag_download_dolby_atmos)
+
+        self.lv_flag_empty_atmos = QVBoxLayout()
+        self.lv_flag_empty_atmos.setObjectName("lv_flag_empty_atmos")
+        self.lh_flags_atmos.addLayout(self.lv_flag_empty_atmos)
+
+        self.lv_flags.addLayout(self.lh_flags_atmos)
+
         self.lv_main.addWidget(self.gb_flags)
 
         self.gb_choices = QGroupBox(DialogSettings)