beets/beetsplug/tidal/__init__.py

from __future__ import annotations

import itertools
import os
import re
import time
from functools import cached_property
from typing import TYPE_CHECKING, ClassVar, Literal, overload

import confuse

from beets import ui
from beets.autotag import AlbumInfo, TrackInfo
from beets.dbcore import types
from beets.exceptions import UserError
from beets.logging import getLogger
from beets.metadata_plugins import MetadataSourcePlugin

from .api import TidalAPI

if TYPE_CHECKING:
    import optparse
    from collections.abc import Callable, Iterable, Sequence

    from beets.autotag import Info
    from beets.dbcore.db import Results
    from beets.library import Library
    from beets.library.models import Album, Item

    from .api_types import (
        AlbumAttributes,
        MediaAttributes,
        ResourceIdentifier,
        TidalAlbum,
        TidalArtist,
        TidalArtwork,
        TidalTrack,
    )


log = getLogger("beets.tidal")


class TidalPlugin(MetadataSourcePlugin):
    item_types: ClassVar[dict[str, types.Type]] = {
        "tidal_track_id": types.STRING,
        "tidal_artist_id": types.STRING,
        "tidal_track_popularity": types.INTEGER,
        "tidal_updated": types.DATE,
    }

    album_types: ClassVar[dict[str, types.Type]] = {
        "tidal_album_id": types.STRING,
        "tidal_artist_id": types.STRING,
        "tidal_album_popularity": types.INTEGER,
        "tidal_updated": types.DATE,
    }

    def __init__(self) -> None:
        super().__init__()

        self.config.add(
            {"client_id": "mcjmpl1bPATJXcBT", "tokenfile": "tidal_token.json"}
        )
        self.config["client_id"].redact = True

        # We need to be authenticated if plugin is used to fetch metadata
        # otherwise the import cannot run.
        self.register_listener("import_begin", self.require_authentication)

    @cached_property
    def api(self) -> TidalAPI:
        return TidalAPI(
            client_id=self.config["client_id"].as_str(),
            token_path=self._tokenfile(),
        )

    def _tokenfile(self) -> str:
        """Return the configured path to the token file in the app directory."""
        return self.config["tokenfile"].get(confuse.Filename(in_app_dir=True))

    def require_authentication(self) -> None:
        if not os.path.isfile(self._tokenfile()):
            raise UserError(
                "Please login to TIDAL"
                " using `beet tidal --auth` or disable tidal plugin"
            )

    def album_for_id(self, album_id: str) -> AlbumInfo | None:
        if not (tidal_id := self._extract_id(album_id)):
            return None

        if album := list(self.search_albums_by_ids(tidal_ids=[tidal_id])):
            return album[0]

        log.warning("Could not find album:{0}", tidal_id)
        return None

    def albums_for_ids(self, ids: Iterable[str]) -> Iterable[AlbumInfo | None]:
        yield from self.search_albums_by_ids(ids=ids)

    def track_for_id(self, track_id: str) -> TrackInfo | None:
        if not (tidal_id := self._extract_id(track_id)):
            return None

        if track := list(self.search_tracks_by_ids(tidal_ids=[tidal_id])):
            return track[0]

        log.warning("Could not find track:{0}", tidal_id)
        return None

    def tracks_for_ids(self, ids: Iterable[str]) -> Iterable[TrackInfo | None]:
        yield from self.search_tracks_by_ids(ids=ids)

    def candidates(
        self, items: Sequence[Item], artist: str, album: str, va_likely: bool
    ) -> Iterable[AlbumInfo]:
        candidates: list[AlbumInfo] = []
        # Tidal allows to lookup via isrc and barcode (nice!)
        # We just return early here as a lookup via isrc should
        # return a 100% match
        barcodes: list[str] = list(
            filter(None, set(i.get("barcode") for i in items))
        )
        if barcodes and (
            candidates := list(
                filter(None, self.search_albums_by_ids(barcode_ids=barcodes))
            )
        ):
            return candidates

        for query in self._album_queries(items):
            candidates += self.search_albums_by_query(query)

        log.debug("Found {0} candidates", len(candidates))
        return candidates

    def item_candidates(
        self, item: Item, artist: str, title: str
    ) -> Iterable[TrackInfo]:
        candidates: list[TrackInfo] = []
        # Tidal allows to lookup via isrc and barcode (nice!)
        # We just return early here as a lookup via isrc should
        # return a 100% match
        if isrc := item.get("isrc"):
            if candidates := list(
                filter(None, self.search_tracks_by_ids(isrcs=[isrc]))
            ):
                return candidates

        for query in self._item_queries(item):
            candidates += self.search_tracks_by_query(query)

        log.debug("Found {0} candidates", len(candidates))
        return candidates

    @staticmethod
    def _item_queries(item: Item) -> Iterable[str]:
        """Search queries for items."""
        yield item.title

        if item.artist:
            yield f"{item.artist} {item.title}"

    @staticmethod
    def _album_queries(items: Sequence[Item]) -> Iterable[str]:
        """Search queries for albums."""

        album_names = set(i.album for i in items)
        artist_names = set(i.artist for i in items)

        for album, artist in itertools.product(album_names, artist_names):
            yield f"{artist} {album}"

    def search_tracks_by_query(self, query: str) -> Iterable[TrackInfo]:
        """Search for tracks given a string query."""
        search_doc = self.api.search_results(query, include=["tracks.artists"])
        track_by_id: dict[str, TidalTrack] = {
            item["id"]: item
            for item in search_doc.get("included", [])
            if item["type"] == "tracks"
        }
        artist_by_id: dict[str, TidalArtist] = {
            item["id"]: item
            for item in search_doc.get("included", [])
            if item["type"] == "artists"
        }
        for track_rel in search_doc["data"]["relationships"]["tracks"]["data"]:
            if track := track_by_id.get(track_rel["id"]):
                yield self._get_track_info(track, artist_by_id=artist_by_id)
            else:
                log.warning(
                    "Track with id {0} not found in lookup", track_rel["id"]
                )

    def search_albums_by_query(self, query: str) -> Iterable[AlbumInfo]:
        """Search for album given a string query."""
        search_doc = self.api.search_results(
            query,
            include=["albums"],
            # include="albums.items.artists" <- not supported
            # This is a bit inconvenient, but we fetch the items and artists
            # for all albums separately.
        )
        album_ids = [
            album_rel["id"]
            for album_rel in search_doc["data"]["relationships"]["albums"][
                "data"
            ]
        ]
        yield from filter(None, self.search_albums_by_ids(tidal_ids=album_ids))

    @overload
    def search_tracks_by_ids(
        self, *, ids: Iterable[str]
    ) -> Iterable[TrackInfo | None]: ...

    @overload
    def search_tracks_by_ids(
        self, *, tidal_ids: Iterable[str]
    ) -> Iterable[TrackInfo | None]: ...

    @overload
    def search_tracks_by_ids(
        self, *, isrcs: Iterable[str]
    ) -> Iterable[TrackInfo | None]: ...

    def search_tracks_by_ids(
        self,
        ids: Iterable[str] | None = None,
        tidal_ids: Iterable[str] | None = None,
        isrcs: Iterable[str] | None = None,
    ) -> Iterable[TrackInfo | None]:
        _ids: list[str | None] = list(tidal_ids or [])
        isrcs = list(isrcs or [])
        if ids:
            _ids = list(map(self._extract_id, ids))

        tracks_doc = self.api.get_tracks(
            ids=list(filter(None, _ids)), isrcs=isrcs, include=["artists"]
        )
        track_by_id: dict[str, TidalTrack] = {
            item["id"]: item
            for item in tracks_doc.get("data", [])
            if item["type"] == "tracks"
        }
        artist_by_id: dict[str, TidalArtist] = {
            item["id"]: item
            for item in tracks_doc.get("included", [])
            if item["type"] == "artists"
        }

        for _id in _ids:
            if _id is not None and (track := track_by_id.get(_id)):
                yield self._get_track_info(track, artist_by_id=artist_by_id)
            else:
                yield None

        if isrcs:
            isrc_to_track: dict[str, TidalTrack] = {
                t["attributes"]["isrc"]: t for t in track_by_id.values()
            }

            for isrc in isrcs:
                if track := isrc_to_track.get(isrc):
                    yield self._get_track_info(track, artist_by_id=artist_by_id)
                else:
                    yield None

    @overload
    def search_albums_by_ids(
        self, *, ids: Iterable[str]
    ) -> Iterable[AlbumInfo | None]: ...

    @overload
    def search_albums_by_ids(
        self, *, tidal_ids: Iterable[str]
    ) -> Iterable[AlbumInfo | None]: ...

    @overload
    def search_albums_by_ids(
        self, *, barcode_ids: Iterable[str]
    ) -> Iterable[AlbumInfo | None]: ...

    def search_albums_by_ids(
        self,
        ids: Iterable[str] | None = None,
        tidal_ids: Iterable[str] | None = None,
        barcode_ids: Iterable[str] | None = None,
    ) -> Iterable[AlbumInfo | None]:
        _ids: list[str | None] = list(tidal_ids or [])
        barcode_ids = list(barcode_ids or [])
        if ids:
            _ids = list(map(self._extract_id, ids))

        albums_doc = self.api.get_albums(
            ids=list(filter(None, _ids)),
            barcode_ids=barcode_ids,
            include=["items.artists", "artists", "coverArt"],
        )
        album_by_id: dict[str, TidalAlbum] = {
            item["id"]: item
            for item in albums_doc.get("data", [])
            if item["type"] == "albums"
        }
        track_by_id: dict[str, TidalTrack] = {
            item["id"]: item
            for item in albums_doc.get("included", [])
            if item["type"] == "tracks"
        }
        artist_by_id: dict[str, TidalArtist] = {
            item["id"]: item
            for item in albums_doc.get("included", [])
            if item["type"] == "artists"
        }
        artwork_by_id: dict[str, TidalArtwork] = {
            item["id"]: item
            for item in albums_doc.get("included", [])
            if item["type"] == "artworks"
        }

        for _id in _ids:
            if _id is not None and (album := album_by_id.get(_id)):
                yield self._get_album_info(
                    album,
                    track_by_id=track_by_id,
                    artist_by_id=artist_by_id,
                    artwork_by_id=artwork_by_id,
                )
            else:
                yield None

        if barcode_ids:
            barcode_to_album: dict[str, TidalAlbum] = {
                a["attributes"]["barcodeId"]: a for a in album_by_id.values()
            }

            for barcode in barcode_ids:
                if album := barcode_to_album.get(barcode):
                    yield self._get_album_info(
                        album,
                        track_by_id=track_by_id,
                        artist_by_id=artist_by_id,
                        artwork_by_id=artwork_by_id,
                    )
                else:
                    yield None

    def _get_album_info(
        self,
        album: TidalAlbum,
        track_by_id: dict[str, TidalTrack],
        artist_by_id: dict[str, TidalArtist],
        artwork_by_id: dict[str, TidalArtwork],
    ) -> AlbumInfo:
        track_infos: list[TrackInfo] = []
        for i, track_rel in enumerate(
            album["relationships"]["items"]["data"], start=1
        ):
            if track := track_by_id.get(track_rel["id"]):
                track_info = self._get_track_info(track, artist_by_id)
                track_info.index = i
                track_infos.append(track_info)

        artist_names, artist_ids = self._parse_artists(
            album["relationships"]["artists"]["data"], artist_by_id
        )
        date_parts = self._parse_release_date(album["attributes"])
        return AlbumInfo(
            # Identifier
            data_source=self.data_source,
            album_id=album["id"],
            artists_ids=artist_ids,
            data_url=self._parse_data_url(album["attributes"]),
            barcode=album["attributes"]["barcodeId"],
            cover_art_url=self._parse_artwork_url(album, artwork_by_id),
            # Meta
            album=self._parse_title(album["attributes"]),
            tracks=track_infos,
            artist=", ".join(artist_names),
            artists=artist_names,
            duration=self._duration_to_seconds(album["attributes"]["duration"]),
            albumtypes=[album["attributes"]["albumType"].lower()],
            label=self._parse_label(album["attributes"]),
            year=date_parts[0] if date_parts else None,
            month=date_parts[1] if date_parts else None,
            day=date_parts[2] if date_parts else None,
            # Flexattrs
            tidal_album_id=album["id"],
            tidal_artist_id=artist_ids[0] if artist_ids else None,
            tidal_album_popularity=self._parse_popularity(album["attributes"]),
            tidal_updated=time.time(),
        )

    @staticmethod
    def _parse_artwork_url(
        album: TidalAlbum, artwork_by_id: dict[str, TidalArtwork]
    ) -> str | None:
        cover_rel = album["relationships"].get("coverArt")
        if cover_rel is None:
            return None
        for rel_data in cover_rel["data"]:
            if (
                rel_data["type"] == "artworks"
                and (artwork := artwork_by_id.get(rel_data["id"]))
                and (files := artwork["attributes"]["files"])
            ):
                return files[0]["href"]
        return None

    def _get_track_info(
        self, track: TidalTrack, artist_by_id: dict[str, TidalArtist]
    ) -> TrackInfo:
        artist_names, artist_ids = self._parse_artists(
            track["relationships"]["artists"]["data"], artist_by_id
        )

        return TrackInfo(
            # Identifier
            data_source=self.data_source,
            track_id=track["id"],
            artists_ids=artist_ids,
            data_url=self._parse_data_url(track["attributes"]),
            # Meta
            title=self._parse_title(track["attributes"]),
            isrc=track["attributes"]["isrc"],
            artist=", ".join(artist_names),
            artists=artist_names,
            duration=self._duration_to_seconds(track["attributes"]["duration"]),
            label=self._parse_label(track["attributes"]),
            # Flexattrs
            tidal_track_id=track["id"],
            tidal_artist_id=artist_ids[0] if artist_ids else None,
            tidal_track_popularity=self._parse_popularity(track["attributes"]),
            tidal_updated=time.time(),
        )

    @staticmethod
    def _parse_artists(
        artist_relationships: list[ResourceIdentifier],
        artist_by_id: dict[str, TidalArtist],
    ) -> tuple[list[str], list[str]]:
        """Extract artists from a relationship.

        Artists are sorted in the track/album response relationship but not in the
        track/album responses included items.
        """
        artist_names = []
        artist_ids = []
        for artist_rel in artist_relationships:
            if artist := artist_by_id.get(artist_rel["id"]):
                artist_ids.append(artist["id"])
                artist_names.append(artist["attributes"]["name"])
            else:
                log.warning(
                    "Artist with id {0} not found in lookup", artist_rel["id"]
                )

        return artist_names, artist_ids

    @staticmethod
    def _parse_title(attributes: MediaAttributes) -> str:
        """
        Tidal UIs append the version string at the end of the title. We do the same here
        by formatting it as ``"{title} ({version})"`` to stay consistent.
        """
        if version := attributes.get("version"):
            return f"{attributes['title']} ({version})"
        return attributes["title"]

    @staticmethod
    def _parse_data_url(attributes: MediaAttributes) -> str | None:
        if external_links := attributes.get("externalLinks"):
            return external_links[0].get("href")
        return None

    @staticmethod
    def _duration_to_seconds(duration: str) -> int | None:
        """Convert ISO 8601 duration to seconds. E.g. 'PT15M2S' -> 902."""
        match = ISO_8601_RE.match(duration)
        if not match:
            log.warning("Invalid ISO 8601 duration: {0}", duration)
            return None
        parts = {k: int(v) if v else 0 for k, v in match.groupdict().items()}
        return parts["seconds"] + parts["minutes"] * 60 + parts["hours"] * 3600

    @staticmethod
    def _parse_label(attributes: MediaAttributes) -> str | None:
        if copyright_ := attributes.get("copyright"):
            return copyright_["text"]
        return None

    @staticmethod
    def _parse_release_date(
        attributes: AlbumAttributes,
    ) -> tuple[int, int, int] | None:
        """Returns year, month, day from iso YYYY-MM-DD"""

        if (
            (release_date := attributes.get("releaseDate"))
            and (parts := release_date.split("-"))
            and len(parts) == 3
        ):
            return int(parts[0]), int(parts[1]), int(parts[2])
        return None

    @staticmethod
    def _parse_popularity(attributes: MediaAttributes) -> int:
        return round(attributes["popularity"] * 100)

    def sync_item_popularity(
        self, results: Results[Item], write: bool, force: bool = False
    ) -> None:
        """Sync Tidal popularity data for library items."""
        self._sync_popularity(
            results=results,
            write=write,
            force=force,
            id_field="tidal_track_id",
            popularity_field="tidal_track_popularity",
            search=lambda ids: self.search_tracks_by_ids(tidal_ids=ids),
            label="item",
        )

    def sync_album_popularity(
        self, results: Results[Album], write: bool, force: bool = False
    ) -> None:
        """Sync Tidal popularity data for library albums."""
        self._sync_popularity(
            results=results,
            write=write,
            force=force,
            id_field="tidal_album_id",
            popularity_field="tidal_album_popularity",
            search=lambda ids: self.search_albums_by_ids(tidal_ids=ids),
            label="album",
        )

    def _sync_popularity(
        self,
        *,
        results: Results[Item] | Results[Album],
        write: bool,
        force: bool,
        id_field: str,
        popularity_field: str,
        search: Callable[[Iterable[str]], Iterable[Info | None]],
        label: Literal["item", "album"],
    ) -> None:
        """Sync Tidal popularity for a generic model (Item or Album)."""
        log.info("Syncing popularity for {0} {1}s", len(results), label)

        model_by_id = {
            tidal_id: model
            for model in results
            if (tidal_id := model.get(id_field))
            and (force or model.get(popularity_field) is None)
        }
        total = len(model_by_id)
        log.debug("{0} {1}s need updates", total, label)
        processed = 0
        for model, new_info in zip(
            model_by_id.values(), search(model_by_id.keys())
        ):
            if not new_info:
                continue

            model[popularity_field] = new_info.get(popularity_field)
            model["tidal_updated"] = time.time()

            model.store()
            if write:
                model.try_write()

            processed += 1
            if processed % 100 == 0 or processed == total:
                log.debug("Synced {0}/{1} {2}s", processed, total, label)

        log.info(
            "Successfully synchronised popularity for {0} {1}s", total, label
        )

    def commands(self) -> list[ui.Subcommand]:
        tidal_cmd = ui.Subcommand(
            "tidal", help="Tidal metadata plugin commands"
        )
        tidal_cmd.parser.add_option(
            "-a",
            "--auth",
            action="store_true",
            help="Authenticate and login to Tidal",
            default=False,
        )

        def auth_func(
            lib: Library, opts: optparse.Values, args: list[str]
        ) -> None:
            if opts.auth:
                self.api.ui_authenticate_flow()
            else:
                tidal_cmd.print_help()

        tidal_cmd.func = auth_func

        # It would be nice to combine both auth and sync commands but
        # currently not really supported in beets. We might be able to
        # revist this if our subcommand parsing changes
        sync_help = (
            "Synchronize Tidal popularity data for library items and albums"
        )
        tidalsync_cmd = ui.Subcommand("tidalsync", help=sync_help)
        tidalsync_cmd.parser.add_album_option()
        tidalsync_cmd.parser.add_option(
            "-f",
            "--force",
            action="store_true",
            dest="force",
            default=False,
            help=(
                "re-fetch popularity even if already present (default: False)"
            ),
        )
        tidalsync_cmd.parser.add_option(
            "-w",
            "--write",
            action="store_true",
            dest="write",
            default=False,
            help="write updated tags to media files (default: False)",
        )
        tidalsync_cmd.parser.set_usage(
            "Usage: beet tidalsync <query> [options]"
        )

        def sync_func(
            lib: Library, opts: optparse.Values, args: list[str]
        ) -> None:
            query = ["data_source:tidal", *args]

            if opts.album:
                self.sync_album_popularity(
                    lib.albums(query), write=opts.write, force=opts.force
                )
            else:
                self.sync_item_popularity(
                    lib.items(query), write=opts.write, force=opts.force
                )

        tidalsync_cmd.func = sync_func

        return [tidal_cmd, tidalsync_cmd]


ISO_8601_RE = re.compile(
    r"^P"
    r"T"
    r"(?:(?P<hours>\d+)H)?"
    r"(?:(?P<minutes>\d+)M)?"
    r"(?:(?P<seconds>\d+)S)?$"
)