Lyra/worker/lyra_worker/_musicbrainz.py

import time
from difflib import SequenceMatcher

from lyra_worker.types import MBTarget

_MIN_TITLE_RATIO = 0.6


def _with_retry(fn, *, attempts: int = 4, base_delay: float = 1.0, sleep=time.sleep):
    """Call `fn`, retrying transient MusicBrainz failures with exponential backoff.

    A dropped TLS connection to the MusicBrainz webservice surfaces as
    musicbrainzngs.NetworkError (`<urlopen error [SSL: UNEXPECTED_EOF_WHILE_READING]>`),
    and MusicBrainz's own 503 rate-limit surfaces as ResponseError; both are
    transient, so we retry rather than failing the whole pipeline for the job.
    Any other exception (bug, non-transient) propagates immediately.

    musicbrainzngs is imported lazily so importing this module stays offline.
    """
    import musicbrainzngs

    transient = (musicbrainzngs.NetworkError, musicbrainzngs.ResponseError)
    for i in range(attempts):
        try:
            return fn()
        except transient:
            if i == attempts - 1:
                raise
            sleep(base_delay * (2**i))


def _ratio(a: str, b: str) -> float:
    return SequenceMatcher(None, a.strip().casefold(), b.strip().casefold()).ratio()


def _is_album(g: dict) -> bool:
    return (g.get("primary-type") or "").casefold() == "album"


def _rg_year(g: dict) -> int:
    """First-release year of a release-group, or 9999 when unknown (so undated groups lose ties)."""
    frd = (g.get("first-release-date") or "")[:4]
    return int(frd) if frd.isdigit() else 9999


def _credited_artist(g: dict) -> str:
    """Primary credited artist name of a release-group search result, or '' if absent."""
    credit = g.get("artist-credit")
    if isinstance(credit, list) and credit and isinstance(credit[0], dict):
        return (credit[0].get("artist") or {}).get("name", "") or ""
    return ""


def _best_release_group(artist: str, album: str, groups: list) -> dict | None:
    """Choose the best release-group match for `artist`/`album`.

    Ranked, highest first, on `(artist_ratio, title_ratio, is_album)`:
      * artist match dominates — MusicBrainz relevance ties same-titled release groups by
        *different* artists (e.g. "Fun Machine" exists as a Lake Street Dive EP and an
        unrelated band's Album), and the artist we followed is trustworthy, so a credited
        artist that matches the request outranks everything else;
      * title similarity is next;
      * an Album primary-type only breaks ties *within* the same artist+title, so a famous
        album (Michael Jackson's "Off the Wall") still isn't resolved to its same-named
        single whose short tracklist would map positionally onto the album's files;
      * finally, the EARLIEST release-group wins — MusicBrainz lists the original album and its
        later reissues/anniversary editions as separate same-titled groups (Linkin Park "Hybrid
        Theory": a 2000 original and a 2023 13-track reissue), which tie on everything above, so
        MB's arbitrary relevance order used to pick a reissue whose longer tracklist/runtime then
        failed completeness checks against the standard album a source delivers. Prefer the original.
    Artist is a ranking signal, never a hard filter — a low match sinks a candidate but
    never drops the release, so credited-name variations (feat., punctuation) still resolve.
    Returns None below the title threshold. Pure — no network I/O."""
    if not groups:
        return None
    best = max(
        groups,
        key=lambda g: (
            _ratio(artist, _credited_artist(g)),
            _ratio(album, g.get("title", "")),
            _is_album(g),
            -_rg_year(g),  # tie-break: the original (earliest) group over later reissues
        ),
    )
    if _ratio(album, best.get("title", "")) < _MIN_TITLE_RATIO:
        return None
    return best


def _pick_release(releases: list) -> dict | None:
    """Choose which release of a release-group to read the tracklist from. MB returns them in an
    arbitrary order and ``releases[0]`` is often a later reissue/deluxe/anniversary edition with
    more tracks and a longer runtime than the standard album a source actually delivers — which
    then makes a *complete* download fail the pipeline's track-count/duration completeness checks
    and re-download forever (Linkin Park "Hybrid Theory": MB's arbitrary pick was a 2023 13-track
    reissue at ~52 min vs the 2000 original's 12 tracks / ~38 min). Prefer the ORIGINAL standard
    edition: an Official release, then the earliest date, then a dated one. Pure — unit-tested."""
    if not releases:
        return None
    def _key(r: dict) -> tuple:
        official = 0 if (r.get("status") == "Official") else 1  # official editions first
        date = r.get("date") or ""
        return (official, 0 if date else 1, date)  # then dated, then earliest date
    return min(releases, key=_key)


class MusicBrainzResolver:
    """Real MbResolver using the MusicBrainz webservice via musicbrainzngs.

    NOT unit-tested offline; see test_musicbrainz_live.py. musicbrainzngs is
    imported lazily so importing/constructing this class stays offline.
    """

    def __init__(self, app_name: str = "Lyra", version: str = "0.1", contact: str = "lyra@localhost"):
        self._app = app_name
        self._version = version
        self._contact = contact

    def resolve(self, artist: str, album: str) -> MBTarget | None:
        import musicbrainzngs

        musicbrainzngs.set_useragent(self._app, self._version, self._contact)

        res = _with_retry(lambda: musicbrainzngs.search_release_groups(query=album, artist=artist, limit=5))
        groups = res.get("release-group-list", [])
        rg = _best_release_group(artist, album, groups)
        if rg is None:
            return None

        canonical_album = rg.get("title", album)
        canonical_artist = artist
        artist_mbid = ""
        credit = rg.get("artist-credit")
        if isinstance(credit, list) and credit and isinstance(credit[0], dict):
            _artist = credit[0].get("artist") or {}
            canonical_artist = _artist.get("name", artist)
            artist_mbid = _artist.get("id", "") or ""

        year = None
        frd = rg.get("first-release-date", "") or ""
        if frd[:4].isdigit():
            year = int(frd[:4])

        track_count = None
        total_duration_s = None
        titles: tuple[str, ...] = ()
        rgid = rg.get("id")
        if rgid:
            rgfull = _with_retry(lambda: musicbrainzngs.get_release_group_by_id(rgid, includes=["releases"]))
            releases = rgfull.get("release-group", {}).get("release-list", [])
            chosen = _pick_release(releases)
            if chosen:
                rel = _with_retry(
                    lambda: musicbrainzngs.get_release_by_id(chosen["id"], includes=["recordings"])
                ).get("release", {})
                tracks = []
                for medium in rel.get("medium-list", []):
                    tracks.extend(medium.get("track-list", []))
                if tracks:
                    track_count = len(tracks)
                    total_ms = sum(int((t.get("recording") or {}).get("length") or 0) for t in tracks)
                    total_duration_s = total_ms // 1000 if total_ms else None
                    titles = tuple((t.get("recording") or {}).get("title", "") for t in tracks)

        return MBTarget(
            artist=canonical_artist,
            album=canonical_album,
            track_count=track_count,
            total_duration_s=total_duration_s,
            year=year,
            tracklist=titles,
            rg_mbid=rg.get("id", "") or "",
            artist_mbid=artist_mbid,
        )