utils.py

from mutagen.id3 import ID3, TIT2, TPE1, TALB, TDRC, TCON, TRCK, TPOS, TCOM, TPE2, COMM, USLT, TBPM, TSRC, APIC
from mutagen.flac import FLAC, Picture
from mutagen.oggvorbis import OggVorbis
from mutagen.mp4 import MP4, MP4Cover
from mutagen.wave import WAVE
from mutagen.oggopus import OggOpus
from mutagen.asf import ASF, ASFByteArrayAttribute
from mutagen.aiff import AIFF
from base64 import b64encode
import os
import subprocess
import struct


def printHelp():
    print("""
  Manual mode:
    Select a single file or a folder of audio files, then fill in each
    metadata field (title, artist, album, year, genre, etc.) by hand.
    In album/batch mode shared fields are entered once and only the
    title and track number are prompted per file.
    You can optionally convert files to another format (requires FFmpeg).

  Automatic mode:
    Search TheAudioDB by artist and album name. Choose whether you are
    tagging a single song or a full album. The retrieved metadata
    (including cover art) is previewed before being applied.
    For albums, each file is matched to a track from the tracklist.

  Supported formats:
    MP3, FLAC, OGG, Opus, MP4, M4A, M4B, M4P, M4V, 3GP, ASF, WMA,
    WMV, AIF, AIFF, WAV

  Notes:
    - Cover art must be a JPEG (.jpg) file.
    - Lyrics are read from a plain text file.
    - WAV tagging and format conversion require FFmpeg on PATH.
""")


_EXTENSION_HANDLERS = {
    ".mp3":  lambda p, m: setID3Metadata(p, m),
    ".flac": lambda p, m: setFlacMetadata(p, m),
    ".ogg":  lambda p, m: setOggMetadata(p, m),
    ".opus": lambda p, m: setOpusMetadata(p, m),
    ".wav":  lambda p, m: setWavMetadata(p, m),
    ".aif":  lambda p, m: setAIFMetadata(p, m),
    ".aiff": lambda p, m: setAIFMetadata(p, m),
    ".asf":  lambda p, m: setASFMetadata(p, m),
    ".wma":  lambda p, m: setASFMetadata(p, m),
    ".wmv":  lambda p, m: setASFMetadata(p, m),
    ".mp4":  lambda p, m: setMP4Metadata(p, m),
    ".m4a":  lambda p, m: setMP4Metadata(p, m),
    ".m4b":  lambda p, m: setMP4Metadata(p, m),
    ".m4p":  lambda p, m: setMP4Metadata(p, m),
    ".m4v":  lambda p, m: setMP4Metadata(p, m),
    ".3gp":  lambda p, m: setMP4Metadata(p, m),
}


def setMetadata(extension, path, metadata):
    handler = _EXTENSION_HANDLERS.get(extension)
    if handler:
        handler(path, metadata)


# --- ID3 (MP3, WAV, AIF/AIFF) ---

def setID3Tags(audio_file, metadata):
    if "title" in metadata:
        audio_file["TIT2"] = TIT2(encoding=3, text=metadata["title"])
    if "artist" in metadata:
        audio_file["TPE1"] = TPE1(encoding=3, text=metadata["artist"])
    if "album" in metadata:
        audio_file["TALB"] = TALB(encoding=3, text=metadata["album"])
    if "year" in metadata:
        audio_file["TDRC"] = TDRC(encoding=3, text=str(metadata["year"]))
    if "genre" in metadata:
        audio_file["TCON"] = TCON(encoding=3, text=metadata["genre"])

    if "track" in metadata or "tracks" in metadata:
        track = metadata.get("track", "")
        total = metadata.get("tracks", "")
        audio_file["TRCK"] = TRCK(encoding=3, text=f"{track}/{total}".strip("/"))

    if "disk" in metadata or "disks" in metadata:
        disk = metadata.get("disk", "")
        total = metadata.get("disks", "")
        audio_file["TPOS"] = TPOS(encoding=3, text=f"{disk}/{total}".strip("/"))

    if "compositor" in metadata:
        audio_file["TCOM"] = TCOM(encoding=3, text=metadata["compositor"])
    if "albumArtist" in metadata:
        audio_file["TPE2"] = TPE2(encoding=3, text=metadata["albumArtist"])
    if "comments" in metadata:
        audio_file["COMM"] = COMM(encoding=3, lang="eng", desc="desc", text=metadata["comments"])
    if "lyrics" in metadata:
        audio_file["USLT"] = USLT(encoding=3, lang="eng", desc="lyrics", text=metadata["lyrics"])
    if "bpm" in metadata:
        audio_file["TBPM"] = TBPM(encoding=3, text=str(metadata["bpm"]))
    if "isrc" in metadata:
        audio_file["TSRC"] = TSRC(encoding=3, text=metadata["isrc"])

    # Cover art
    if "image" in metadata and metadata["image"]:
        with open(metadata["image"], "rb") as img:
            audio_file["APIC"] = APIC(
                encoding=3,
                mime="image/jpeg",
                type=3,
                desc="Cover",
                data=img.read(),
            )


def setID3Metadata(path, metadata):
    audio_file = ID3(path)
    setID3Tags(audio_file, metadata)
    audio_file.save(path)


