tidal-dl/tidal_dl_ng/gui/search.py

# tidal_dl_ng/search.py

from typing import TYPE_CHECKING, Any

from tidalapi import Album, Artist, Mix, Playlist, Track, Video
from tidalapi import VideoQuality as QualityVideo
from tidalapi.media import AudioMode, Quality
from tidalapi.playlist import Folder
from tidalapi.session import SearchTypes

from tidal_dl_ng.constants import QueueDownloadStatus
from tidal_dl_ng.helper.tidal import (
    get_tidal_media_id,
    get_tidal_media_type,
    instantiate_media,
    name_builder_artist,
    name_builder_title,
    quality_audio_highest,
    search_results_all,
    url_ending_clean,
)
from tidal_dl_ng.logger import logger_gui
from tidal_dl_ng.model.gui_data import QueueDownloadItem, ResultItem

if TYPE_CHECKING:
    from tidal_dl_ng.gui import MainWindow


class GuiSearchManager:
    """Manages the search GUI and logic."""

    def __init__(self, main_window: "MainWindow"):
        """Initialize the search manager."""
        self.main_window: "MainWindow" = main_window

    def search_populate_results(self, query: str, type_media: Any) -> None:
        """Populate the results tree with search results."""
        results = self.search(query, [type_media])
        self.main_window.populate_tree_results(results)

    def search(self, query: str, types_media: list[Any]) -> list[ResultItem]:
        """Perform a search and return a list of ResultItems.

        Args:
            query (str): The search query.
            types_media (list[Any]): The types of media to search for.

        Returns:
            list[ResultItem]: The search results.
        """
        query_clean: str = query.strip()

        # If a direct link was searched for, skip search and create the object from the link directly.
        if "http" in query_clean:
            query_clean: str = url_ending_clean(query_clean)
            media_type = get_tidal_media_type(query_clean)
            item_id = get_tidal_media_id(query_clean)

            try:
                media = instantiate_media(self.main_window.tidal.session, media_type, item_id)
            except:
                logger_gui.error(f"Media not found (ID: {item_id}). Maybe it is not available anymore.")

                media = None

            result_search = {"direct": [media]}
        else:
            result_search: dict[str, list[SearchTypes]] = search_results_all(
                session=self.main_window.tidal.session, needle=query_clean, types_media=types_media
            )

        result: list[ResultItem] = []

        for _media_type, l_media in result_search.items():
            if isinstance(l_media, list):
                result = result + self.search_result_to_model(l_media)

        return result

    def search_result_to_model(self, items: list[SearchTypes]) -> list[ResultItem]:
        """Convert search results to ResultItem models.

        Args:
            items (list[SearchTypes]): List of search result items.

        Returns:
            list[ResultItem]: List of ResultItem models.
        """
        result: list[ResultItem] = []

        for idx, item in enumerate(items):
            result_item = self._to_result_item(idx, item)

            if result_item is not None:
                result.append(result_item)

        return result

    def _to_result_item(self, idx: int, item) -> ResultItem | None:
        """Helper to convert a single item to ResultItem, or None if not valid.

        Args:
            idx (int): Index of the item.
            item: The item to convert.

        Returns:
            ResultItem | None: The converted ResultItem or None if not valid.
        """
        if not item or (hasattr(item, "available") and not item.available):
            return None

        # Prepare common data
        explicit = " 🅴" if isinstance(item, Track | Video | Album) and item.explicit else ""
        date_user_added = (
            item.user_date_added.strftime("%Y-%m-%d_%H:%M") if getattr(item, "user_date_added", None) else ""
        )
        date_release = self._get_date_release(item)

        # Map item types to their conversion methods
        type_handlers = {
            Track: lambda: self._result_item_from_track(idx, item, explicit, date_user_added, date_release),
            Video: lambda: self._result_item_from_video(idx, item, explicit, date_user_added, date_release),
            Playlist: lambda: self._result_item_from_playlist(idx, item, date_user_added, date_release),
            Album: lambda: self._result_item_from_album(idx, item, explicit, date_user_added, date_release),
            Mix: lambda: self._result_item_from_mix(idx, item, date_user_added, date_release),
            Artist: lambda: self._result_item_from_artist(idx, item, date_user_added, date_release),
            Folder: lambda: self._result_item_from_folder(idx, item, date_user_added),
        }

        # Find and execute the appropriate handler
        for item_type, handler in type_handlers.items():
            if isinstance(item, item_type):
                return handler()

        return None

    def _get_date_release(self, item) -> str:
        """Get the release date string for an item.

        Args:
            item: The item to extract the release date from.

        Returns:
            str: The formatted release date or empty string.
        """
        if hasattr(item, "album") and item.album and getattr(item.album, "release_date", None):
            return item.album.release_date.strftime("%Y-%m-%d_%H:%M")

        if hasattr(item, "release_date") and item.release_date:
            return item.release_date.strftime("%Y-%m-%d_%H:%M")

        return ""

    def _result_item_from_track(
        self, idx: int, item, explicit: str, date_user_added: str, date_release: str
    ) -> ResultItem:
        """Create a ResultItem from a Track.

        Args:
            idx (int): Index of the item.
            item: The Track item.
            explicit (str): Explicit tag.
            date_user_added (str): Date user added.
            date_release (str): Release date.

        Returns:
            ResultItem: The constructed ResultItem.
        """

        final_quality = quality_audio_highest(item)
        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,
            artist=name_builder_artist(item),
            title=f"{name_builder_title(item)}{explicit}",
            album=item.album.name,
            duration_sec=item.duration,
            obj=item,
            quality=final_quality,
            explicit=bool(item.explicit),
            date_user_added=date_user_added,
            date_release=date_release,
        )

    def _result_item_from_video(
        self, idx: int, item, explicit: str, date_user_added: str, date_release: str
    ) -> ResultItem:
        """Create a ResultItem from a Video.

        Args:
            idx (int): Index of the item.
            item: The Video item.
            explicit (str): Explicit tag.
            date_user_added (str): Date user added.
            date_release (str): Release date.

        Returns:
            ResultItem: The constructed ResultItem.
        """
        return ResultItem(
            position=idx,
            artist=name_builder_artist(item),
            title=f"{name_builder_title(item)}{explicit}",
            album=item.album.name if item.album else "",
            duration_sec=item.duration,
            obj=item,
            quality=item.video_quality,
            explicit=bool(item.explicit),
            date_user_added=date_user_added,
            date_release=date_release,
        )

    def _result_item_from_playlist(self, idx: int, item, date_user_added: str, date_release: str) -> ResultItem:
        """Create a ResultItem from a Playlist.

        Args:
            idx (int): Index of the item.
            item: The Playlist item.
            date_user_added (str): Date user added.
            date_release (str): Release date.

        Returns:
            ResultItem: The constructed ResultItem.
        """
        return ResultItem(
            position=idx,
            artist=", ".join(artist.name for artist in item.promoted_artists) if item.promoted_artists else "",
            title=item.name,
            album="",
            duration_sec=item.duration,
            obj=item,
            quality="",
            explicit=False,
            date_user_added=date_user_added,
            date_release=date_release,
        )

    def _result_item_from_album(
        self, idx: int, item, explicit: str, date_user_added: str, date_release: str
    ) -> ResultItem:
        """Create a ResultItem from an Album.

        Args:
            idx (int): Index of the item.
            item: The Album item.
            explicit (str): Explicit tag.
            date_user_added (str): Date user added.
            date_release (str): Release date.

        Returns:
            ResultItem: The constructed ResultItem.
        """
        return ResultItem(
            position=idx,
            artist=name_builder_artist(item),
            title="",
            album=f"{item.name}{explicit}",
            duration_sec=item.duration,
            obj=item,
            quality=quality_audio_highest(item),
            explicit=bool(item.explicit),
            date_user_added=date_user_added,
            date_release=date_release,
        )

    def _result_item_from_mix(self, idx: int, item, date_user_added: str, date_release: str) -> ResultItem:
        """Create a ResultItem from a Mix.

        Args:
            idx (int): Index of the item.
            item: The Mix item.
            date_user_added (str): Date user added.
            date_release (str): Release date.

        Returns:
            ResultItem: The constructed ResultItem.
        """
        return ResultItem(
            position=idx,
            artist=item.sub_title,
            title=item.title,
            album="",
            duration_sec=-1,  # TODO: Calculate total duration.
            obj=item,
            quality="",
            explicit=False,
            date_user_added=date_user_added,
            date_release=date_release,
        )

    def _result_item_from_artist(self, idx: int, item, date_user_added: str, date_release: str) -> ResultItem:
        """Create a ResultItem from an Artist.

        Args:
            idx (int): Index of the item.
            item: The Artist item.
            date_user_added (str): Date user added.
            date_release (str): Release date.

        Returns:
            ResultItem: The constructed ResultItem.
        """
        return ResultItem(
            position=idx,
            artist=item.name,
            title="",
            album="",
            duration_sec=-1,
            obj=item,
            quality="",
            explicit=False,
            date_user_added=date_user_added,
            date_release=date_release,
        )

    def _result_item_from_folder(self, idx: int, item: Folder, date_user_added: str) -> ResultItem:
        """Create a ResultItem from a Folder.

        Args:
            idx (int): Index of the item.
            item (Folder): The Folder item.
            date_user_added (str): Date user added.

        Returns:
            ResultItem: The constructed ResultItem.
        """
        total_items: int = item.total_number_of_items if hasattr(item, "total_number_of_items") else 0
        return ResultItem(
            position=idx,
            artist="",
            title=f"📁 {item.name} ({total_items} items)",
            album="",
            duration_sec=-1,
            obj=item,
            quality="",
            explicit=False,
            date_user_added=date_user_added,
            date_release="",
        )

    def media_to_queue_download_model(
        self, media: Artist | Track | Video | Album | Playlist | Mix
    ) -> QueueDownloadItem | bool:
        """Convert a media object to a QueueDownloadItem for the download queue.

        Args:
            media (Artist | Track | Video | Album | Playlist | Mix): The media object.

        Returns:
            QueueDownloadItem | bool: The queue item or False if not available.
        """
        result: QueueDownloadItem | False
        name: str = ""
        quality_audio: Quality = self.main_window.settings.data.quality_audio
        quality_video: QualityVideo = self.main_window.settings.data.quality_video
        explicit: str = ""

        # Check if item is available on TIDAL.
        # Note: Some albums have available=None, which should be treated as available
        if hasattr(media, "available") and media.available is False:
            return False

        # Set "Explicit" tag
        if isinstance(media, Track | Video | Album):
            explicit = " 🅴" if media.explicit else ""

        # Build name and set quality
        if isinstance(media, Track | Video):
            name = f"{name_builder_artist(media)} - {name_builder_title(media)}{explicit}"
        elif isinstance(media, Playlist | Artist):
            name = media.name
        elif isinstance(media, Album):
            name = f"{name_builder_artist(media)} - {media.name}{explicit}"
        elif isinstance(media, Mix):
            name = media.title

        # Determine actual quality.
        if isinstance(media, Track | Album):
            quality_highest: Quality = quality_audio_highest(media)

            if (
                self.main_window.settings.data.quality_audio == quality_highest
                or self.main_window.settings.data.quality_audio == Quality.hi_res_lossless
            ):
                quality_audio = quality_highest

        if name:
            result = QueueDownloadItem(
                name=name,
                quality_audio=quality_audio,
                quality_video=quality_video,
                type_media=type(media).__name__,
                status=QueueDownloadStatus.Waiting,
                obj=media,
            )
        else:
            result = False

        return result