neo-pod-desktop/src/artwork/extractor.py

"""Extract embedded album art from media files using mutagen.

Ported from iOpenPod ArtworkDB_Writer/art_extractor.py.
Supports: MP3, M4A/AAC, M4B, M4V/MP4, FLAC, OGG, OPUS, AIFF/WAV.
"""

from __future__ import annotations

import os
from pathlib import Path

import mutagen


def extract_cover(file_path: str) -> bytes | None:
    """Extract the first embedded album art image from a media file.

    Returns raw image bytes (JPEG/PNG) or None.
    """
    path = Path(file_path)
    ext = path.suffix.lower()

    try:
        if ext in _IMAGE_EXTS:
            return path.read_bytes()
        if ext == ".mp3":
            return _extract_mp3(file_path)
        elif ext in (".m4a", ".m4p", ".m4b", ".aac", ".alac"):
            return _extract_mp4(file_path)
        elif ext in (".m4v", ".mp4", ".mov"):
            art = _extract_mp4(file_path)
            if art:
                return art
            return None
        elif ext == ".flac":
            return _extract_flac(file_path)
        elif ext == ".ogg":
            return _extract_ogg(file_path)
        elif ext == ".opus":
            return _extract_opus(file_path)
        elif ext in (".aif", ".aiff"):
            return _extract_aiff(file_path)
        else:
            return _extract_generic(file_path)
    except Exception:
        return None


def find_folder_art(file_path: str) -> str | None:
    """Return path of folder artwork image next to *file_path*, or None."""
    directory = str(Path(file_path).parent)
    try:
        entries = os.listdir(directory)
    except OSError:
        return None
    lower_map = {e.lower(): e for e in entries}
    for stem in _FOLDER_ART_NAMES:
        for ext in _IMAGE_EXTS:
            if (stem + ext) in lower_map:
                return os.path.join(directory, lower_map[stem + ext])
    return None


def extract_cover_with_fallback(file_path: str) -> bytes | None:
    """Extract cover art, falling back to folder.jpg if no embedded art."""
    art = extract_cover(file_path)
    if art is not None:
        return art
    folder_img = find_folder_art(file_path)
    if folder_img:
        try:
            return Path(folder_img).read_bytes()
        except OSError:
            pass
    return None


_FOLDER_ART_NAMES = ("cover", "folder", "album", "front", "artwork", "thumb")
_IMAGE_EXTS = (".jpg", ".jpeg", ".png", ".bmp", ".gif", ".webp")


def _extract_mp3(path: str) -> bytes | None:
    from mutagen.mp3 import MP3
    audio = MP3(path)
    if audio.tags is None:
        return None
    for key in audio.tags:
        if key.startswith("APIC"):
            return audio.tags[key].data
    return None


def _extract_mp4(path: str) -> bytes | None:
    from mutagen.mp4 import MP4
    audio = MP4(path)
    if audio.tags is None:
        return None
    covers = audio.tags.get("covr", [])
    if covers:
        return bytes(covers[0])
    return None


def _extract_flac(path: str) -> bytes | None:
    from mutagen.flac import FLAC
    audio = FLAC(path)
    if audio.pictures:
        return audio.pictures[0].data
    return None


def _extract_ogg(path: str) -> bytes | None:
    from mutagen.oggvorbis import OggVorbis
    audio = OggVorbis(path)
    return _extract_vorbis_picture(audio)


def _extract_opus(path: str) -> bytes | None:
    from mutagen.oggopus import OggOpus
    audio = OggOpus(path)
    return _extract_vorbis_picture(audio)


def _extract_vorbis_picture(audio) -> bytes | None:
    import base64
    pictures = audio.get("metadata_block_picture", [])
    if pictures:
        try:
            from mutagen.flac import Picture
            pic = Picture(base64.b64decode(pictures[0]))
            return pic.data
        except Exception:
            pass
    return None


def _extract_aiff(path: str) -> bytes | None:
    from mutagen.aiff import AIFF
    audio = AIFF(path)
    if audio.tags is None:
        return None
    for key in audio.tags:
        if key.startswith("APIC"):
            return audio.tags[key].data
    return None


def _extract_generic(path: str) -> bytes | None:
    audio = mutagen.File(path)
    if audio is None or audio.tags is None:
        return None
    for key in audio.tags:
        if hasattr(key, "startswith") and key.startswith("APIC"):
            frame = audio.tags[key]
            if hasattr(frame, "data"):
                return frame.data
    covers = audio.tags.get("covr", [])
    if covers:
        return bytes(covers[0])
    return None