Lyra/worker/lyra_worker/adapters/_streamrip.py

import asyncio
import hashlib
import http.client
import os
from typing import Callable

# Qobuz's CDN returns responses with more than Python's default 100-header limit for some
# tracks, so http.client raises HTTPException("got more than 100 headers"), streamrip skips
# that track, and it lands as a 0-byte file (seen on Bruno Mars "24K Magic" — track 1). Raise
# the limit process-wide so those downloads succeed. (Set on import of the Qobuz adapter.)
http.client._MAXHEADERS = 1000


def _hashed_password(password: str) -> str:
    """streamrip's Qobuz email/password login expects the password as an MD5 hex digest."""
    return hashlib.md5(password.encode("utf-8")).hexdigest()


_AUDIO_EXT = {".flac", ".mp3", ".m4a", ".opus", ".ogg", ".aac", ".wav"}

# streamrip holds module-global asyncio primitives (a download-concurrency Semaphore) that
# bind to the running event loop on first use. `asyncio.run` creates a NEW loop per call and
# closes it, orphaning that Semaphore, so the 2nd+ download fails with "Semaphore ... is bound
# to a different event loop" and its tracks silently fail (-> incomplete download). Running every
# streamrip coroutine on ONE persistent loop keeps those globals valid across downloads. The
# worker processes one job at a time, so a single shared loop is safe.
_loop: "asyncio.AbstractEventLoop | None" = None


def _run(coro):
    global _loop
    if _loop is None or _loop.is_closed():
        _loop = asyncio.new_event_loop()
    return _loop.run_until_complete(coro)


def _flatten_audio(dest: str) -> int:
    """streamrip writes into a nested `Artist - Album [...]/` subfolder. Move any audio
    files up into `dest`, remove the now-empty subfolders, and return the count of audio
    files directly in `dest` — the REAL number downloaded, so a partial download (streamrip
    continues past per-track failures) reports its true short count to the completeness check.
    """
    dest = os.path.abspath(dest)
    for root, _dirs, files in os.walk(dest):
        if os.path.abspath(root) == dest:
            continue
        for name in files:
            if os.path.splitext(name)[1].lower() not in _AUDIO_EXT:
                continue
            src = os.path.join(root, name)
            target = os.path.join(dest, name)
            if os.path.exists(target):
                base, ext = os.path.splitext(name)
                target = os.path.join(dest, f"{base} ({os.path.basename(root)}){ext}")
            os.rename(src, target)
    for root, _dirs, _files in os.walk(dest, topdown=False):
        if os.path.abspath(root) != dest:
            try:
                os.rmdir(root)
            except OSError:
                pass  # subfolder still has non-audio files (e.g. cover art) — leave it
    return sum(1 for f in os.listdir(dest) if os.path.splitext(f)[1].lower() in _AUDIO_EXT)


class StreamripClient:
    """Real Qobuz client backed by streamrip (v2, async → sync via asyncio.run).

    NOT unit-tested offline; see test_qobuz_live.py. streamrip is imported lazily
    inside the async helpers so importing/constructing this class stays offline.
    """

    def __init__(self, config: dict):
        self._email = config.get("qobuz.email", "")
        self._password = config.get("qobuz.password", "")
        self._user_id = config.get("qobuz.user_id", "")
        self._token = config.get("qobuz.auth_token", "")

    def is_configured(self) -> bool:
        # either a user_id + auth token, or an email + password
        return bool((self._user_id and self._token) or (self._email and self._password))

    def _make_config(self, download_folder: str | None = None):
        from streamrip.config import Config

        cfg = Config.defaults()
        q = cfg.session.qobuz
        if self._user_id and self._token:
            # token auth: more reliable than email/password, which Qobuz often rejects
            q.use_auth_token = True
            q.email_or_userid = self._user_id
            q.password_or_token = self._token
        else:
            q.use_auth_token = False
            q.email_or_userid = self._email
            q.password_or_token = _hashed_password(self._password)
        q.quality = 3
        if download_folder is not None:
            cfg.session.downloads.folder = download_folder
        return cfg

    def search_album(self, artist: str, album: str) -> list[dict]:
        return _run(self._search(f"{artist} {album}"))

    async def _search(self, query: str) -> list[dict]:
        from streamrip.client import QobuzClient

        client = QobuzClient(self._make_config())
        try:
            try:
                await client.login()
            except Exception as e:  # never let streamrip's credential-bearing error escape
                raise RuntimeError(f"Qobuz login failed ({type(e).__name__}) — check credentials") from None
            pages = await client.search("album", query, limit=5)
        finally:
            await client.session.close()
        results: list[dict] = []
        for page in pages:
            for item in page.get("albums", {}).get("items", []):
                results.append(
                    {
                        "source_ref": str(item.get("id", "")),
                        "title": item.get("title", ""),
                        "artist": (item.get("artist") or {}).get("name", ""),
                        "track_count": item.get("tracks_count", 0) or 0,
                        "bit_depth": item.get("maximum_bit_depth", 16) or 16,
                        "sampling_rate": item.get("maximum_sampling_rate", 44.1) or 44.1,
                    }
                )
        return results

    def download(self, source_ref: str, dest: str, on_progress: Callable[[float], None]) -> dict:
        _run(self._download(source_ref, dest, on_progress))
        # Count the files actually written and flatten streamrip's nested output into dest
        # so the tag stage can organize it. A partial download -> short count -> needs_attention.
        return {"track_count": _flatten_audio(dest), "path": dest}

    async def _download(self, album_id: str, dest: str, on_progress: Callable[[float], None]) -> None:
        os.makedirs(dest, exist_ok=True)
        from streamrip.client import QobuzClient
        from streamrip.db import Database, Dummy
        from streamrip.media import PendingAlbum

        cfg = self._make_config(download_folder=dest)
        client = QobuzClient(cfg)
        try:
            try:
                await client.login()
            except Exception as e:  # never let streamrip's credential-bearing error escape
                raise RuntimeError(f"Qobuz login failed ({type(e).__name__}) — check credentials") from None
            on_progress(0.0)
            db = Database(downloads=Dummy(), failed=Dummy())
            pending = PendingAlbum(album_id, client, cfg, db)
            album = await pending.resolve()
            if album is None:
                raise RuntimeError(f"could not resolve Qobuz album {album_id}")
            await album.rip()
            on_progress(1.0)
        finally:
            await client.session.close()