MTS-music-another-client/tags.py

"""ID3-теги для скачанных треков: максимум метаданных.

Пишем ID3v2.3: название, исполнителей, альбом и альбомного исполнителя,
позицию трека и диска, год и дату, жанр, лейбл, длительность, ссылку,
служебные поля (id в Яндекс.Музыке и MTS, explicit), обложку 600×600 и
текст песни (обычный USLT и синхронный SYLT, если он есть).

Модуль необязательный: без `mutagen` теги просто не пишутся.
"""

from __future__ import annotations

import logging
import re
from pathlib import Path

import requests

from media import track_cover_url
from mtsmusic import (
    USER_AGENT,
    force_ipv4,
    lrclib_lyrics,
    lyrics_response,
    lyrics_text,
)

try:
    from mutagen.id3 import (
        APIC,
        COMM,
        SYLT,
        TALB,
        TCON,
        TDAT,
        TIT2,
        TLEN,
        TPE1,
        TPE2,
        TPOS,
        TPUB,
        TRCK,
        TXXX,
        TYER,
        USLT,
        WOAR,
        ID3,
    )

    HAVE_MUTAGEN = True
except Exception:  # pragma: no cover - зависит от окружения
    HAVE_MUTAGEN = False

log = logging.getLogger("tags")

COVER_SIZE = 600
_LRC_LINE = re.compile(r"\[(\d{1,2}):(\d{2})(?:[.:](\d{1,3}))?\]")


def _names(items) -> str:
    return ", ".join(str(item.get("name") or "") for item in items or [])


def fetch_cover(track: dict) -> bytes | None:
    """Скачивает обложку 600×600 для встраивания в APIC."""
    url = track_cover_url(track, COVER_SIZE)
    if not url:
        return None
    try:
        force_ipv4()
        response = requests.get(
            url, headers={"User-Agent": USER_AGENT}, timeout=30)
        response.raise_for_status()
        return response.content
    except requests.RequestException:
        return None


def _mts_lyrics(api, track_id: str, fmt: str) -> str:
    try:
        response = lyrics_response(api, track_id, fmt)
    except Exception:
        return ""
    if response.status_code != 200:
        return ""
    try:
        text = lyrics_text(api, response)
    except Exception:
        return ""
    if "Подключите подписку" in text:
        return ""
    return text.strip()


def fetch_lyrics(api, track: dict) -> tuple[str, str]:
    """(обычный текст, синхронный LRC). Пробуем MTS, затем LRCLIB."""
    track_id = str(track.get("id"))
    plain = _mts_lyrics(api, track_id, "TEXT")
    synced = _mts_lyrics(api, track_id, "LRC")
    if not plain:
        plain = (lrclib_lyrics(track, want_sync=False) or "").strip()
    if not synced:
        synced = (lrclib_lyrics(track, want_sync=True) or "").strip()
    if synced and not _LRC_LINE.search(synced):
        synced = ""
    return plain, synced


def parse_lrc(text: str) -> list[tuple[str, int]]:
    """LRC → [(строка, миллисекунды)] для синхронного SYLT."""
    entries: list[tuple[str, int]] = []
    for line in text.splitlines():
        stamps = list(_LRC_LINE.finditer(line))
        if not stamps:
            continue
        content = _LRC_LINE.sub("", line).strip()
        for stamp in stamps:
            minutes = int(stamp.group(1))
            seconds = int(stamp.group(2))
            fraction = (stamp.group(3) or "0").ljust(3, "0")[:3]
            millis = (minutes * 60 + seconds) * 1000 + int(fraction)
            entries.append((content, millis))
    entries.sort(key=lambda item: item[1])
    return entries


def _cover_mime(data: bytes) -> str:
    return "image/png" if data[:8].startswith(b"\x89PNG") else "image/jpeg"


def write_tags(
    track: dict,
    path: Path,
    plain_lyrics: str = "",
    synced_lyrics: str = "",
    cover: bytes | None = None,
) -> None:
    """Проставляет ID3v2.3-теги в уже скачанном файле."""
    if not HAVE_MUTAGEN:
        return
    album = (track.get("albums") or [{}])[0]
    artist = _names(track.get("artists"))
    album_artist = _names(album.get("artists")) or artist
    track_id = str(track.get("id") or "")

    try:
        tags = ID3(path)
    except Exception:
        tags = ID3()

    def add(frame) -> None:
        tags.add(frame)

    if track.get("title"):
        add(TIT2(encoding=3, text=str(track["title"])))
    if artist:
        add(TPE1(encoding=3, text=artist))
    if album_artist:
        add(TPE2(encoding=3, text=album_artist))
    if album.get("title"):
        add(TALB(encoding=3, text=str(album["title"])))

    position = album.get("trackPosition") or {}
    index = position.get("index")
    if index:
        count = album.get("trackCount")
        add(TRCK(encoding=3, text=f"{index}/{count}" if count else str(index)))
    volume = position.get("volume")
    if volume:
        add(TPOS(encoding=3, text=str(volume)))

    if album.get("year"):
        add(TYER(encoding=3, text=str(album["year"])))
    release = str(album.get("releaseDate") or "")
    if release[:10].count("-") == 2:
        _year, month, day = release[:10].split("-")
        add(TDAT(encoding=3, text=f"{day}{month}"))  # формат DDMM

    if album.get("genre"):
        add(TCON(encoding=3, text=str(album["genre"])))
    labels = _names(album.get("labels"))
    if labels:
        add(TPUB(encoding=3, text=labels))
    if track.get("durationMs"):
        add(TLEN(encoding=3, text=str(int(track["durationMs"]))))

    if track_id:
        add(COMM(
            encoding=3,
            lang="XXX",
            desc="",
            text=f"https://music.yandex.ru/track/{track_id}",
        ))
        add(WOAR(url=f"https://music.yandex.ru/track/{track_id}"))

    def add_text(desc: str, value) -> None:
        if value not in (None, ""):
            add(TXXX(encoding=3, desc=desc, text=str(value)))

    add_text("YANDEX_MUSIC_TRACK_ID", track.get("id"))
    add_text("YANDEX_MUSIC_REAL_ID", track.get("realId"))
    add_text("YANDEX_MUSIC_ALBUM_ID", album.get("id"))
    add_text("MTS_MUSIC_SOURCE", track.get("trackSource"))
    add_text("YANDEX_MUSIC_TYPE", track.get("type"))
    if track.get("contentWarning") == "explicit":
        add_text("ITUNESADVISORY", 1)

    if cover:
        add(APIC(
            encoding=3,
            mime=_cover_mime(cover),
            type=3,
            desc="Cover",
            data=cover,
        ))
    if plain_lyrics:
        add(USLT(encoding=3, lang="XXX", desc="", text=plain_lyrics))
    if synced_lyrics:
        entries = parse_lrc(synced_lyrics)
        if entries:
            add(SYLT(
                encoding=3,
                lang="XXX",
                format=2,
                type=1,
                desc="",
                text=entries,
            ))

    tags.save(path, v2_version=3)


def tag_track(api, track: dict, path: Path) -> None:
    """Собирает метаданные (текст, обложка) и пишет их в файл."""
    if not HAVE_MUTAGEN:
        return
    plain, synced = fetch_lyrics(api, track)
    cover = fetch_cover(track)
    write_tags(track, Path(path), plain, synced, cover)