def setWavMetadata(path, metadata):
    # Rename temporarily, run through FFmpeg to get a WAV structure Mutagen can tag
    pre_path = os.path.join(os.path.dirname(path), "pre_" + os.path.basename(path))
    os.rename(path, pre_path)
    touchWav(pre_path, path)
    os.remove(pre_path)

    audio = WAVE(path)
    if audio.tags is None:
        audio.add_tags()
        audio.save()

    setID3Tags(audio, metadata)
    audio.save()


def setAIFMetadata(path, metadata):
    audio_file = AIFF(path)
    setID3Tags(audio_file, metadata)
    audio_file.save(path)


def touchWav(inputPath, outputPath):
    subprocess.run(["ffmpeg", "-i", inputPath, outputPath], check=True)


# --- Vorbis (FLAC, OGG, Opus) ---

simple_vorbis_map = {
    "title": "TITLE",
    "artist": "ARTIST",
    "album": "ALBUM",
    "year": "DATE",
    "genre": "GENRE",
    "isrc": "ISRC",
    "track": "TRACKNUMBER",
    "tracks": "TRACKTOTAL",
    "disk": "DISKNUMBER",
    "disks": "DISKTOTAL",
}


def setFlacMetadata(path, metadata):
    audio = FLAC(path)

    for key, vorbis_key in simple_vorbis_map.items():
        if key in metadata and metadata[key]:
            audio[vorbis_key] = str(metadata[key])

    if "image" in metadata and metadata["image"]:
        picture = Picture()
        picture.type = 3
        picture.mime = "image/jpeg"
        with open(metadata["image"], "rb") as img:
            picture.data = img.read()
        audio.clear_pictures()
        audio.add_picture(picture)

    audio.save()


def setOggTags(oggFile, metadata):
    if oggFile.tags is None:
        oggFile.add_tags()

    for key, vorbis_key in simple_vorbis_map.items():
        if key in metadata and metadata[key]:
            oggFile[vorbis_key] = str(metadata[key])

    if "image" in metadata and metadata["image"]:
        picture = Picture()
        picture.type = 3
        picture.mime = "image/jpeg"
        with open(metadata["image"], "rb") as img:
            picture.data = img.read()
            oggFile["METADATA_BLOCK_PICTURE"] = [b64encode(picture.write()).decode("ascii")]


def setOggMetadata(path, metadata):
    audio = OggVorbis(path)
    setOggTags(audio, metadata)
    audio.save()


def setOpusMetadata(path, metadata):
    audio = OggOpus(path)
    setOggTags(audio, metadata)
    audio.save()


# --- MP4 (MP4, M4A, M4B, M4P, M4V, 3GP) ---

simple_mp4_map = {
    "title": "\xa9nam",
    "artist": "\xa9ART",
    "album": "\xa9alb",
    "year": "\xa9day",
    "genre": "\xa9gen",
    "track": "trkn",
    "disk": "disk",
    "compositor": "\xa9wrt",
    "albumArtist": "aART",
    "comments": "\xa9cmt",
    "bpm": "tmpo",
    "isrc": "\xa9isrc",
}


def setMP4Metadata(path, metadata):
    audio = MP4(path)

    for key, mp4_key in simple_mp4_map.items():
        if key in metadata and metadata[key]:
            if key == "track":
                audio["trkn"] = [(int(metadata["track"]), int(metadata.get("tracks", 0)))]
            elif key == "disk":
                audio["disk"] = [(int(metadata["disk"]), int(metadata.get("disks", 0)))]
            else:
                audio[mp4_key] = [metadata[key]]

    if "image" in metadata and metadata["image"]:
        with open(metadata["image"], "rb") as img:
            cover = MP4Cover(img.read(), imageformat=MP4Cover.FORMAT_JPEG)
            audio["covr"] = [cover]

    audio.save()


# --- ASF (ASF, WMA, WMV) ---

asf_map = {
    "title": "Title",
    "artist": "Author",
    "album": "WM/AlbumTitle",
    "genre": "WM/Genre",
    "year": "WM/Year",
    "compositor": "WM/Composer",
    "albumArtist": "WM/AlbumArtist",
    "comments": "Description",
    "bpm": "WM/BeatsPerMinute",
    "isrc": "WM/ISRC",
    "track": "WM/TrackNumber",
}


def setASFMetadata(path, metadata):
    audio = ASF(path)

    for key, asf_key in asf_map.items():
        if key in metadata and metadata[key]:
            audio[asf_key] = metadata[key]

    audio["WM/MediaClassPrimaryID"] = "{D1607DDE-7B67-4B88-9E0A-FB37E2049096}"
    audio["WM/MediaClassSecondaryID"] = "{00000000-0000-0000-0000-000000000000}"
    audio["WM/ContentType"] = "audio"

    if "image" in metadata and metadata["image"]:
        with open(metadata["image"], "rb") as img:
            image_data = img.read()
            type_byte = struct.pack("<b", 3)
            size_bytes = struct.pack("<i", len(image_data))
            mime_utf16 = "image/jpeg".encode("utf-16-le") + b"\x00\x00"
            desc_utf16 = b"\x00\x00"
            packed = type_byte + size_bytes + mime_utf16 + desc_utf16 + image_data
            audio["WM/Picture"] = [ASFByteArrayAttribute(packed)]

    audio.save()