/
/
/
1"""API client wrapper for KION Music (MTS Music)."""
2
3from __future__ import annotations
4
5import logging
6from typing import TYPE_CHECKING, Any, cast
7
8from music_assistant_models.errors import (
9 LoginFailed,
10 ProviderUnavailableError,
11 ResourceTemporarilyUnavailable,
12)
13from yandex_music import Album as YandexAlbum
14from yandex_music import Artist as YandexArtist
15from yandex_music import ClientAsync, Search, TrackShort
16from yandex_music import Playlist as YandexPlaylist
17from yandex_music import Track as YandexTrack
18from yandex_music.exceptions import BadRequestError, NetworkError, UnauthorizedError
19from yandex_music.utils.sign_request import get_sign_request
20
21if TYPE_CHECKING:
22 from yandex_music import DownloadInfo
23
24from .constants import DEFAULT_LIMIT
25
26# get-file-info with quality=lossless returns FLAC; default /tracks/.../download-info often does not
27# Prefer flac-mp4/aac-mp4
28GET_FILE_INFO_CODECS = "flac-mp4,flac,aac-mp4,aac,he-aac,mp3,he-aac-mp4"
29# get-file-info: same host as library (all requests go through one API)
30KION_BASE_URL = "https://music.mts.ru/ya_api"
31
32LOGGER = logging.getLogger(__name__)
33
34
35class KionMusicClient:
36 """Wrapper around yandex-music-api ClientAsync for KION Music."""
37
38 def __init__(self, token: str) -> None:
39 """Initialize the KION Music client.
40
41 :param token: KION Music OAuth token.
42 """
43 self._token = token
44 self._client: ClientAsync | None = None
45 self._user_id: int | None = None
46
47 @property
48 def user_id(self) -> int:
49 """Return the user ID."""
50 if self._user_id is None:
51 raise ProviderUnavailableError("Client not initialized, call connect() first")
52 return self._user_id
53
54 async def connect(self) -> bool:
55 """Initialize the client and verify token validity.
56
57 :return: True if connection was successful.
58 :raises LoginFailed: If the token is invalid.
59 """
60 try:
61 self._client = await ClientAsync(self._token, base_url=KION_BASE_URL).init()
62 if self._client.me is None or self._client.me.account is None:
63 raise LoginFailed("Failed to get account info")
64 self._user_id = self._client.me.account.uid
65 LOGGER.debug("Connected to KION Music as user %s", self._user_id)
66 return True
67 except UnauthorizedError as err:
68 raise LoginFailed("Invalid KION Music token") from err
69 except NetworkError as err:
70 msg = "Network error connecting to KION Music"
71 raise ResourceTemporarilyUnavailable(msg) from err
72
73 async def disconnect(self) -> None:
74 """Disconnect the client."""
75 self._client = None
76 self._user_id = None
77
78 def _ensure_connected(self) -> ClientAsync:
79 """Ensure the client is connected and return it."""
80 if self._client is None:
81 raise ProviderUnavailableError("Client not connected, call connect() first")
82 return self._client
83
84 # Library methods
85
86 async def get_liked_tracks(self) -> list[TrackShort]:
87 """Get user's liked tracks.
88
89 :return: List of liked track objects.
90 """
91 client = self._ensure_connected()
92 try:
93 result = await client.users_likes_tracks()
94 if result is None:
95 return []
96 return result.tracks or []
97 except (BadRequestError, NetworkError) as err:
98 LOGGER.error("Error fetching liked tracks: %s", err)
99 raise ResourceTemporarilyUnavailable("Failed to fetch liked tracks") from err
100
101 async def get_liked_albums(self) -> list[YandexAlbum]:
102 """Get user's liked albums with full details (including cover art).
103
104 The users_likes_albums endpoint returns minimal album data without
105 cover_uri, so we fetch full album details in batches afterwards.
106
107 :return: List of liked album objects with full details.
108 """
109 client = self._ensure_connected()
110 try:
111 result = await client.users_likes_albums()
112 if result is None:
113 return []
114 album_ids = [
115 str(like.album.id) for like in result if like.album is not None and like.album.id
116 ]
117 if not album_ids:
118 return []
119 # Fetch full album details in batches to get cover_uri and other metadata
120 batch_size = 50
121 full_albums: list[YandexAlbum] = []
122 for i in range(0, len(album_ids), batch_size):
123 batch = album_ids[i : i + batch_size]
124 try:
125 batch_result = await client.albums(batch)
126 if batch_result:
127 full_albums.extend(batch_result)
128 except (BadRequestError, NetworkError) as batch_err:
129 LOGGER.warning("Error fetching album details batch: %s", batch_err)
130 # Fall back to minimal data for this batch
131 batch_set = set(batch)
132 for like in result:
133 if (
134 like.album is not None
135 and like.album.id
136 and str(like.album.id) in batch_set
137 ):
138 full_albums.append(like.album)
139 return full_albums
140 except (BadRequestError, NetworkError) as err:
141 LOGGER.error("Error fetching liked albums: %s", err)
142 raise ResourceTemporarilyUnavailable("Failed to fetch liked albums") from err
143
144 async def get_liked_artists(self) -> list[YandexArtist]:
145 """Get user's liked artists.
146
147 :return: List of liked artist objects.
148 """
149 client = self._ensure_connected()
150 try:
151 result = await client.users_likes_artists()
152 if result is None:
153 return []
154 return [like.artist for like in result if like.artist is not None]
155 except (BadRequestError, NetworkError) as err:
156 LOGGER.error("Error fetching liked artists: %s", err)
157 raise ResourceTemporarilyUnavailable("Failed to fetch liked artists") from err
158
159 async def get_user_playlists(self) -> list[YandexPlaylist]:
160 """Get user's playlists.
161
162 :return: List of playlist objects.
163 """
164 client = self._ensure_connected()
165 try:
166 result = await client.users_playlists_list()
167 if result is None:
168 return []
169 return list(result)
170 except (BadRequestError, NetworkError) as err:
171 LOGGER.error("Error fetching playlists: %s", err)
172 raise ResourceTemporarilyUnavailable("Failed to fetch playlists") from err
173
174 # Search
175
176 async def search(
177 self,
178 query: str,
179 search_type: str = "all",
180 limit: int = DEFAULT_LIMIT,
181 ) -> Search | None:
182 """Search for tracks, albums, artists, or playlists.
183
184 :param query: Search query string.
185 :param search_type: Type of search ('all', 'track', 'album', 'artist', 'playlist').
186 :param limit: Maximum number of results per type.
187 :return: Search results object.
188 """
189 client = self._ensure_connected()
190 try:
191 return await client.search(query, type_=search_type, page=0, nocorrect=False)
192 except (BadRequestError, NetworkError) as err:
193 LOGGER.error("Search error: %s", err)
194 raise ResourceTemporarilyUnavailable("Search failed") from err
195
196 # Get single items
197
198 async def get_track(self, track_id: str) -> YandexTrack | None:
199 """Get a single track by ID.
200
201 :param track_id: Track ID.
202 :return: Track object or None if not found.
203 """
204 client = self._ensure_connected()
205 try:
206 tracks = await client.tracks([track_id])
207 return tracks[0] if tracks else None
208 except (BadRequestError, NetworkError) as err:
209 LOGGER.error("Error fetching track %s: %s", track_id, err)
210 return None
211
212 async def get_tracks(self, track_ids: list[str]) -> list[YandexTrack]:
213 """Get multiple tracks by IDs.
214
215 :param track_ids: List of track IDs.
216 :return: List of track objects.
217 :raises ResourceTemporarilyUnavailable: On network errors after retry.
218 """
219 client = self._ensure_connected()
220 try:
221 result = await client.tracks(track_ids)
222 return result or []
223 except NetworkError as err:
224 # Retry once on network errors (timeout, disconnect, etc.)
225 LOGGER.warning("Network error fetching tracks, retrying once: %s", err)
226 try:
227 result = await client.tracks(track_ids)
228 return result or []
229 except NetworkError as retry_err:
230 LOGGER.error("Error fetching tracks (retry failed): %s", retry_err)
231 raise ResourceTemporarilyUnavailable("Failed to fetch tracks") from retry_err
232 except BadRequestError as err:
233 LOGGER.error("Error fetching tracks: %s", err)
234 return []
235
236 async def get_album(self, album_id: str) -> YandexAlbum | None:
237 """Get a single album by ID.
238
239 :param album_id: Album ID.
240 :return: Album object or None if not found.
241 """
242 client = self._ensure_connected()
243 try:
244 albums = await client.albums([album_id])
245 return albums[0] if albums else None
246 except (BadRequestError, NetworkError) as err:
247 LOGGER.error("Error fetching album %s: %s", album_id, err)
248 return None
249
250 async def get_album_with_tracks(self, album_id: str) -> YandexAlbum | None:
251 """Get an album with its tracks.
252
253 Uses the same semantics as the web client: albums/{id}/with-tracks
254 with resumeStream, richTracks, withListeningFinished when the library
255 passes them through.
256
257 :param album_id: Album ID.
258 :return: Album object with tracks or None if not found.
259 """
260 client = self._ensure_connected()
261 try:
262 return await client.albums_with_tracks(
263 album_id,
264 resumeStream=True,
265 richTracks=True,
266 withListeningFinished=True,
267 )
268 except TypeError:
269 # Older yandex-music may not accept these kwargs
270 return await client.albums_with_tracks(album_id)
271 except (BadRequestError, NetworkError) as err:
272 LOGGER.error("Error fetching album with tracks %s: %s", album_id, err)
273 return None
274
275 async def get_artist(self, artist_id: str) -> YandexArtist | None:
276 """Get a single artist by ID.
277
278 :param artist_id: Artist ID.
279 :return: Artist object or None if not found.
280 """
281 client = self._ensure_connected()
282 try:
283 artists = await client.artists([artist_id])
284 return artists[0] if artists else None
285 except (BadRequestError, NetworkError) as err:
286 LOGGER.error("Error fetching artist %s: %s", artist_id, err)
287 return None
288
289 async def get_artist_albums(
290 self, artist_id: str, limit: int = DEFAULT_LIMIT
291 ) -> list[YandexAlbum]:
292 """Get artist's albums.
293
294 :param artist_id: Artist ID.
295 :param limit: Maximum number of albums.
296 :return: List of album objects.
297 """
298 client = self._ensure_connected()
299 try:
300 result = await client.artists_direct_albums(artist_id, page=0, page_size=limit)
301 if result is None:
302 return []
303 return result.albums or []
304 except (BadRequestError, NetworkError) as err:
305 LOGGER.error("Error fetching artist albums %s: %s", artist_id, err)
306 return []
307
308 async def get_artist_tracks(
309 self, artist_id: str, limit: int = DEFAULT_LIMIT
310 ) -> list[YandexTrack]:
311 """Get artist's top tracks.
312
313 :param artist_id: Artist ID.
314 :param limit: Maximum number of tracks.
315 :return: List of track objects.
316 """
317 client = self._ensure_connected()
318 try:
319 result = await client.artists_tracks(artist_id, page=0, page_size=limit)
320 if result is None:
321 return []
322 return result.tracks or []
323 except (BadRequestError, NetworkError) as err:
324 LOGGER.error("Error fetching artist tracks %s: %s", artist_id, err)
325 return []
326
327 async def get_playlist(self, user_id: str, playlist_id: str) -> YandexPlaylist | None:
328 """Get a playlist by ID.
329
330 :param user_id: User ID (owner of the playlist).
331 :param playlist_id: Playlist ID (kind).
332 :return: Playlist object or None if not found.
333 """
334 client = self._ensure_connected()
335 try:
336 result = await client.users_playlists(kind=int(playlist_id), user_id=user_id)
337 if isinstance(result, list):
338 return result[0] if result else None
339 return result
340 except (BadRequestError, NetworkError) as err:
341 LOGGER.error("Error fetching playlist %s/%s: %s", user_id, playlist_id, err)
342 return None
343
344 # Streaming
345
346 async def get_track_download_info(
347 self, track_id: str, get_direct_links: bool = True
348 ) -> list[DownloadInfo]:
349 """Get download info for a track.
350
351 :param track_id: Track ID.
352 :param get_direct_links: Whether to get direct download links.
353 :return: List of download info objects.
354 """
355 client = self._ensure_connected()
356 try:
357 result = await client.tracks_download_info(track_id, get_direct_links=get_direct_links)
358 return result or []
359 except (BadRequestError, NetworkError) as err:
360 LOGGER.error("Error fetching download info for track %s: %s", track_id, err)
361 return []
362
363 async def get_track_file_info_lossless(self, track_id: str) -> dict[str, Any] | None:
364 """Request lossless stream via get-file-info (quality=lossless).
365
366 The /tracks/{id}/download-info endpoint often returns only MP3; get-file-info
367 with quality=lossless and codecs=flac,... returns FLAC when available.
368
369 :param track_id: Track ID.
370 :return: Parsed downloadInfo dict (url, codec, urls, ...) or None on error.
371 """
372 client = self._ensure_connected()
373 sign = get_sign_request(track_id)
374 base_params = {
375 "ts": sign.timestamp,
376 "trackId": track_id,
377 "quality": "lossless",
378 "codecs": GET_FILE_INFO_CODECS,
379 "sign": sign.value,
380 }
381
382 def _parse_file_info_result(raw: dict[str, Any] | None) -> dict[str, Any] | None:
383 if not raw or not isinstance(raw, dict):
384 return None
385 download_info = raw.get("download_info")
386 if not download_info or not download_info.get("url"):
387 return None
388 return cast("dict[str, Any]", download_info)
389
390 url = f"{KION_BASE_URL}/get-file-info"
391 params_encraw = {**base_params, "transports": "encraw"}
392 try:
393 result = await client._request.get(url, params=params_encraw)
394 return _parse_file_info_result(result)
395 except (BadRequestError, NetworkError) as err:
396 LOGGER.debug(
397 "get-file-info lossless for track %s: %s %s",
398 track_id,
399 type(err).__name__,
400 getattr(err, "message", str(err)) or repr(err),
401 )
402 return None
403 except UnauthorizedError as err:
404 LOGGER.debug(
405 "get-file-info lossless for track %s (transports=encraw): %s %s",
406 track_id,
407 type(err).__name__,
408 getattr(err, "message", str(err)) or repr(err),
409 )
410 LOGGER.debug(
411 "If you have KION Music Plus and this track has lossless, "
412 "try a token from the web client (music.mts.ru)."
413 )
414 params_raw = {**base_params, "transports": "raw"}
415 try:
416 result = await client._request.get(url, params=params_raw)
417 return _parse_file_info_result(result)
418 except (BadRequestError, NetworkError, UnauthorizedError) as retry_err:
419 LOGGER.debug(
420 "get-file-info lossless for track %s (transports=raw): %s %s",
421 track_id,
422 type(retry_err).__name__,
423 getattr(retry_err, "message", str(retry_err)) or repr(retry_err),
424 )
425 return None
426
427 # Library modifications
428
429 async def like_track(self, track_id: str) -> bool:
430 """Add a track to liked tracks.
431
432 :param track_id: Track ID to like.
433 :return: True if successful.
434 """
435 client = self._ensure_connected()
436 try:
437 result = await client.users_likes_tracks_add(track_id)
438 return result is not None
439 except (BadRequestError, NetworkError) as err:
440 LOGGER.error("Error liking track %s: %s", track_id, err)
441 return False
442
443 async def unlike_track(self, track_id: str) -> bool:
444 """Remove a track from liked tracks.
445
446 :param track_id: Track ID to unlike.
447 :return: True if successful.
448 """
449 client = self._ensure_connected()
450 try:
451 result = await client.users_likes_tracks_remove(track_id)
452 return result is not None
453 except (BadRequestError, NetworkError) as err:
454 LOGGER.error("Error unliking track %s: %s", track_id, err)
455 return False
456
457 async def like_album(self, album_id: str) -> bool:
458 """Add an album to liked albums.
459
460 :param album_id: Album ID to like.
461 :return: True if successful.
462 """
463 client = self._ensure_connected()
464 try:
465 result = await client.users_likes_albums_add(album_id)
466 return result is not None
467 except (BadRequestError, NetworkError) as err:
468 LOGGER.error("Error liking album %s: %s", album_id, err)
469 return False
470
471 async def unlike_album(self, album_id: str) -> bool:
472 """Remove an album from liked albums.
473
474 :param album_id: Album ID to unlike.
475 :return: True if successful.
476 """
477 client = self._ensure_connected()
478 try:
479 result = await client.users_likes_albums_remove(album_id)
480 return result is not None
481 except (BadRequestError, NetworkError) as err:
482 LOGGER.error("Error unliking album %s: %s", album_id, err)
483 return False
484
485 async def like_artist(self, artist_id: str) -> bool:
486 """Add an artist to liked artists.
487
488 :param artist_id: Artist ID to like.
489 :return: True if successful.
490 """
491 client = self._ensure_connected()
492 try:
493 result = await client.users_likes_artists_add(artist_id)
494 return result is not None
495 except (BadRequestError, NetworkError) as err:
496 LOGGER.error("Error liking artist %s: %s", artist_id, err)
497 return False
498
499 async def unlike_artist(self, artist_id: str) -> bool:
500 """Remove an artist from liked artists.
501
502 :param artist_id: Artist ID to unlike.
503 :return: True if successful.
504 """
505 client = self._ensure_connected()
506 try:
507 result = await client.users_likes_artists_remove(artist_id)
508 return result is not None
509 except (BadRequestError, NetworkError) as err:
510 LOGGER.error("Error unliking artist %s: %s", artist_id, err)
511 return False
512