Lyra/worker/lyra_worker/pipeline.py

import shutil
from dataclasses import replace
from typing import Sequence

import psycopg

from lyra_worker.adapters.base import SourceAdapter
from lyra_worker.confidence import score_confidence
from lyra_worker.library import album_dir, import_album, staging_dir
from lyra_worker.quality import quality_class
from lyra_worker.ranker import rank_candidates
from lyra_worker.types import Candidate, MBTarget


def _set_state(conn: psycopg.Connection, job_id: str, state: str, stage: str) -> None:
    with conn.cursor() as cur:
        cur.execute(
            'UPDATE "Job" SET state = %s, "currentStage" = %s, "updatedAt" = now() WHERE id = %s',
            (state, stage, job_id),
        )
    conn.commit()


def _request_id(conn: psycopg.Connection, job_id: str) -> str:
    with conn.cursor() as cur:
        cur.execute('SELECT "requestId" FROM "Job" WHERE id = %s', (job_id,))
        return cur.fetchone()[0]


def _fail(conn: psycopg.Connection, job_id: str, reason: str) -> None:
    with conn.cursor() as cur:
        cur.execute(
            'UPDATE "Job" SET state = \'needs_attention\', error = %s, "updatedAt" = now() WHERE id = %s',
            (reason, job_id),
        )
        cur.execute(
            'UPDATE "Request" SET status = \'needs_attention\' WHERE id = (SELECT "requestId" FROM "Job" WHERE id = %s)',
            (job_id,),
        )
    conn.commit()


def _load_target(conn: psycopg.Connection, job_id: str) -> MBTarget:
    with conn.cursor() as cur:
        cur.execute(
            'SELECT artist, album FROM "Request" WHERE id = (SELECT "requestId" FROM "Job" WHERE id = %s)',
            (job_id,),
        )
        artist, album = cur.fetchone()
    return MBTarget(artist=artist, album=album)


def _already_in_library(conn: psycopg.Connection, target: MBTarget) -> bool:
    with conn.cursor() as cur:
        cur.execute(
            'SELECT 1 FROM "LibraryItem" WHERE artist = %s AND album = %s LIMIT 1',
            (target.artist, target.album),
        )
        return cur.fetchone() is not None


def _is_upgrade_job(conn: psycopg.Connection, job_id: str) -> bool:
    with conn.cursor() as cur:
        cur.execute(
            'SELECT "monitoredReleaseId" IS NOT NULL FROM "Request" '
            'WHERE id = (SELECT "requestId" FROM "Job" WHERE id = %s)',
            (job_id,),
        )
        row = cur.fetchone()
        return bool(row and row[0])


def _already_in_library_at_cutoff(conn: psycopg.Connection, target: MBTarget, cutoff: int) -> bool:
    with conn.cursor() as cur:
        cur.execute(
            'SELECT 1 FROM "LibraryItem" WHERE artist = %s AND album = %s '
            'AND "qualityClass" >= %s LIMIT 1',
            (target.artist, target.album, cutoff),
        )
        return cur.fetchone() is not None


def _library_quality(conn: psycopg.Connection, target: MBTarget) -> int | None:
    with conn.cursor() as cur:
        cur.execute(
            'SELECT "qualityClass" FROM "LibraryItem" WHERE artist = %s AND album = %s',
            (target.artist, target.album),
        )
        row = cur.fetchone()
        return row[0] if row else None


def _persist_candidates(conn: psycopg.Connection, job_id: str, candidates: list[Candidate]) -> None:
    with conn.cursor() as cur:
        for c in candidates:
            cur.execute(
                'INSERT INTO "Candidate" (id, "jobId", source, format, "qualityClass", '
                '"trackCount", confidence, "sourceRef", chosen, "createdAt") '
                "VALUES (gen_random_uuid()::text, %s, %s, %s, %s, %s, %s, %s, false, now())",
                (job_id, c.source, c.quality.fmt, quality_class(c.quality),
                 c.track_count, c.confidence, c.source_ref),
            )
    conn.commit()


def _mark_chosen(conn: psycopg.Connection, job_id: str, source_ref: str) -> None:
    with conn.cursor() as cur:
        cur.execute(
            'UPDATE "Candidate" SET chosen = true WHERE "jobId" = %s AND "sourceRef" = %s',
            (job_id, source_ref),
        )
    conn.commit()


