Lyra/worker/lyra_worker/adapters/_slskd.py

import json
import os
import shutil
import time
import unicodedata
from typing import Callable

import requests

_DEFAULT_DOWNLOADS_ROOT = "/slskd-downloads"

_AUDIO_EXT = {"flac", "mp3", "m4a", "ogg", "opus", "wav", "aac"}
_LOSSLESS_EXT = {"flac", "wav"}
_SEARCH_POLLS = 40    # * 3s ≈ 2 min
_MAX_XFER_POLLS = 1200  # * 3s ≈ 60 min absolute backstop for a slow-but-progressing peer
_STALL_POLLS = 30     # * 3s ≈ 90s of no byte progress → abandon a queued/dead peer and fall through
_POLL_SECONDS = 3


def _basename(path: str) -> str:
    return path.rsplit("\\", 1)[-1].rsplit("/", 1)[-1]


def _dirname(path: str) -> str:
    if "\\" in path:
        return path.rsplit("\\", 1)[0]
    if "/" in path:
        return path.rsplit("/", 1)[0]
    return ""


def _eta_seconds(total_bytes: int, bytes_now: int, speed_bps: float) -> int | None:
    """Seconds until the transfer finishes at the current measured rate, or None when it can't be
    estimated (no throughput sample yet, unknown total, or already complete). Pure — unit-tested."""
    if speed_bps <= 0 or total_bytes <= 0 or bytes_now >= total_bytes:
        return None
    return int((total_bytes - bytes_now) / speed_bps)


def _norm(s: str) -> str:
    """Lowercase, strip diacritics, and reduce to alphanumeric tokens, so 'Beyoncé' matches a
    peer's 'Beyonce' folder and punctuation ('I Am… Sasha Fierce') doesn't defeat the match."""
    s = unicodedata.normalize("NFKD", s or "")
    s = "".join(ch for ch in s if not unicodedata.combining(ch))
    s = "".join(ch if ch.isalnum() else " " for ch in s.lower())
    return " ".join(s.split())


def _path_matches(needle: str, path: str) -> bool:
    """True if `needle` (an artist or album name) appears in a file path, tolerant of accents and
    punctuation. Every significant token (≥3 chars) must be present; for a very short needle
    (e.g. 'U2') the whole normalized string must appear."""
    npath = _norm(path)
    tokens = [t for t in _norm(needle).split() if len(t) >= 3]
    if not tokens:
        whole = _norm(needle)
        return bool(whole) and whole in npath
    return all(t in npath for t in tokens)


def _keep_matching(responses: list, needle: str) -> list:
    """Drop files — then now-empty responses — whose path doesn't match `needle`. Makes an album-
    or artist-only fallback search precise so we never grab a same-titled release by a different
    artist. Pure — unit-tested offline."""
    out = []
    for resp in responses:
        files = [f for f in (resp.get("files") or []) if _path_matches(needle, f.get("filename", ""))]
        if files:
            out.append({**resp, "files": files})
    return out


def _parse_search_responses(responses: list) -> list[dict]:
    """Turn slskd search responses into album candidates (one per peer+directory), ordered so the
    peers most likely to deliver quickly come first: a free upload slot, then higher upload speed,
    then a shorter queue. The pipeline tries candidates in this order, so ranking fast/free peers
    first (over slow or queued ones) is what makes the Soulseek fall-through actually converge.
    Pure — no I/O — so it is unit-tested offline."""
    scored: list[tuple] = []
    for resp in responses:
        username = resp.get("username", "")
        has_slot = 1 if resp.get("hasFreeUploadSlot") else 0
        speed = int(resp.get("uploadSpeed") or 0)
        queue = int(resp.get("queueLength") or 0)
        by_dir: dict[str, list] = {}
        for f in resp.get("files", []) or []:
            name = f.get("filename", "")
            ext = name.rsplit(".", 1)[-1].lower() if "." in name else ""
            if ext not in _AUDIO_EXT:
                continue
            by_dir.setdefault(_dirname(name), []).append(
                {"filename": name, "size": f.get("size", 0), "ext": ext}
            )
        for directory, dfiles in by_dir.items():
            if not dfiles:
                continue
            lossless = all(f["ext"] in _LOSSLESS_EXT for f in dfiles)
            dirname = _basename(directory)
            if " - " in dirname:
                guess_artist, guess_album = (p.strip() for p in dirname.split(" - ", 1))
            else:
                guess_artist, guess_album = "", dirname
            candidate = {
                "source_ref": json.dumps({
                    "username": username,
                    "files": [{"filename": f["filename"], "size": f["size"]} for f in dfiles],
                }),
                "title": guess_album,
                "artist": guess_artist,
                "track_count": len(dfiles),
                "format": "FLAC" if lossless else "MP3",
                "bitrate": None,
            }
            scored.append((has_slot, speed, -queue, candidate))
    scored.sort(key=lambda t: (t[0], t[1], t[2]), reverse=True)  # free + fast + short-queue first
    return [c for *_rest, c in scored]


