from mutagen import mp3, id3
from yt_dlp.postprocessor import PostProcessor


ENC_UTF8 = 3


def add_tags(file: mp3.MP3, info: dict[str, str | int]) -> None:

    file['TIT2'] = id3.TIT2(encoding=ENC_UTF8, text=info['track'])
    file['TPE1'] = id3.TPE1(encoding=ENC_UTF8, text=info['artist'])
    if 'album' in info:
        file['TALB'] = id3.TALB(encoding=ENC_UTF8, text=info['album'])
    if 'release_year' in info:
        file['TDRC'] = id3.TDRC(encoding=ENC_UTF8, text=str(info['release_year']))
    if 'track_number' in info:
        file['TRCK'] = id3.TRCK(encoding=ENC_UTF8, text=str(info['track_number']))


class InfoYouTubePP(PostProcessor):
    '''Generates YT Music fields in info if necessary.
    Must be run before downloading and post-processing with
    FFmpegExtractAudioPP and ID3YouTubePP, so use only with
    when<="before_dl" ("pre_process" also suits, see yt_dlp.utils.POSTPROCESS_WHEN)
    for correct file path and ID3 tags'''

    def run(self, information):

        if not 'track' in information:
            information['track'] = information['title']
        if not 'artist' in information:
            information['artist'] = information['channel'].removesuffix(' - Topic')

        if not 'album' in information:
            try:
                information['album'] = information['playlist'].removeprefix('Album \u2013 ')
            except KeyError:
                pass

        if not 'track_number' in information:
            try:
                information['track_number'] = information['playlist_index']
            except KeyError:
                pass
        # here was some code for disc_number,
        # but it is even not in mutagen ID3 fields list

        return [], information


class ID3YouTubePP(PostProcessor):
    '''Inserts ID3 tags after InfoYouTubePP and FFmpegExtractAudioPP,
    triggers searching and parsing lyrics from Genius'''

    def run(self, information):

        add_tags(
            mp3.MP3(information['filepath']),
            information,
        )

        return [], information