lintunes/lintunes/tagging.py

"""Read and write audio file tags (mp3/m4a/flac) via mutagen.

A small uniform field dict is used throughout:
name, artist, album_artist, album, genre, composer, grouping, comments,
year (int), track_number, track_count, disc_number, disc_count,
compilation (bool), total_time (ms, read-only), bit_rate, sample_rate, size.
"""
from pathlib import Path

import mutagen
from mutagen.easyid3 import EasyID3
from mutagen.id3 import ID3, COMM, GRP1, TCMP, TBPM
from mutagen.mp3 import MP3
from mutagen.mp4 import MP4
from mutagen.flac import FLAC

from lintunes.perf import timed


# EasyID3 has no 'comment' key by default; COMM frames need explicit handling.
def _comm_get(id3, _key):
    frame = id3.get("COMM::eng") or id3.get("COMM::XXX")
    return list(frame.text) if frame else []


def _comm_set(id3, _key, value):
    id3.delall("COMM")
    id3.add(COMM(encoding=3, lang="eng", desc="", text=value))


def _comm_delete(id3, _key):
    id3.delall("COMM")


EasyID3.RegisterKey("comment", _comm_get, _comm_set, _comm_delete)


# Register grouping/compilation/bpm as Easy keys (iTunes-style frames) so
# write_tags can set every field on one object and save the file ONCE —
# a second parse+save per edit is real I/O on big files / slow disks.
def _register_id3_frame_key(key: str, frame_cls):
    frame_id = frame_cls.__name__

    def getter(id3, _key):
        frame = id3.get(frame_id)
        return list(frame.text) if frame else []

    def setter(id3, _key, value):
        id3.delall(frame_id)
        id3.add(frame_cls(encoding=3, text=value))

    def deleter(id3, _key):
        id3.delall(frame_id)

    EasyID3.RegisterKey(key, getter, setter, deleter)


_register_id3_frame_key("grouping", GRP1)      # iTunes GRP1, not TIT1
_register_id3_frame_key("compilation", TCMP)
_register_id3_frame_key("bpm", TBPM)

# EasyMP4 is missing 'composer' (©wrt) out of the box
from mutagen.easymp4 import EasyMP4Tags  # noqa: E402
if "composer" not in EasyMP4Tags.Set:
    EasyMP4Tags.RegisterTextKey("composer", "\xa9wrt")
if "grouping" not in EasyMP4Tags.Set:
    EasyMP4Tags.RegisterTextKey("grouping", "\xa9grp")
if "bpm" not in EasyMP4Tags.Set:
    EasyMP4Tags.RegisterIntKey("bpm", "tmpo")


def _cpil_get(tags, _key):
    return ["1"] if tags.get("cpil") else []


def _cpil_set(tags, _key, value):
    tags["cpil"] = bool(value and str(value[0]) not in ("", "0"))


def _cpil_delete(tags, _key):
    tags.pop("cpil", None)


if "compilation" not in EasyMP4Tags.Set:
    EasyMP4Tags.RegisterKey("compilation", _cpil_get, _cpil_set, _cpil_delete)

# Easy-interface tag names shared by EasyID3 / EasyMP4 / FLAC (Vorbis)
_TEXT_FIELDS = {
    "name": "title",
    "artist": "artist",
    "album_artist": "albumartist",
    "album": "album",
    "genre": "genre",
    "composer": "composer",
    "comments": "comment",
}

EDITABLE_FIELDS = list(_TEXT_FIELDS) + [
    "grouping", "year", "track_number", "track_count",
    "disc_number", "disc_count", "compilation", "bpm",
]


def read_tags(path: str | Path) -> dict:
    """Read tags + audio properties into Track-style fields."""
    audio = mutagen.File(str(path), easy=True)
    if audio is None:
        return {}

    fields = {}
    for our_key, tag_key in _TEXT_FIELDS.items():
        values = audio.tags.get(tag_key) if audio.tags else None
        if values:
            fields[our_key] = str(values[0])

    if audio.tags:
        for src, num_key, count_key in (("tracknumber", "track_number", "track_count"),
                                        ("discnumber", "disc_number", "disc_count")):
            values = audio.tags.get(src)
            if values:
                num, _, count = str(values[0]).partition("/")
                if num.strip().isdigit():
                    fields[num_key] = int(num)
                if count.strip().isdigit():
                    fields[count_key] = int(count)
        date_values = audio.tags.get("date") or audio.tags.get("originaldate")
        if date_values:
            year_part = str(date_values[0])[:4]
            if year_part.isdigit():
                fields["year"] = int(year_part)

    info = audio.info
    if info is not None:
        if getattr(info, "length", 0):
            fields["total_time"] = int(info.length * 1000)
        if getattr(info, "bitrate", 0):
            fields["bit_rate"] = info.bitrate // 1000
        if getattr(info, "sample_rate", 0):
            fields["sample_rate"] = info.sample_rate
    fields["size"] = Path(path).stat().st_size
    fields["kind"] = _kind_for(audio)
    return fields


