swingmusic-extended/swingmusic/lib/searchlib.py

"""
This library contains all the functions related to the search functionality.
"""

from typing import Any, Generator, List, TypeVar

from rapidfuzz import process, utils, fuzz
from unidecode import unidecode

from swingmusic import models
from swingmusic.config import UserConfig

# from swingmusic.db.libdata import AlbumTable, ArtistTable, TrackTable

# from swingmusic.db.sqlite.favorite import SQLiteFavoriteMethods as favdb
from swingmusic.models.album import Album
from swingmusic.models.artist import Artist
from swingmusic.models.enums import FavType
from swingmusic.models.playlist import Playlist
from swingmusic.models.track import Track
from swingmusic.serializers.album import serialize_for_card as serialize_album
from swingmusic.serializers.album import serialize_for_card_many as serialize_albums
from swingmusic.serializers.artist import serialize_for_card, serialize_for_cards
from swingmusic.serializers.track import serialize_track, serialize_tracks

from swingmusic.store.albums import AlbumStore
from swingmusic.store.artists import ArtistStore
from swingmusic.store.tracks import TrackStore

from swingmusic.utils.remove_duplicates import remove_duplicates

# ratio = fuzz.ratio
# wratio = fuzz.WRatio


class Cutoff:
    """
    Holds all the default cutoff values.
    """

    tracks: int = 50
    albums: int = 50
    artists: int = 50
    playlists: int = 50


class Limit:
    """
    Holds all the default limit values.
    """

    tracks: int = 150
    albums: int = 150
    artists: int = 150
    playlists: int = 150


class SearchTracks:
    def __init__(self, query: str) -> None:
        self.query = query
        self.tracks = TrackStore.get_flat_list()

    def __call__(self, limit: int = Limit.tracks) -> List[models.Track]:
        """
        Gets all songs with a given title.
        """

        track_titles = [unidecode(track.title).lower() for track in self.tracks]
        results = process.extract(
            self.query,
            track_titles,
            score_cutoff=Cutoff.tracks,
            limit=limit,
            processor=utils.default_process,
            scorer=fuzz.WRatio,
        )

        tracks: list[Track] = []

        for item in results:
            track = self.tracks[item[2]]
            track._score = item[1]
            tracks.append(track)

        return remove_duplicates(tracks)


class SearchArtists:
    def __init__(self, query: str) -> None:
        self.query = query
        self.artists = ArtistStore.get_flat_list()

    def __call__(self, limit: int = Limit.artists):
        """
        Gets all artists with a given name.
        """
        choices = [unidecode(a.name).lower() for a in self.artists]

        results = process.extract(
            self.query,
            choices,
            score_cutoff=Cutoff.artists,
            limit=limit,
            processor=utils.default_process,
            scorer=fuzz.WRatio,
        )

        artists: list[Artist] = []

        for item in results:
            artist = self.artists[item[2]]
            artist._score = item[1]
            artists.append(artist)

        return artists


class SearchAlbums:
    def __init__(self, query: str) -> None:
        self.query = query
        self.albums = AlbumStore.get_flat_list()

    def __call__(self, limit: int = Limit.albums):
        """
        Gets all albums with a given title.
        """

        choices = [unidecode(a.title).lower() for a in self.albums]

        results = process.extract(
            self.query,
            choices,
            score_cutoff=Cutoff.albums,
            limit=limit,
            processor=utils.default_process,
            scorer=fuzz.token_sort_ratio,
        )

        albums: list[Album] = []

        for item in results:
            album = self.albums[item[2]]
            album._score = item[1]
            albums.append(album)

        return albums


class SearchPlaylists:
    def __init__(self, playlists: List[models.Playlist], query: str) -> None:
        self.playlists = playlists
        self.query = query

    def __call__(self, limit: int = Limit.playlists):
        choices = [p.name for p in self.playlists]
        results = process.extract(
            self.query,
            choices,
            score_cutoff=Cutoff.playlists,
            limit=limit,
            processor=utils.default_process,
            scorer=fuzz.WRatio,
        )

        playlists: list[Playlist] = []

        for item in results:
            playlist = self.playlists[item[2]]
            playlist._score = item[1]
            playlists.append(playlist)

        return playlists


_type = models.Track | models.Album | models.Artist


