from yt_dlp import YoutubeDL
import urllib.parse
from ytmusicapi import YTMusic
import re
import concurrent.futures
import json
import argparse
import sys
import os
import logging
logging.basicConfig(stream=sys.stdout)
logger = logging.getLogger()
LOGLEVEL = os.environ.get('LOG_LEVEL', 'INFO').upper()
logger.setLevel(LOGLEVEL)
logger.debug("Starting...")

def album_info(data):

    logger.info(json.dumps(data['playlist'] if 'playlist' in data else "N/A", indent=4))
    logger.info(json.dumps(data['playlist_id'] if 'playlist_id' in data else "N/A", indent=4))
    return data['playlist_id'] if 'playlist_id' in data else "N/A"

def vid_info(data):

    logger.info(json.dumps(data['title'] if 'title' in data else "N/A", indent=4))
    logger.info(json.dumps(data['album'] if 'album' in data else "N/A", indent=4))
    return data['album'] if 'album' in data else "N/A"

ydl_opts = {
    'format': 'bestaudio/best',
    'postprocessors': [{
        'key': 'FFmpegExtractAudio',
        'preferredcodec': 'flac',
        'preferredquality': 'best',
    },
    {'key': 'FFmpegMetadata'},
    {'key': 'EmbedThumbnail'}],
    'writethumbnail': True,
    'embedthumbnail': True,
     'add_metadata': True,
    'logger': logger,
    'outtmpl': '%(playlist)s/%(playlist_index)s - %(title)s.%(ext)s',
}


def download_album(album_id):
    data = yt.download(album_id)
    return data


if __name__ == "__main__":
    parser = argparse.ArgumentParser(prog="YT DLP Music downloader", description="Download all albums corresponding to songs in a playlist")
    parser.add_argument("playlist")

    args = parser.parse_args()

    logger.debug("Create YT object...")
    albums = []
    yt = YoutubeDL(ydl_opts)
    data = yt.extract_info(args.playlist, download=False)
    i = 0
    if data:
        albums = set()
        for datum in data["entries"]:
            logger.debug(json.dumps(datum, indent=4))
            albums.add(vid_info(datum))
            logger.info(f"ALBUM NAME: {vid_info(datum)}")

        logger.info(f"albums: {albums}")
        re = re.compile("OLAK.*")
        ytmusic = YTMusic()
        album_ids = []
        for album in list(albums):
            results = ytmusic.search(albums.pop())
            for item in results:
                if item['resultType'] == 'album':
                    album_ids.append(item['playlistId'])
                    break
        logger.info(f"Album IDs: {album_ids}")
        tasks = []
        with concurrent.futures.ProcessPoolExecutor() as executor:
            for result in executor.map(download_album, album_ids):
                tasks.append(result)
            logger.info(f"DONE")
        logger.info(f"DONE")