Lyra/worker/lyra_worker/pipeline.py

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
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 _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 NOTHING',
            (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,
    min_confidence: float = 0.7,
    dest_root: str = "/music",
) -> 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
    if _already_in_library(conn, target):
        _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)
    _set_state(conn, job_id, "downloading", "download")
    by_source = {a.name: a for a in adapters}
    dest = album_dir(dest_root, target)
    winner = None
    result = None
    for candidate in ranked:
        adapter = by_source.get(candidate.source)
        if adapter is None:
            continue
        result = adapter.download(candidate, dest, 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. tag (integrity check against the canonical track count when known)
    _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

    # 6. import
    _set_state(conn, job_id, "imported", "import")
    _import(conn, job_id, target, winner, result.path or dest)