def get_titles(items: list[_type]):
    for item in items:
        if isinstance(item, models.Track):
            text = item.og_title
        elif isinstance(item, models.Album):
            text = item.title
        elif isinstance(item, models.Artist):
            text = item.name
        else:
            text = None

        yield text


class TopResults:
    """
    Joins all tracks, albums and artists
    then fuzzy searches them as a single unit.
    """

    @staticmethod
    def collect_all():
        all_items: list[_type] = []

        all_items.extend(ArtistStore.get_flat_list())
        all_items.extend(TrackStore.get_flat_list())
        all_items.extend(AlbumStore.get_flat_list())

        return all_items, get_titles(all_items)

    @staticmethod
    def get_track_items(item: Track | Album | Artist, limit=5):
        tracks: list[Track] = []

        # INFO: If the item is a track, return empty list
        # to be filled by the results from the top search
        if isinstance(item, Track):
            return tracks

        # INFO: If the item is an album, get the tracks from the album
        if isinstance(item, Album):
            tracks = TrackStore.get_tracks_by_albumhash(item.albumhash)[:limit]
            tracks.sort(key=lambda x: x.playduration, reverse=True)
            return tracks

        # INFO: If the item is an artist, get the tracks from the artist
        if isinstance(item, Artist):
            tracks = TrackStore.get_tracks_by_artisthash(item.artisthash)[:limit]
            tracks.sort(key=lambda x: x.playduration, reverse=True)

        return tracks

    @staticmethod
    def get_album_items(item: Track | Album | Artist, limit=6):
        albums: list[Album] = []

        # INFO: If the item is a track or album, search for albums
        if isinstance(item, Track) or isinstance(item, Album):
            return albums

        # INFO: If the item is an artist, get the albums from the artist
        if isinstance(item, Artist):
            albums = AlbumStore.get_albums_by_artisthash(item.artisthash)[:limit]

        return albums

    @staticmethod
    def search(
        query: str,
        limit: int = None,
        albums_only=False,
        tracks_only=False,
    ):
        tracks_limit = Limit.tracks if tracks_only else 4
        albums_limit = Limit.albums if albums_only else limit
        artists_limit = limit

        # INFO: Individually search all stores as each type has a different scorer
        tracks = SearchTracks(query)(limit=tracks_limit) if not albums_only else []
        albums = SearchAlbums(query)(limit=albums_limit)
        artists = SearchArtists(query)(limit=artists_limit)

        # INFO: Combine all results and sort them by score
        all_results = artists + tracks + albums
        all_results = sorted(all_results, key=lambda x: int(x._score), reverse=True)

        # INFO: Get the top result
        top_result = all_results[0]
        top_tracks = []

        if not albums_only:
            top_tracks = TopResults.get_track_items(top_result, limit=tracks_limit)

            # INFO: If there are not enough tracks, fill with search results
            if len(top_tracks) < tracks_limit:
                found_tracks_set = {track.trackhash for track in top_tracks}

                for track in tracks:
                    if track.trackhash not in found_tracks_set:
                        top_tracks.append(track)

                    if len(top_tracks) >= tracks_limit:
                        break

            top_tracks = serialize_tracks(top_tracks)

            if tracks_only:
                return top_tracks

        top_albums = TopResults.get_album_items(top_result, limit=albums_limit)

        # INFO: If there are not enough albums, fill with search results
        if len(top_albums) < albums_limit:
            found_albums_set = {album.albumhash for album in top_albums}

            for album in albums:
                if album.albumhash not in found_albums_set:
                    top_albums.append(album)

                    if len(top_albums) >= albums_limit:
                        break

        top_albums = serialize_albums(top_albums)

        if albums_only:
            return top_albums

        artists = serialize_for_cards(artists)

        if isinstance(top_result, Track):
            top_result = serialize_track(top_result)
            top_result["type"] = "track"

        if isinstance(top_result, Album):
            top_result = serialize_album(top_result)
            top_result["type"] = "album"

        if isinstance(top_result, Artist):
            top_result = serialize_for_card(
                top_result, include={"albumcount", "trackcount"}
            )
            top_result["type"] = "artist"

        return {
            "top_result": top_result,
            "tracks": top_tracks,
            "artists": artists,
            "albums": top_albums,
        }