import os

from yandex_music import Client, Track

client = Client(os.getenv('YANDEX_TOKEN')).init()


def search(_str: str, _type: str = 'all') -> dict:
    album_id = None
    album_title = None
    artist_id = None
    artist_name = None
    result_json = {}

    search_result = client.search(_str, type_=_type)
    type_ = search_result.best.type
    best = search_result.best.result

    if type_ == 'track':
        artists = ''
        if best.artists:
            artists = ', '.join(artist.name for artist in best.artists)
        if best.albums:
            for i in best.albums:
                album_id = i.id
                album_title = i.title

        # Generate json for track
        result_json = {
            "artist": {
                "id": best.id,
                "name": artists
            },
            "track": {
                "id": best.trackId,
                "title": best.title
            },
            "album": {
                "id": album_id,
                "title": album_title
            }
        }
    elif type_ == 'artist':
        result_json = {
            "artist": {
                "id": best.id,
                "name": best.name
            }
        }
    elif type_ == 'album':
        if best.artists:
            for i in best.artists:
                artist_id = i.id
                artist_name = i.name
        _tracks
        list[Track] = []
        tracks = {}
        album = client.albums_with_tracks(best.id)
        for i, volume in enumerate(album.volumes):
            if len(album.volumes) > 1:
                _tracks.append(f'💿 Диск {i + 1}')
            _tracks += volume
        _i = 1
        for track in _tracks:
            tracks[f"track_{_i}"] = {
                "id": track.id,
                "tittle": track.title
            }
            _i += 1

        result_json = {
            "artist": {
                "id": artist_id,
                "name": artist_name
            },
            "album": {
                "id": best.id,
                "title": best.title
            },
            "tracks": tracks
        }
    print(_type)
    return result_json