class SlskdClient:
    """Real Soulseek client talking to an slskd daemon's REST API.

    NOT unit-tested offline; see test_soulseek_live.py. slskd response/transfer
    JSON shapes may vary by version — validate with the live test.
    """

    def __init__(self, config: dict):
        self._url = (config.get("slskd.url", "") or "").rstrip("/")
        self._key = config.get("slskd.api_key", "")
        # Where slskd (an external daemon) writes finished downloads. slskd saves into ITS
        # OWN directory, so this path must be a share/mount both slskd and this worker see.
        # Precedence: the slskd.downloadsDir Config row (Settings-overridable) > the
        # SLSKD_DOWNLOADS_ROOT env (the container mount point) > a sensible default.
        self._downloads_dir = (
            config.get("slskd.downloadsDir")
            or os.environ.get("SLSKD_DOWNLOADS_ROOT")
            or _DEFAULT_DOWNLOADS_ROOT
        )

    def is_configured(self) -> bool:
        return bool(self._url and self._key)

    def _headers(self) -> dict:
        return {"X-API-Key": self._key, "Content-Type": "application/json"}

    def _get(self, path: str) -> dict | list:
        r = requests.get(f"{self._url}{path}", headers=self._headers(), timeout=30)
        r.raise_for_status()
        return r.json()

    def _run_search(self, text: str) -> list:
        """POST a search, wait for it to complete, and return the raw responses (a list)."""
        r = requests.post(
            f"{self._url}/api/v0/searches",
            headers=self._headers(),
            json={"searchText": text},
            timeout=30,
        )
        r.raise_for_status()
        search_id = r.json()["id"]

        for _ in range(_SEARCH_POLLS):
            time.sleep(_POLL_SECONDS)
            state = self._get(f"/api/v0/searches/{search_id}")
            if state.get("isComplete") or "Completed" in str(state.get("state", "")):
                break

        responses = self._get(f"/api/v0/searches/{search_id}/responses")
        return responses if isinstance(responses, list) else []

    def search_album(self, artist: str, album: str) -> list[dict]:
        responses = self._run_search(f"{artist} {album}")
        if not responses:
            # The Soulseek server silently drops searches containing certain blocklisted terms
            # (major-label artists / album titles like Adele, Rihanna, Beyoncé, "Lemonade") →
            # 0 responses. The blocked term poisons the whole query, but the OTHER field usually
            # isn't blocked, so retry with each field alone (normalized) and keep only results
            # whose path still matches the field we dropped. That recovers e.g. Beyoncé's
            # "I Am… Sasha Fierce" via the album title, without ever grabbing a same-titled album
            # by someone else.
            for query, needle in ((album, artist), (artist, album)):
                q = _norm(query)
                if not q:
                    continue
                alt = _keep_matching(self._run_search(q), needle)
                if alt:
                    responses = alt
                    break
        return _parse_search_responses(responses)

    def download(self, source_ref: str, dest: str, on_progress: Callable[[float], None]) -> dict:
        ref = json.loads(source_ref)
        username = ref["username"]
        files = ref["files"]
        wanted = {f["filename"] for f in files}
        os.makedirs(dest, exist_ok=True)
        xfer_ids: set[str] = set()  # slskd transfer ids we started, for cancel-on-abandon
        try:
            return self._download(username, files, wanted, dest, on_progress, xfer_ids)
        except Exception:
            # Abandoning this peer (timeout / error / the pipeline moving on): cancel the transfers
            # we enqueued so they don't keep running and pile up in slskd as the job falls through
            # to other peers/sources for the same album.
            self._cancel_transfers(username, xfer_ids)
            raise

    def _download(self, username, files, wanted, dest, on_progress, xfer_ids: set) -> dict:
        r = requests.post(
            f"{self._url}/api/v0/transfers/downloads/{username}",
            headers=self._headers(),
            json=[{"filename": f["filename"], "size": f["size"]} for f in files],
            timeout=30,
        )
        r.raise_for_status()
        on_progress(0.0)

        total = len(files)
        total_bytes = sum(int(f.get("size") or 0) for f in files)
        completed = False
        last_bytes = -1
        stall = 0
        speed = 0.0  # bytes/sec, EWMA-smoothed so the ETA doesn't jitter poll to poll
        prev_bytes = 0
        prev_t = time.monotonic()
        for _ in range(_MAX_XFER_POLLS):
            time.sleep(_POLL_SECONDS)
            data = self._get(f"/api/v0/transfers/downloads/{username}")
            states: list[str] = []
            bytes_now = 0
            for directory in data.get("directories", []) if isinstance(data, dict) else []:
                for f in directory.get("files", []):
                    if f.get("filename") in wanted:  # only the files we enqueued
                        states.append(str(f.get("state", "")))
                        bytes_now += int(f.get("bytesTransferred") or 0)
                        if f.get("id"):  # remember the transfer so we can cancel it if abandoned
                            xfer_ids.add(str(f["id"]))
            if states:
                if any(bad in s for s in states for bad in ("Errored", "Rejected", "Cancelled")):
                    raise RuntimeError("soulseek transfer failed")
                now = time.monotonic()
                dt = now - prev_t
                if dt > 0 and bytes_now >= prev_bytes:
                    inst = (bytes_now - prev_bytes) / dt
                    speed = inst if speed == 0.0 else 0.6 * speed + 0.4 * inst
                prev_bytes, prev_t = bytes_now, now
                done = sum(1 for s in states if "Completed" in s)
                on_progress(min(1.0, done / max(total, 1)), _eta_seconds(total_bytes, bytes_now, speed))
                if done >= total:
                    completed = True
                    break
                # Stall detection: a peer making byte progress keeps its slot (up to the absolute
                # backstop above) so a slow-but-working transfer can finish; one that sends no bytes
                # for _STALL_POLLS (~90s) is queued/dead — abandon it fast and fall through.
                if bytes_now > last_bytes:
                    last_bytes = bytes_now
                    stall = 0
                else:
                    stall += 1
                    if stall >= _STALL_POLLS:
                        raise TimeoutError(f"soulseek transfer stalled for {username}")

        if not completed:
            raise TimeoutError(f"soulseek download did not complete for {username}")

        # slskd saved the files into its OWN downloads dir — retrieve them into `dest` so the
        # pipeline (which verifies + imports against `dest`/staging) actually sees the audio.
        moved = self._retrieve(files, dest)
        if moved == 0:
            raise RuntimeError(
                f"soulseek transfer completed but no files were found under "
                f"{self._downloads_dir!r} to move into staging — check that the slskd "
                f"downloads dir is shared/mounted (SLSKD_DOWNLOADS_ROOT / slskd.downloadsDir)"
            )
        return {"track_count": moved, "path": dest}

    def _cancel_transfers(self, username: str, xfer_ids: set) -> None:
        """Cancel the given in-flight slskd transfers for a peer. Best-effort: a cleanup failure
        must never mask the original download failure."""
        for tid in xfer_ids:
            try:
                requests.delete(
                    f"{self._url}/api/v0/transfers/downloads/{username}/{tid}",
                    headers=self._headers(),
                    timeout=15,
                )
            except Exception:
                pass

    def _retrieve(self, files: list[dict], dest: str) -> int:
        """Move the just-finished slskd files out of its downloads dir into `dest`.

        slskd's on-disk layout can vary by version/config, so match tolerantly: for each
        wanted file, find on-disk candidates by basename anywhere under the downloads root,
        then prefer the one whose parent-folder name matches the remote directory and whose
        size matches. Returns the number of files actually moved. Offline-testable.
        """
        root = self._downloads_dir
        if not root or not os.path.isdir(root):
            raise RuntimeError(
                f"slskd downloads dir not found: {root!r} — mount slskd's downloads "
                f"directory into the worker (SLSKD_DOWNLOADS_ROOT / slskd.downloadsDir)"
            )
        os.makedirs(dest, exist_ok=True)

        wanted: dict[str, dict] = {}  # basename(lower) -> {parent, size}
        for f in files:
            name = f.get("filename", "")
            wanted[_basename(name).lower()] = {
                "parent": _basename(_dirname(name)).lower(),
                "size": f.get("size"),
            }

        found: dict[str, list[str]] = {b: [] for b in wanted}
        for cur_root, _dirs, fnames in os.walk(root):
            for fn in fnames:
                key = fn.lower()
                if key in found:
                    found[key].append(os.path.join(cur_root, fn))

        moved = 0
        for basename, paths in found.items():
            if not paths:
                continue
            want = wanted[basename]

            def _score(p: str) -> tuple[int, int]:
                parent_ok = int(os.path.basename(os.path.dirname(p)).lower() == want["parent"])
                size_ok = 0
                try:
                    size_ok = int(bool(want["size"]) and os.path.getsize(p) == want["size"])
                except OSError:
                    pass
                return (parent_ok, size_ok)

            best = max(paths, key=_score)
            shutil.move(best, os.path.join(dest, os.path.basename(best)))
            moved += 1
        return moved