#!/bin/python3
import argparse
import itertools
import logging
import re
import time
import typing
from collections.abc import Generator, Iterable
from pathlib import Path
from typing import Optional, Union
from urllib.parse import urlparse

from yandex_music import Album, Client, Playlist, Track

from ymd import core

DEFAULT_DELAY = 0

TRACK_RE = re.compile(r"track/(\d+)")
ALBUM_RE = re.compile(r"album/(\d+)$")
ARTIST_RE = re.compile(r"artist/(\d+)$")
PLAYLIST_RE = re.compile(r"([\w\-._]+)/playlists/(\d+)$")

logger = logging.getLogger("yandex-music-downloader")


def show_default(text: Optional[str] = None) -> str:
    default = "по умолчанию: %(default)s"
    if text is None:
        return default
    return f"{text} ({default})"


def main():
    parser = argparse.ArgumentParser(
        description="Загрузчик музыки с сервиса Яндекс.Музыка",
        formatter_class=argparse.RawTextHelpFormatter,
    )

    common_group = parser.add_argument_group("Общие параметры")
    common_group.add_argument(
        "--quality",
        metavar="<Качество>",
        default=0,
        type=int,
        help="Качество трека:\n0 - Низкое (mp3 128kbps)\n1 - Высокое (mp3 320kbps)\n(по умолчанию: %(default)s)",
    )
    common_group.add_argument(
        "--skip-existing", action="store_true", help="Пропускать уже загруженные треки"
    )
    common_group.add_argument(
        "--add-lyrics", action="store_true", help="Загружать тексты песен"
    )
    common_group.add_argument(
        "--embed-cover", action="store_true", help="Встраивать обложку в .mp3 файл"
    )
    common_group.add_argument(
        "--cover-resolution",
        default=core.DEFAULT_COVER_RESOLUTION,
        metavar="<Разрешение обложки>",
        type=int,
        help=show_default(None),
    )
    common_group.add_argument(
        "--delay",
        default=DEFAULT_DELAY,
        metavar="<Задержка>",
        type=int,
        help=show_default("Задержка между запросами, в секундах"),
    )
    common_group.add_argument(
        "--stick-to-artist",
        action="store_true",
        help="Загружать альбомы, созданные только данным исполнителем",
    )
    common_group.add_argument(
        "--only-music",
        action="store_true",
        help="Загружать только музыкальные альбомы"
        " (пропускать подкасты и аудиокниги)",
    )
    common_group.add_argument("--debug", action="store_true", help=argparse.SUPPRESS)

    id_group_meta = parser.add_argument_group("ID")
    id_group = id_group_meta.add_mutually_exclusive_group(required=True)
    id_group.add_argument("--artist-id", metavar="<ID исполнителя>")
    id_group.add_argument("--album-id", metavar="<ID альбома>")
    id_group.add_argument("--track-id", metavar="<ID трека>")
    id_group.add_argument(
        "--playlist-id",
        metavar="<владелец плейлиста>/<тип плейлиста>",
    )
    id_group.add_argument("-u", "--url", help="URL исполнителя/альбома/трека/плейлиста")

    path_group = parser.add_argument_group("Указание пути")
    path_group.add_argument(
        "--unsafe-path",
        action="store_true",
        help="Не очищать путь от недопустимых символов",
    )
    path_group.add_argument(
        "--dir",
        default=".",
        metavar="<Папка>",
        help=show_default("Папка для загрузки музыки"),
        type=Path,
    )
    path_group.add_argument(
        "--path-pattern",
        default=core.DEFAULT_PATH_PATTERN,
        metavar="<Паттерн>",
        type=Path,
        help=show_default(
            "Поддерживает следующие заполнители:"
            " #number, #artist, #album-artist, #title,"
            " #album, #year, #artist-id, #album-id, #track-id, #number-padded"
        ),
    )

    auth_group = parser.add_argument_group("Авторизация")
    auth_group.add_argument(
        "--token",
        required=True,
        metavar="<Токен>",
        help="Токен для авторизации. См. README для способов получения",
    )

    args = parser.parse_args()

    logging.basicConfig(
        format="%(asctime)s |%(levelname)s| %(name)s: %(message)s",
        datefmt="%H:%M:%S",
        level=logging.DEBUG if args.debug else logging.ERROR,
    )

    if args.url is not None:
        parsed_url = urlparse(args.url)
        path = parsed_url.path
        if match := ARTIST_RE.search(path):
            args.artist_id = match.group(1)
        elif match := ALBUM_RE.search(path):
            args.album_id = match.group(1)
        elif match := TRACK_RE.search(path):
            args.track_id = match.group(1)
        elif match := PLAYLIST_RE.search(path):
            args.playlist_id = match.group(1) + "/" + match.group(2)
        else:
            print("Параметер url указан в неверном формате")
            return 1

    client = Client(args.token).init()
    result_tracks: Iterable[Track] = []

    def album_tracks_gen(album_ids: Iterable[Union[int, str]]) -> Generator[Track]:
        for album_id in album_ids:
            if full_album := client.albums_with_tracks(album_id):
                if volumes := full_album.volumes:
                    yield from itertools.chain.from_iterable(volumes)

    if args.artist_id is not None:

        def filter_album(album: Album) -> bool:
            title = album.title
            if album.id is None or not album.available:
                print(f'Альбом "{title}" не доступен для скачивания')
            elif args.only_music and album.meta_type != "music":
                print(f'Альбом "{title}" пропущен' " т.к. не является музыкальным")
            elif args.stick_to_artist and album.artists[0].id != int(args.artist_id):
                print(f'Альбом "{title}" пропущен' " из-за флага --stick-to-artist")
            else:
                return True
            return False

        def albums_gen() -> Generator[Album]:
            has_next = True
            page = 0
            while has_next:
                if albums := client.artists_direct_albums(args.artist_id, page):
                    yield from albums.albums
                else:
                    break
                if pager := albums.pager:
                    page = pager.page + 1
                    has_next = pager.per_page * page < pager.total
                else:
                    break

        result_tracks = album_tracks_gen(
            a.id for a in albums_gen() if filter_album(a) and a.id is not None
        )
    elif args.album_id is not None:
        result_tracks = album_tracks_gen((args.album_id,))
    elif args.track_id is not None:
        track = client.tracks(args.track_id)
        result_tracks = track
    elif args.playlist_id is not None:
        user, kind = args.playlist_id.split("/")
        playlist = typing.cast(Playlist, client.users_playlists(kind, user))

        def playlist_tracks_gen() -> Generator[Track]:
            tracks = playlist.fetch_tracks()
            for track in tracks:
                yield track.fetch_track()

        result_tracks = playlist_tracks_gen()

    covers_cache: dict[int, bytes] = {}
    for track in result_tracks:
        if not track.available:
            print(f"Трек {track.title} не доступен для скачивания")
            continue

        save_path = args.dir / core.prepare_track_path(
            args.path_pattern,
            track,
            args.unsafe_path,
        )
        if args.skip_existing and save_path.is_file():
            continue

        save_dir = save_path.parent
        if not save_dir.is_dir():
            save_dir.mkdir(parents=True)

        print(f"Загружается {save_path}")
        core.download_track(
            track=track,
            target_path=save_path,
            quality=args.quality,
            add_lyrics=args.add_lyrics,
            embed_cover=args.embed_cover,
            cover_resolution=args.cover_resolution,
            covers_cache=covers_cache,
        )
        if args.delay > 0:
            time.sleep(args.delay)