/
/
/
1"""KION Music provider implementation."""
2
3from __future__ import annotations
4
5from typing import TYPE_CHECKING
6
7from music_assistant_models.enums import MediaType
8from music_assistant_models.errors import (
9 InvalidDataError,
10 LoginFailed,
11 MediaNotFoundError,
12 ProviderUnavailableError,
13 ResourceTemporarilyUnavailable,
14)
15from music_assistant_models.media_items import (
16 Album,
17 Artist,
18 ItemMapping,
19 MediaItemType,
20 Playlist,
21 SearchResults,
22 Track,
23)
24
25from music_assistant.controllers.cache import use_cache
26from music_assistant.models.music_provider import MusicProvider
27
28from .api_client import KionMusicClient
29from .constants import CONF_TOKEN, PLAYLIST_ID_SPLITTER
30from .parsers import parse_album, parse_artist, parse_playlist, parse_track
31from .streaming import KionMusicStreamingManager
32
33if TYPE_CHECKING:
34 from collections.abc import AsyncGenerator
35
36 from music_assistant_models.streamdetails import StreamDetails
37
38
39class KionMusicProvider(MusicProvider):
40 """Implementation of a KION Music MusicProvider."""
41
42 _client: KionMusicClient | None = None
43 _streaming: KionMusicStreamingManager | None = None
44
45 @property
46 def client(self) -> KionMusicClient:
47 """Return the KION Music client."""
48 if self._client is None:
49 raise ProviderUnavailableError("Provider not initialized")
50 return self._client
51
52 @property
53 def streaming(self) -> KionMusicStreamingManager:
54 """Return the streaming manager."""
55 if self._streaming is None:
56 raise ProviderUnavailableError("Provider not initialized")
57 return self._streaming
58
59 async def handle_async_init(self) -> None:
60 """Handle async initialization of the provider."""
61 token = self.config.get_value(CONF_TOKEN)
62 if not token:
63 raise LoginFailed("No KION Music token provided")
64
65 self._client = KionMusicClient(str(token))
66 await self._client.connect()
67 self._streaming = KionMusicStreamingManager(self)
68 self.logger.info("Successfully connected to KION Music")
69
70 async def unload(self, is_removed: bool = False) -> None:
71 """Handle unload/close of the provider.
72
73 :param is_removed: Whether the provider is being removed.
74 """
75 if self._client:
76 await self._client.disconnect()
77 self._client = None
78 self._streaming = None
79 await super().unload(is_removed)
80
81 def get_item_mapping(self, media_type: MediaType | str, key: str, name: str) -> ItemMapping:
82 """Create a generic item mapping.
83
84 :param media_type: The media type.
85 :param key: The item ID.
86 :param name: The item name.
87 :return: An ItemMapping instance.
88 """
89 if isinstance(media_type, str):
90 media_type = MediaType(media_type)
91 return ItemMapping(
92 media_type=media_type,
93 item_id=key,
94 provider=self.instance_id,
95 name=name,
96 )
97
98 # Search
99
100 @use_cache(3600 * 24 * 14)
101 async def search(
102 self, search_query: str, media_types: list[MediaType], limit: int = 5
103 ) -> SearchResults:
104 """Perform search on KION Music.
105
106 :param search_query: The search query.
107 :param media_types: List of media types to search for.
108 :param limit: Maximum number of results per type.
109 :return: SearchResults with found items.
110 """
111 result = SearchResults()
112
113 # Determine search type based on requested media types
114 # Map MediaType to KION API search type
115 type_mapping = {
116 MediaType.TRACK: "track",
117 MediaType.ALBUM: "album",
118 MediaType.ARTIST: "artist",
119 MediaType.PLAYLIST: "playlist",
120 }
121 requested_types = [type_mapping[mt] for mt in media_types if mt in type_mapping]
122
123 # Use specific type if only one requested, otherwise search all
124 search_type = requested_types[0] if len(requested_types) == 1 else "all"
125
126 search_result = await self.client.search(search_query, search_type=search_type, limit=limit)
127 if not search_result:
128 return result
129
130 # Parse tracks
131 if MediaType.TRACK in media_types and search_result.tracks:
132 for track in search_result.tracks.results[:limit]:
133 try:
134 result.tracks = [*result.tracks, parse_track(self, track)]
135 except InvalidDataError as err:
136 self.logger.debug("Error parsing track: %s", err)
137
138 # Parse albums
139 if MediaType.ALBUM in media_types and search_result.albums:
140 for album in search_result.albums.results[:limit]:
141 try:
142 result.albums = [*result.albums, parse_album(self, album)]
143 except InvalidDataError as err:
144 self.logger.debug("Error parsing album: %s", err)
145
146 # Parse artists
147 if MediaType.ARTIST in media_types and search_result.artists:
148 for artist in search_result.artists.results[:limit]:
149 try:
150 result.artists = [*result.artists, parse_artist(self, artist)]
151 except InvalidDataError as err:
152 self.logger.debug("Error parsing artist: %s", err)
153
154 # Parse playlists
155 if MediaType.PLAYLIST in media_types and search_result.playlists:
156 for playlist in search_result.playlists.results[:limit]:
157 try:
158 result.playlists = [*result.playlists, parse_playlist(self, playlist)]
159 except InvalidDataError as err:
160 self.logger.debug("Error parsing playlist: %s", err)
161
162 return result
163
164 # Get single items
165
166 @use_cache(3600 * 24 * 30)
167 async def get_artist(self, prov_artist_id: str) -> Artist:
168 """Get artist details by ID.
169
170 :param prov_artist_id: The provider artist ID.
171 :return: Artist object.
172 :raises MediaNotFoundError: If artist not found.
173 """
174 artist = await self.client.get_artist(prov_artist_id)
175 if not artist:
176 raise MediaNotFoundError(f"Artist {prov_artist_id} not found")
177 return parse_artist(self, artist)
178
179 @use_cache(3600 * 24 * 30)
180 async def get_album(self, prov_album_id: str) -> Album:
181 """Get album details by ID.
182
183 :param prov_album_id: The provider album ID.
184 :return: Album object.
185 :raises MediaNotFoundError: If album not found.
186 """
187 album = await self.client.get_album(prov_album_id)
188 if not album:
189 raise MediaNotFoundError(f"Album {prov_album_id} not found")
190 return parse_album(self, album)
191
192 @use_cache(3600 * 24 * 30)
193 async def get_track(self, prov_track_id: str) -> Track:
194 """Get track details by ID.
195
196 :param prov_track_id: The provider track ID.
197 :return: Track object.
198 :raises MediaNotFoundError: If track not found.
199 """
200 yandex_track = await self.client.get_track(prov_track_id)
201 if not yandex_track:
202 raise MediaNotFoundError(f"Track {prov_track_id} not found")
203 return parse_track(self, yandex_track)
204
205 @use_cache(3600 * 24 * 30)
206 async def get_playlist(self, prov_playlist_id: str) -> Playlist:
207 """Get playlist details by ID.
208
209 :param prov_playlist_id: The provider playlist ID (format: "owner_id:kind").
210 :return: Playlist object.
211 :raises MediaNotFoundError: If playlist not found.
212 """
213 # Parse the playlist ID (format: owner_id:kind)
214 if PLAYLIST_ID_SPLITTER in prov_playlist_id:
215 owner_id, kind = prov_playlist_id.split(PLAYLIST_ID_SPLITTER, 1)
216 else:
217 owner_id = str(self.client.user_id)
218 kind = prov_playlist_id
219
220 playlist = await self.client.get_playlist(owner_id, kind)
221 if not playlist:
222 raise MediaNotFoundError(f"Playlist {prov_playlist_id} not found")
223 return parse_playlist(self, playlist)
224
225 # Get related items
226
227 @use_cache(3600 * 24 * 30)
228 async def get_album_tracks(self, prov_album_id: str) -> list[Track]:
229 """Get album tracks.
230
231 :param prov_album_id: The provider album ID.
232 :return: List of Track objects.
233 """
234 album = await self.client.get_album_with_tracks(prov_album_id)
235 if not album or not album.volumes:
236 return []
237
238 tracks = []
239 for volume_index, volume in enumerate(album.volumes):
240 for track_index, track in enumerate(volume):
241 try:
242 parsed_track = parse_track(self, track)
243 parsed_track.disc_number = volume_index + 1
244 parsed_track.track_number = track_index + 1
245 tracks.append(parsed_track)
246 except InvalidDataError as err:
247 self.logger.debug("Error parsing album track: %s", err)
248 return tracks
249
250 @use_cache(3600 * 3)
251 async def get_playlist_tracks(self, prov_playlist_id: str, page: int = 0) -> list[Track]:
252 """Get playlist tracks.
253
254 :param prov_playlist_id: The provider playlist ID (format: "owner_id:kind").
255 :param page: Page number for pagination.
256 :return: List of Track objects.
257 """
258 # KION Music API returns all playlist tracks in one call (no server-side pagination)
259 if page > 0:
260 return []
261
262 self.logger.debug("get_playlist_tracks called: %s", prov_playlist_id)
263 # Parse the playlist ID (format: owner_id:kind)
264 if PLAYLIST_ID_SPLITTER in prov_playlist_id:
265 owner_id, kind = prov_playlist_id.split(PLAYLIST_ID_SPLITTER, 1)
266 else:
267 owner_id = str(self.client.user_id)
268 kind = prov_playlist_id
269
270 self.logger.debug("Fetching playlist %s/%s from API...", owner_id, kind)
271 playlist = await self.client.get_playlist(owner_id, kind)
272 if not playlist:
273 self.logger.debug("Playlist %s/%s not found", owner_id, kind)
274 return []
275
276 # API sometimes returns playlist without tracks; fetch them explicitly if needed
277 tracks_list = playlist.tracks or []
278 track_count = getattr(playlist, "track_count", None) or 0
279 self.logger.debug(
280 "Playlist %s/%s: track_count=%s, tracks_in_response=%s",
281 owner_id,
282 kind,
283 track_count,
284 len(tracks_list),
285 )
286 if not tracks_list and track_count > 0:
287 self.logger.debug("No tracks in response, calling fetch_tracks_async...")
288 try:
289 tracks_list = await playlist.fetch_tracks_async()
290 self.logger.debug("fetch_tracks_async returned %s tracks", len(tracks_list or []))
291 except Exception as err:
292 self.logger.warning("fetch_tracks_async failed: %s", err)
293 if not tracks_list:
294 self.logger.warning(
295 "Playlist %s/%s: expected %s tracks but got none",
296 owner_id,
297 kind,
298 track_count,
299 )
300 raise ResourceTemporarilyUnavailable(
301 "Playlist tracks not available; try again later"
302 ) from None
303
304 if not tracks_list:
305 self.logger.debug("Playlist %s/%s has no tracks", owner_id, kind)
306 return []
307
308 # API returns TrackShort objects, we need to fetch full track info
309 track_ids = [
310 str(track.track_id) if hasattr(track, "track_id") else str(track.id)
311 for track in tracks_list
312 if track
313 ]
314 if not track_ids:
315 return []
316
317 self.logger.debug("Fetching full details for %s tracks...", len(track_ids))
318 # Fetch full track details in batches to avoid timeouts
319 batch_size = 50
320 full_tracks = []
321 for i in range(0, len(track_ids), batch_size):
322 batch = track_ids[i : i + batch_size]
323 self.logger.debug("Fetching batch %s-%s...", i, i + len(batch))
324 batch_result = await self.client.get_tracks(batch)
325 self.logger.debug("Batch returned %s tracks", len(batch_result or []))
326 full_tracks.extend(batch_result or [])
327
328 if track_ids and not full_tracks:
329 self.logger.warning("Got 0 full tracks for %s IDs", len(track_ids))
330 raise ResourceTemporarilyUnavailable(
331 "Failed to load track details; try again later"
332 ) from None
333
334 tracks = []
335 for track in full_tracks:
336 try:
337 tracks.append(parse_track(self, track))
338 except InvalidDataError as err:
339 self.logger.debug("Error parsing playlist track: %s", err)
340 self.logger.debug("Returning %s parsed tracks", len(tracks))
341 return tracks
342
343 @use_cache(3600 * 24 * 7)
344 async def get_artist_albums(self, prov_artist_id: str) -> list[Album]:
345 """Get artist's albums.
346
347 :param prov_artist_id: The provider artist ID.
348 :return: List of Album objects.
349 """
350 albums = await self.client.get_artist_albums(prov_artist_id)
351 result = []
352 for album in albums:
353 try:
354 result.append(parse_album(self, album))
355 except InvalidDataError as err:
356 self.logger.debug("Error parsing artist album: %s", err)
357 return result
358
359 @use_cache(3600 * 24 * 7)
360 async def get_artist_toptracks(self, prov_artist_id: str) -> list[Track]:
361 """Get artist's top tracks.
362
363 :param prov_artist_id: The provider artist ID.
364 :return: List of Track objects.
365 """
366 tracks = await self.client.get_artist_tracks(prov_artist_id)
367 result = []
368 for track in tracks:
369 try:
370 result.append(parse_track(self, track))
371 except InvalidDataError as err:
372 self.logger.debug("Error parsing artist track: %s", err)
373 return result
374
375 # Library methods
376
377 async def get_library_artists(self) -> AsyncGenerator[Artist, None]:
378 """Retrieve library artists from KION Music."""
379 artists = await self.client.get_liked_artists()
380 for artist in artists:
381 try:
382 yield parse_artist(self, artist)
383 except InvalidDataError as err:
384 self.logger.debug("Error parsing library artist: %s", err)
385
386 async def get_library_albums(self) -> AsyncGenerator[Album, None]:
387 """Retrieve library albums from KION Music."""
388 albums = await self.client.get_liked_albums()
389 for album in albums:
390 try:
391 yield parse_album(self, album)
392 except InvalidDataError as err:
393 self.logger.debug("Error parsing library album: %s", err)
394
395 async def get_library_tracks(self) -> AsyncGenerator[Track, None]:
396 """Retrieve library tracks from KION Music."""
397 track_shorts = await self.client.get_liked_tracks()
398 if not track_shorts:
399 return
400
401 # Fetch full track details in batches
402 track_ids = [str(ts.track_id) for ts in track_shorts if ts.track_id]
403 batch_size = 50
404 for i in range(0, len(track_ids), batch_size):
405 batch_ids = track_ids[i : i + batch_size]
406 full_tracks = await self.client.get_tracks(batch_ids)
407 for track in full_tracks:
408 try:
409 yield parse_track(self, track)
410 except InvalidDataError as err:
411 self.logger.debug("Error parsing library track: %s", err)
412
413 async def get_library_playlists(self) -> AsyncGenerator[Playlist, None]:
414 """Retrieve library playlists from KION Music."""
415 playlists = await self.client.get_user_playlists()
416 for playlist in playlists:
417 try:
418 yield parse_playlist(self, playlist)
419 except InvalidDataError as err:
420 self.logger.debug("Error parsing library playlist: %s", err)
421
422 # Library edit methods
423
424 async def library_add(self, item: MediaItemType) -> bool:
425 """Add item to library.
426
427 :param item: The media item to add.
428 :return: True if successful.
429 """
430 prov_item_id = self._get_provider_item_id(item)
431 if not prov_item_id:
432 return False
433
434 if item.media_type == MediaType.TRACK:
435 return await self.client.like_track(prov_item_id)
436 if item.media_type == MediaType.ALBUM:
437 return await self.client.like_album(prov_item_id)
438 if item.media_type == MediaType.ARTIST:
439 return await self.client.like_artist(prov_item_id)
440 return False
441
442 async def library_remove(self, prov_item_id: str, media_type: MediaType) -> bool:
443 """Remove item from library.
444
445 :param prov_item_id: The provider item ID.
446 :param media_type: The media type.
447 :return: True if successful.
448 """
449 if media_type == MediaType.TRACK:
450 return await self.client.unlike_track(prov_item_id)
451 if media_type == MediaType.ALBUM:
452 return await self.client.unlike_album(prov_item_id)
453 if media_type == MediaType.ARTIST:
454 return await self.client.unlike_artist(prov_item_id)
455 return False
456
457 def _get_provider_item_id(self, item: MediaItemType) -> str | None:
458 """Get provider item ID from media item."""
459 for mapping in item.provider_mappings:
460 if mapping.provider_instance == self.instance_id:
461 return mapping.item_id
462 return item.item_id if item.provider == self.instance_id else None
463
464 # Streaming
465
466 async def get_stream_details(
467 self, item_id: str, media_type: MediaType = MediaType.TRACK
468 ) -> StreamDetails:
469 """Get stream details for a track.
470
471 :param item_id: The track ID.
472 :param media_type: The media type (should be TRACK).
473 :return: StreamDetails for the track.
474 """
475 return await self.streaming.get_stream_details(item_id)
476