def write_tags(path: str | Path, fields: dict):
    """Write the given Track-style fields to the file's tags.

    One parse and ONE save per call: grouping/compilation/bpm are covered by
    the Easy keys registered above, so no second raw-tag pass is needed."""
    path = str(path)
    audio = mutagen.File(path, easy=True)
    if audio is None:
        raise ValueError(f"Unsupported audio file: {path}")
    if audio.tags is None:
        audio.add_tags()

    for our_key, tag_key in _TEXT_FIELDS.items():
        if our_key in fields:
            value = fields[our_key]
            if value:
                audio.tags[tag_key] = [str(value)]
            else:
                audio.tags.pop(tag_key, None)

    if "year" in fields:
        if fields["year"]:
            audio.tags["date"] = [str(fields["year"])]
        else:
            audio.tags.pop("date", None)

    for num_key, count_key, tag_key in (("track_number", "track_count", "tracknumber"),
                                        ("disc_number", "disc_count", "discnumber")):
        if num_key in fields or count_key in fields:
            num = fields.get(num_key, 0)
            count = fields.get(count_key, 0)
            if num and count:
                audio.tags[tag_key] = [f"{num}/{count}"]
            elif num:
                audio.tags[tag_key] = [str(num)]
            else:
                audio.tags.pop(tag_key, None)

    if "grouping" in fields:
        _set_or_pop(audio.tags, "grouping",
                    str(fields["grouping"]) if fields["grouping"] else None)
    if "bpm" in fields:
        _set_or_pop(audio.tags, "bpm",
                    str(fields["bpm"]) if fields["bpm"] else None)
    if "compilation" in fields:
        _set_or_pop(audio.tags, "compilation",
                    "1" if fields["compilation"] else None)

    with timed("tag save %s", path):
        audio.save()


def _set_or_pop(tags, key, value):
    if value:
        tags[key] = [value]
        return
    try:
        del tags[key]  # FLAC's VCommentDict has no pop-with-default
    except KeyError:
        pass


def _kind_for(audio) -> str:
    name = type(audio).__name__
    return {
        "EasyMP3": "MPEG audio file",
        "MP3": "MPEG audio file",
        "EasyMP4": "AAC audio file",
        "MP4": "AAC audio file",
        "FLAC": "FLAC audio file",
    }.get(name, name)


def write_artwork(path: str | Path, image_bytes: bytes, mime: str = "image/jpeg"):
    """Replace the file's embedded cover art (front cover) with image_bytes."""
    from mutagen.id3 import APIC
    from mutagen.mp4 import MP4Cover
    from mutagen.flac import Picture

    audio = mutagen.File(str(path))
    if audio is None:
        raise ValueError(f"Unsupported audio file: {path}")
    if isinstance(audio, MP4):
        fmt = (MP4Cover.FORMAT_PNG if mime == "image/png"
               else MP4Cover.FORMAT_JPEG)
        if audio.tags is None:
            audio.add_tags()
        audio.tags["covr"] = [MP4Cover(image_bytes, imageformat=fmt)]
    elif isinstance(audio, FLAC):
        picture = Picture()
        picture.type = 3  # front cover
        picture.mime = mime
        picture.desc = "Cover"
        picture.data = image_bytes
        audio.clear_pictures()
        audio.add_picture(picture)
    elif hasattr(audio, "tags") and (isinstance(audio, MP3) or
                                     isinstance(audio.tags, ID3)):
        if audio.tags is None:
            audio.add_tags()
        audio.tags.delall("APIC")
        audio.tags.add(APIC(encoding=3, mime=mime, type=3, desc="Cover",
                            data=image_bytes))
    else:
        raise ValueError(f"Don't know how to embed artwork in {path}")
    with timed("artwork save %s", path):
        audio.save()


def read_embedded_artwork(path: str | Path) -> bytes | None:
    """Return the first embedded cover image's bytes, or None."""
    try:
        audio = mutagen.File(str(path))
    except Exception:
        return None
    if audio is None:
        return None
    if isinstance(audio, MP4):
        covers = audio.tags.get("covr") if audio.tags else None
        if covers:
            return bytes(covers[0])
    elif isinstance(audio, FLAC):
        if audio.pictures:
            return audio.pictures[0].data
    elif audio.tags is not None:
        # ID3 (mp3, aiff): APIC frames
        for key in audio.tags.keys():
            if key.startswith("APIC"):
                return audio.tags[key].data
    return None