def _import(conn: psycopg.Connection, job_id: str, target: MBTarget,
            winner: Candidate, path: str) -> None:
    request_id = _request_id(conn, job_id)
    with conn.cursor() as cur:
        cur.execute(
            'INSERT INTO "LibraryItem" (id, "requestId", artist, album, path, source, '
            'format, "qualityClass", "importedAt") '
            "VALUES (gen_random_uuid()::text, %s, %s, %s, %s, %s, %s, %s, now()) "
            'ON CONFLICT (artist, album) DO UPDATE SET '
            '  "requestId" = EXCLUDED."requestId", path = EXCLUDED.path, source = EXCLUDED.source, '
            '  format = EXCLUDED.format, "qualityClass" = EXCLUDED."qualityClass", '
            '  "importedAt" = now() '
            'WHERE EXCLUDED."qualityClass" > "LibraryItem"."qualityClass"',
            (request_id, target.artist, target.album, path, winner.source,
             winner.quality.fmt, quality_class(winner.quality)),
        )
        cur.execute(
            "UPDATE \"Request\" SET status = 'completed' WHERE id = %s", (request_id,)
        )
    conn.commit()


def run_pipeline(
    conn: psycopg.Connection,
    job_id: str,
    adapters: Sequence[SourceAdapter],
    resolver=None,
    tagger=None,
    min_confidence: float = 0.7,
    dest_root: str = "/music",
    staging_root: str | None = None,
    upgrade_cutoff: int | None = None,
) -> None:
    """Real staged acquisition using source-agnostic adapters. Fakes in this plan."""
    names = [a.name for a in adapters]
    if len(names) != len(set(names)):
        raise ValueError(f"adapter names must be unique, got {names}")

    # 1. intake
    _set_state(conn, job_id, "matching", "intake")
    target = _load_target(conn, job_id)
    if resolver is not None:
        resolved = resolver.resolve(target.artist, target.album)
        if resolved is not None:
            target = resolved
    # dedupe: a monitor-driven upgrade job proceeds unless the existing copy already meets
    # the cutoff; a plain request short-circuits on any existing copy (slice-1 behavior).
    if upgrade_cutoff is not None and _is_upgrade_job(conn, job_id):
        _dedupe_hit = _already_in_library_at_cutoff(conn, target, upgrade_cutoff)
    else:
        _dedupe_hit = _already_in_library(conn, target)
    if _dedupe_hit:
        _set_state(conn, job_id, "imported", "import")
        with conn.cursor() as cur:
            cur.execute(
                "UPDATE \"Request\" SET status = 'completed' WHERE id = %s",
                (_request_id(conn, job_id),),
            )
        conn.commit()
        return

    # 2. match — search all adapters; score + persist EVERY candidate found
    _set_state(conn, job_id, "matching", "match")
    found: list[Candidate] = []
    for adapter in adapters:
        try:
            results = adapter.search(target)
        except Exception as e:  # a down source contributes no candidates, never crashes the job
            print(f"pipeline: adapter {adapter.name} search failed: {e}", flush=True)
            continue
        for c in results:
            found.append(replace(c, source_tier=adapter.tier, confidence=score_confidence(target, c)))
    _persist_candidates(conn, job_id, found)

    # 3. rank
    _set_state(conn, job_id, "matched", "rank")
    ranked = rank_candidates(target, found, min_confidence)
    if not ranked:
        _fail(conn, job_id, "no candidate above confidence threshold")
        return

    # 4. download (fall-through) into an isolated per-job staging dir
    _set_state(conn, job_id, "downloading", "download")
    by_source = {a.name: a for a in adapters}
    staging = staging_dir(staging_root or f"{dest_root}/.staging", job_id)
    winner = None
    result = None
    try:
        for candidate in ranked:
            adapter = by_source.get(candidate.source)
            if adapter is None:
                continue
            shutil.rmtree(staging, ignore_errors=True)  # clean slate per attempt
            result = adapter.download(candidate, staging, lambda _pct: None)
            if result.ok:
                winner = candidate
                _mark_chosen(conn, job_id, candidate.source_ref)
                break
        if winner is None or result is None or not result.ok:
            _fail(conn, job_id, "all downloads failed")
            return

        # 5. verify completeness on the staging copy, then promote + tag
        _set_state(conn, job_id, "tagging", "tag")
        expected = target.track_count if target.track_count is not None else winner.track_count
        if result.track_count < expected:
            _fail(conn, job_id, "incomplete download")
            return

        final = album_dir(dest_root, target)
        existing_q = _library_quality(conn, target)
        new_q = quality_class(winner.quality)
        _set_state(conn, job_id, "imported", "import")
        if existing_q is None or new_q > existing_q:
            import_album(staging, final)  # clean audio(+cover) move, atomically replacing any prior copy
            if tagger is not None:
                try:
                    tagger.tag_album(final, target)
                except Exception as e:  # a tagging failure must not discard a good download
                    print(f"pipeline: tagging failed for job {job_id}: {e}", flush=True)
            _import(conn, job_id, target, winner, final)
        else:
            # an existing copy is already at >= this quality — keep it untouched, just complete the request
            with conn.cursor() as cur:
                cur.execute(
                    "UPDATE \"Request\" SET status = 'completed' WHERE id = %s",
                    (_request_id(conn, job_id),),
                )
            conn.commit()
    finally:
        shutil.rmtree(staging, ignore_errors=True)  # a partial/failed download never persists