Lyra/worker/lyra_worker/scan.py

import os
import re
from dataclasses import dataclass

import psycopg

from lyra_worker.probe import AudioProbe

_AUDIO_EXT = {".flac", ".mp3", ".m4a", ".opus", ".ogg", ".aac", ".wav"}
_YEAR = re.compile(r"^(.*) \((\d{4})\)$")


@dataclass(frozen=True)
class ScanResult:
    imported: int
    skipped: int


def _parse_album_dir(name: str) -> str:
    m = _YEAR.match(name)
    return m.group(1) if m else name


def _first_audio(album_path: str) -> str | None:
    for name in sorted(os.listdir(album_path)):
        if os.path.splitext(name)[1].lower() in _AUDIO_EXT:
            return os.path.join(album_path, name)
    return None


def _ensure_artist(conn, target) -> str | None:
    """Upsert the WatchedArtist (by mbid) and return its id, or None if the target has no artist mbid."""
    if not target.artist_mbid:
        return None
    with conn.cursor() as cur:
        cur.execute(
            'INSERT INTO "WatchedArtist" (id, mbid, name, "autoMonitorFuture", "monitorFrom", "createdAt") '
            "VALUES (gen_random_uuid()::text, %s, %s, false, now(), now()) "
            'ON CONFLICT (mbid) DO NOTHING',
            (target.artist_mbid, target.artist),
        )
        cur.execute('SELECT id FROM "WatchedArtist" WHERE mbid = %s', (target.artist_mbid,))
        row = cur.fetchone()
    conn.commit()
    return row[0] if row else None


def _populate_discography(conn, browser, watched_id, artist_mbid: str, artist_name: str) -> None:
    """Store the artist's full release-group list as unmonitored rows (like a web follow)."""
    for rg in browser.browse_release_groups(artist_mbid):
        with conn.cursor() as cur:
            cur.execute(
                'INSERT INTO "MonitoredRelease" (id, "watchedArtistId", "artistMbid", "artistName", '
                '"rgMbid", album, "primaryType", "secondaryTypes", "firstReleaseDate", monitored, state, "createdAt") '
                "VALUES (gen_random_uuid()::text, %s, %s, %s, %s, %s, %s, %s, %s, false, 'wanted', now()) "
                'ON CONFLICT ("rgMbid") DO NOTHING',
                (watched_id, artist_mbid, artist_name, rg.rg_mbid, rg.title, rg.primary_type or None,
                 list(rg.secondary_types), rg.first_release_date or None),
            )
    conn.commit()


def _record_owned(conn, target, quality_class: int, fmt: str, path: str, watched_id) -> bool:
    """Mark the owned release monitored+fulfilled and record the LibraryItem. Returns True if newly imported."""
    with conn.cursor() as cur:
        if target.rg_mbid:
            cur.execute(
                'INSERT INTO "MonitoredRelease" (id, "watchedArtistId", "artistMbid", "artistName", '
                '"rgMbid", album, "secondaryTypes", monitored, state, "currentQualityClass", '
                '"firstGrabbedAt", "createdAt") '
                "VALUES (gen_random_uuid()::text, %s, %s, %s, %s, %s, '{}', true, 'fulfilled', %s, now(), now()) "
                'ON CONFLICT ("rgMbid") DO UPDATE SET monitored = true, state = \'fulfilled\', '
                '  "currentQualityClass" = EXCLUDED."currentQualityClass", '
                '  "firstGrabbedAt" = COALESCE("MonitoredRelease"."firstGrabbedAt", now()), '
                '  "watchedArtistId" = COALESCE("MonitoredRelease"."watchedArtistId", EXCLUDED."watchedArtistId")',
                (watched_id, target.artist_mbid, target.artist, target.rg_mbid, target.album, quality_class),
            )
        cur.execute(
            'INSERT INTO "LibraryItem" (id, "requestId", artist, album, path, source, format, '
            '"qualityClass", "importedAt") '
            "VALUES (gen_random_uuid()::text, NULL, %s, %s, %s, 'scan', %s, %s, now()) "
            'ON CONFLICT (artist, album) DO NOTHING',
            (target.artist, target.album, path, fmt, quality_class),
        )
        newly = cur.rowcount == 1  # a new LibraryItem means this album wasn't recorded before
    conn.commit()
    return newly


def scan_library(conn: psycopg.Connection, resolver, probe: AudioProbe, browser, dest_root: str = "/music") -> ScanResult:
    """Walk `{dest_root}/{Artist}/{Album}/`, resolve each album against MusicBrainz, and record
    matched albums as have + monitored + followed. Also browses each followed artist's full
    discography (once per scan run) and stores it as unmonitored releases, matching web-follow
    behavior. Unmatched or audio-less folders are skipped."""
    imported = 0
    skipped = 0
    browsed: set[str] = set()  # artists whose discography we've populated this run
    if not os.path.isdir(dest_root):
        return ScanResult(0, 0)
    for artist_name in sorted(os.listdir(dest_root)):
        artist_path = os.path.join(dest_root, artist_name)
        if artist_name.startswith(".") or not os.path.isdir(artist_path):
            continue
        for album_folder in sorted(os.listdir(artist_path)):
            album_path = os.path.join(artist_path, album_folder)
            if album_folder.startswith(".") or not os.path.isdir(album_path):
                continue
            audio = _first_audio(album_path)
            if audio is None:
                continue  # no audio → not an album
            pr = probe.probe(audio)
            album_name = _parse_album_dir(album_folder)
            target = resolver.resolve(artist_name, album_name) if artist_name and album_name else None
            if target is None:  # fall back to the file's embedded tags
                if pr.artist and pr.album and (pr.artist, pr.album) != (artist_name, album_name):
                    target = resolver.resolve(pr.artist, pr.album)
            if target is None:
                skipped += 1
                continue
            watched_id = _ensure_artist(conn, target)
            if watched_id is not None and target.artist_mbid not in browsed:
                _populate_discography(conn, browser, watched_id, target.artist_mbid, target.artist)
                browsed.add(target.artist_mbid)
            if _record_owned(conn, target, pr.quality_class, pr.fmt, album_path, watched_id):
                imported += 1
    return ScanResult(imported, skipped)